`core.models.vision.vit_model`#

Native Megatron-Core Vision Transformer.

Provides ViTModel (Pixtral/CLIP), QwenVLViTModel (Qwen3.5-MoE VL), and KimiViTModel (Kimi-K2). All use mcore TransformerBlock — no HuggingFace at runtime.

Module Contents#

Classes#

`PatchEmbedding`	Conv2d patch extractor — no positional embedding.
`Pixtral2DRotaryEmbedding`	2D RoPE for Mistral-native Pixtral-family vision encoders.
`PixtralLargePatchMerger`	Spatial 2×2 patch merger for Pixtral-Large.
`ViTModel`	Generic Vision Transformer — native mcore, no transformers dependency.
`QwenVL2DRotaryEmbedding`	2D RoPE for Qwen VL.
`QwenLearnedPosEmbed`	Bilinear-interpolatable learned 2D position embedding.
`QwenPatchMerger`	Spatial 2×2 patch grouping for Qwen VL.
`QwenPatchEmbedding`	Patch embedding with .proj Conv3d, matching checkpoint key naming (patch_embed.proj.*).
`QwenVLViTModel`	Native mcore vision encoder for Qwen3.5-MoE VL.
`Kimi2DRotaryEmbedding`	2D RoPE for Kimi-K2 (interleaved complex style).
`KimiLearned2DPosEmbed`	Bicubic-interpolatable learned 2D spatial position embedding for Kimi.
`KimiPatchMerger`	Spatial 2×2 patch grouping for Kimi.
`KimiViTModel`	Native mcore vision encoder for Kimi-K2.

Functions#

`_dynamic_patch_grid`
`_dynamic_patch_grid_lists`	Return dynamic patch metadata without synchronizing CUDA to the host.
`_cat_rope`

Data#

_NORM_IMPL

API#

core.models.vision.vit_model._NORM_IMPL#: None

class core.models.vision.vit_model.PatchEmbedding( in_channels: int, hidden_size: int, patch_dim: int, bias: bool = False, )#

Bases: torch.nn.Module

Conv2d patch extractor — no positional embedding.

Initialization

forward(x: torch.Tensor) → Tuple[torch.Tensor, int, int]#: (B, C, H, W) → (B, N, hidden), (h_patches, w_patches)

forward_patches(x: torch.Tensor) → torch.Tensor#: Pre-patchified input -> (B, N, hidden).

core.models.vision.vit_model._dynamic_patch_grid( imgs_sizes: Union[List[Tuple[int, int]], torch.Tensor], patch_dim: int, device: torch.device, ) → Tuple[torch.Tensor, torch.Tensor]#

core.models.vision.vit_model._dynamic_patch_grid_lists( imgs_sizes: Union[List[Tuple[int, int]], torch.Tensor], patch_dim: int, ) → Tuple[List[Tuple[int, int]], List[int]]#: Return dynamic patch metadata without synchronizing CUDA to the host.

core.models.vision.vit_model._cat_rope(chunks: List[torch.Tensor]) → torch.Tensor#

class core.models.vision.vit_model.Pixtral2DRotaryEmbedding( head_dim: int, max_patches_per_side: int, rope_theta: float = 10000.0, )#

Bases: torch.nn.Module

2D RoPE for Mistral-native Pixtral-family vision encoders.

Mistral/vLLM applies RoPE as complex multiplication over adjacent hidden dimension pairs. This returns repeat-interleaved angles and must be used with rotary_interleaved=True in TransformerConfig.

Initialization

forward( h_patches: int, w_patches: int, device: torch.device, ) → torch.Tensor#: Returns freqs of shape (h_patches * w_patches, 1, 1, head_dim). This is the format mcore’s apply_rotary_pos_emb expects for rotary_pos_emb. mcore handles cos/sin internally.

class core.models.vision.vit_model.PixtralLargePatchMerger( transformer_config: megatron.core.transformer.transformer_config.TransformerConfig, spatial_merge_size: int = 2, )#

Bases: torch.nn.Module

Spatial 2×2 patch merger for Pixtral-Large.

Applies RMSNorm on pre-merge tokens, folds 2×2 spatial blocks into a single 4*hidden vector, then projects back to hidden_size with a bias-less Linear. Matches Mistral-Large-3’s pre_mm_projector_norm + patch_merger.merging_layer.

Initialization

forward( x: torch.Tensor, h_patches: int, w_patches: int, ) → torch.Tensor#: (B, h_pw_p, H) → (B, h_p/mw_p/m, H).

class core.models.vision.vit_model.ViTModel( transformer_config: megatron.core.transformer.transformer_config.TransformerConfig, transformer_layer_spec: megatron.core.transformer.spec_utils.ModuleSpec, patch_dim: int = 16, img_h: int = 1024, img_w: int = 1024, in_channels: int = 3, patch_embed_bias: bool = False, add_class_token: bool = False, class_token_len: int = 0, ln_pre: bool = True, ln_pre_eps: Optional[float] = None, pos_emb_type: str = 'rope2d', rope_theta: float = 10000.0, use_merger: bool = False, spatial_merge_size: int = 2, pg_collection=None, vp_stage: Optional[int] = None, )#

Bases: megatron.core.transformer.module.MegatronModule

Generic Vision Transformer — native mcore, no transformers dependency.

Supports:

Pixtral (2D RoPE, SwiGLU, RMSNorm, no CLS, no bias)
Pixtral-Large (same + 2×2 patch merger after the transformer stack)
CLIP/SigLIP (learned absolute pos, GELU, LayerNorm, CLS token) when add_class_token=True and pos_emb_type=’learned_absolute’

Parameters:

transformer_config –
Standard mcore TransformerConfig. Configure:
- normalization, norm_epsilon
- gated_linear_unit, activation_func
- add_bias_linear
- num_layers, hidden_size, num_attention_heads, ffn_hidden_size
transformer_layer_spec – Layer spec for TransformerBlock (bidirectional).
patch_dim – Patch size in pixels (16 for Pixtral, 14 for CLIP/SigLIP).
img_w (img_h /) – Maximum image dimensions (controls RoPE table size).
add_class_token – Prepend a learnable CLS token (CLIP style).
class_token_len – Width of CLS token.
ln_pre – Apply RMSNorm/LayerNorm before the transformer stack.
ln_pre_eps – Epsilon for pre-transformer norm (defaults to norm_epsilon).
pos_emb_type – ‘rope2d’ | ‘learned_absolute’ | ‘none’.
rope_theta – RoPE base frequency.
use_merger – If True, append a Pixtral-Large-style 2×2 patch merger that reduces sequence length by 4 and keeps the output width at hidden_size.
spatial_merge_size – Merger block size (default 2 → 2×2 → 4× reduction).

Initialization

property num_patches_per_image: int#: Return the number of vision patches for a single fixed-size image.

set_input_tensor(input_tensor)#: Set the input tensor for the decoder (pipeline-parallel entrypoint).

forward( pixel_values: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, imgs_sizes=None, packed_seq_params=None, ) → torch.Tensor#

Parameters:

pixel_values – (B, C, H, W)
attention_mask – optional mask for TransformerBlock

Returns:

(B, N, hidden_size) where N includes CLS tokens if add_class_token=True

class core.models.vision.vit_model.QwenVL2DRotaryEmbedding( head_dim: int, max_patches_per_side: int, rope_theta: float = 10000.0, )#

Bases: torch.nn.Module

2D RoPE for Qwen VL.

Uses the same inv_freq for both H and W; concatenates row_freqs + col_freqs then duplicates to fill head_dim, matching Qwen3_5MoeVisionRotaryEmbedding

rot_pos_emb exactly.

Initialization

forward(row_ids: torch.Tensor, col_ids: torch.Tensor) → torch.Tensor#

Parameters:

row_ids – (N,) integer row indices of each patch, in [0, max_patches_per_side).
col_ids – (N,) integer col indices of each patch, in [0, max_patches_per_side).

Returns:

(N, 1, 1, head_dim) — mcore rotary_pos_emb format (raw freqs, not cos/sin)

Return type:

freqs

class core.models.vision.vit_model.QwenLearnedPosEmbed(num_grid_per_side: int, hidden_size: int)#

Bases: torch.nn.Module

Bilinear-interpolatable learned 2D position embedding.

Stores a (num_grid_per_side x num_grid_per_side) grid of vectors and interpolates bilinearly to the actual (h_patches, w_patches) resolution. Matches Qwen3_5MoeVisionModel.fast_pos_embed_interpolate.

Initialization

forward( h_patches: int, w_patches: int, device: torch.device, ) → torch.Tensor#: Returns (h_patches * w_patches, hidden_size) bilinear-interpolated embeddings.

class core.models.vision.vit_model.QwenPatchMerger( hidden_size: int, spatial_merge_size: int, out_hidden_size: int, )#

Bases: torch.nn.Module

Spatial 2×2 patch grouping for Qwen VL.

The mcore path stops at encoder_tokens(), exposing pre-projector tokens. linear_fc1 / linear_fc2 are present only to accept source-checkpoint keys (visual.merger.mlp.*); they aren’t traversed at runtime. Their forward pass is available via forward() for callers that want Qwen’s full merger, but the mcore inference path doesn’t call it, and DDP/FSDP consumers should treat these params as static (they receive no gradient in the mcore path).

Reshapes spatially adjacent 2×2 patches into single vectors, applies LayerNorm, then projects to out_hidden_size via 2-layer MLP. Matches Qwen3_5MoeVisionPatchMerger (use_postshuffle_norm=False default).

Initialization

encoder_tokens( x: torch.Tensor, h_patches: int, w_patches: int, ) → torch.Tensor#

Parameters:: x – (B, h_patches * w_patches, hidden)
Returns:: (B, h_out * w_out, hidden * merge_size^2) before Qwen’s merger MLP.

forward( x: torch.Tensor, h_patches: int, w_patches: int, ) → torch.Tensor#: Return Qwen’s source-model visual merger output.

class core.models.vision.vit_model.QwenPatchEmbedding( in_channels: int, hidden_size: int, patch_dim: int, temporal_patch_size: int, )#

Bases: torch.nn.Module

Patch embedding with .proj Conv3d, matching checkpoint key naming (patch_embed.proj.*).

Initialization

forward(x: torch.Tensor) → torch.Tensor#: Project pixel patches into the transformer hidden dimension.

forward_patches(x: torch.Tensor) → torch.Tensor#: Pre-patchified (B, N, CPP) input -> (B, N, hidden).

class core.models.vision.vit_model.QwenVLViTModel( transformer_config: megatron.core.transformer.transformer_config.TransformerConfig, transformer_layer_spec: megatron.core.transformer.spec_utils.ModuleSpec, patch_dim: int = 16, temporal_patch_size: int = 2, img_h: int = 768, img_w: int = 768, in_channels: int = 3, spatial_merge_size: int = 2, out_hidden_size: int = 2048, num_pos_per_side: int = 48, rope_theta: float = 10000.0, pg_collection=None, vp_stage: Optional[int] = None, )#

Bases: megatron.core.transformer.module.MegatronModule

Native mcore vision encoder for Qwen3.5-MoE VL.

Accepts standard (B, C, H, W) pixel values. Internally performs block-first patch reordering so that spatial 2×2 merge groups are consecutive in the sequence — this matches Qwen’s fast_pos_embed_interpolate ordering.

Parameters:

transformer_config – TransformerConfig with: normalization=’LayerNorm’, layernorm_epsilon=1e-6, add_bias_linear=True, gated_linear_unit=False, activation_func=gelu_tanh, apply_rope_fusion=False
transformer_layer_spec – Layer spec with no_mask attention.
patch_dim – Spatial patch size in pixels (16 for Qwen).
temporal_patch_size – Temporal merge (2 for Qwen).
img_w (img_h /) – Max image dims for RoPE table (default 768 → 48 patches).
spatial_merge_size – Spatial downsampling in merger (2 for Qwen).
out_hidden_size – Source Qwen merger MLP output dimension.
num_pos_per_side – Learned pos embed grid size (48 for Qwen).
rope_theta – RoPE base frequency.

Initialization

set_input_tensor(input_tensor)#: Set the input tensor for the decoder (pipeline-parallel entrypoint).

forward( pixel_values: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, imgs_sizes=None, packed_seq_params=None, ) → torch.Tensor#

Parameters:: pixel_values – (B, C, H, W)
Returns:: (B, h_out * w_out, hidden * spatial_merge_size^2)

class core.models.vision.vit_model.Kimi2DRotaryEmbedding( head_dim: int, max_patches_per_side: int, rope_theta: float = 10000.0, )#

Bases: torch.nn.Module

2D RoPE for Kimi-K2 (interleaved complex style).

Alternates x-position and y-position frequencies for complex pairs: pairs (0,1), (4,5), … get col (x) rotation; pairs (2,3), (6,7), … get row (y) rotation.

Used with rotary_interleaved=True in TransformerConfig so that mcore’s _rotate_half correctly implements complex multiplication by e^(i*theta).

Initialization

forward(row_ids: torch.Tensor, col_ids: torch.Tensor) → torch.Tensor#

Returns freqs of shape (N, 1, 1, head_dim) for mcore interleaved RoPE.

row_ids / col_ids must be in [0, max_patches_per_side).

class core.models.vision.vit_model.KimiLearned2DPosEmbed(height: int, width: int, hidden_size: int)#

Bases: torch.nn.Module

Bicubic-interpolatable learned 2D spatial position embedding for Kimi.

Matches Kimi’s Learnable2DInterpPosEmbDivided_fixed for static images (T=1). For video (T>1), sinusoidal temporal embeddings would be added; we skip that for the static-image parity test.

Initialization

forward( h_patches: int, w_patches: int, device: torch.device, ) → torch.Tensor#: Returns (h_patches * w_patches, hidden_size) on device.

class core.models.vision.vit_model.KimiPatchMerger#

Bases: torch.nn.Module

Spatial 2×2 patch grouping for Kimi.

Groups spatially adjacent 2×2 patches onto a new axis without reducing them, giving (B, h_out * w_out, 4 * hidden). Matches tpool_patch_merger at T=1, where the temporal mean degenerates to a reshape; a T>1 variant would mean over the temporal axis here.

forward( x: torch.Tensor, h_patches: int, w_patches: int, ) → torch.Tensor#

Parameters:: x – (B, h_patches * w_patches, hidden)
Returns:: (B, h_out * w_out, 4 * hidden)

class core.models.vision.vit_model.KimiViTModel( transformer_config: megatron.core.transformer.transformer_config.TransformerConfig, transformer_layer_spec: megatron.core.transformer.spec_utils.ModuleSpec, patch_dim: int = 14, img_h: int = 896, img_w: int = 896, in_channels: int = 3, pos_embed_height: int = 64, pos_embed_width: int = 64, rope_theta: float = 10000.0, pg_collection=None, vp_stage: Optional[int] = None, )#

Bases: megatron.core.transformer.module.MegatronModule

Native mcore vision encoder for Kimi-K2.

Accepts (B, C, H, W) pixel values. Internally:

Conv2d patch embed (14×14)
Learnable 2D spatial position embedding (bicubic interpolated from 64×64 grid)
2D RoPE with interleaved complex-style rotation (rotary_interleaved=True)
27-layer bidirectional TransformerBlock
Final LayerNorm
2×2 spatial group merger (patch grouping, no reduction)

Parameters:

transformer_config – TransformerConfig with: normalization=’LayerNorm’, add_bias_linear=True, gated_linear_unit=False, activation_func=gelu_tanh, rotary_interleaved=True, apply_rope_fusion=False
transformer_layer_spec – Layer spec with no_mask attention.
patch_dim – Spatial patch size (14 for Kimi).
img_w (img_h /) – Max image dims for RoPE table (default 896 → 64 patches).
pos_embed_width (pos_embed_height /) – Learned pos embed grid size (64 for Kimi).
rope_theta – RoPE base frequency (10000 for Kimi).

Initialization

set_input_tensor(input_tensor)#: Set the input tensor for the decoder (pipeline-parallel entrypoint).

forward( pixel_values: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, imgs_sizes=None, packed_seq_params=None, ) → torch.Tensor#

Parameters:: pixel_values – (B, C, H, W)
Returns:: (B, h_out * w_out, 4 * hidden_size)

core.models.vision.vit_model#

Module Contents#

Classes#

Functions#

Data#

API#

`core.models.vision.vit_model`#