nemo_automodel.components.models.mistral_large4.vision

Native-layout Pixtral correctness path; no checkpoint matrix permutation.

The released tower uses FP32 interleaved axial rotation and BF16 source weights. Packing images once per forward keeps FSDP collective counts independent of the number of images on a rank. Attention remains isolated between images.

Module Contents

Classes

Name	Description
`MistralLarge4PatchMerger`	Channel-major spatial patch merging, before the GELU projector.
`MistralLarge4VisionAdapter`	Bias-free vision-to-language GELU projection with released parameter names.
`MistralLarge4VisionBlock`	A single native Pixtral block and FSDP/recomputation boundary.
`MistralLarge4VisionConfig`	Typed native vision settings, parsed once at the checkpoint boundary.
`MistralLarge4VisionEncoder`	Pack variable images into one set of per-layer calls with native weight names.
`_VisionAttention`	-
`_VisionMLP`	-

Functions

Name	Description
`_rotate`	Apply axial rotation without rounding trigonometric tables to BF16.

API

class nemo_automodel.components.models.mistral_large4.vision.MistralLarge4PatchMerger(
    config: nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionConfig,
    dtype: torch.dtype
)

Bases: Module

Channel-major spatial patch merging, before the GELU projector.

merge

= config.spatial_merge_size

merging_layer

nemo_automodel.components.models.mistral_large4.vision.MistralLarge4PatchMerger.forward(
    features: tuple[torch.Tensor, ...],
    sizes: list[tuple[int, int]]
) -> torch.Tensor

Merge nonoverlapping complete patch groups.

Parameters:

features

tuple[torch.Tensor, ...]

Per-image floating [patch_rows * patch_columns, vision_hidden] tensors.

sizes

list[tuple[int, int]]

Per-image patch grid (rows, columns), divisible by merge size.

Returns: torch.Tensor

Independent [total_merged_patches, vision_hidden] tensor, in image order.

class nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionAdapter(
    config: nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionConfig,
    text_hidden_size: int,
    dtype: torch.dtype
)

Bases: Module

Bias-free vision-to-language GELU projection with released parameter names.

w_in

w_out

nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionAdapter.forward(
    features: torch.Tensor
) -> torch.Tensor

Project merged image features.

Parameters:

features

torch.Tensor

Floating [merged_patches, vision_hidden] features.

Returns: torch.Tensor

Independent floating [merged_patches, text_hidden] features.

class nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionBlock(
    config: nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionConfig,
    dtype: torch.dtype
)

Bases: Module

A single native Pixtral block and FSDP/recomputation boundary.

attention

= _VisionAttention(config, dtype=dtype)

attention_norm

feed_forward

= _VisionMLP(config, dtype=dtype)

ffn_norm

nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionBlock.forward(
    x: torch.Tensor,
    angles: torch.Tensor,
    keep: torch.Tensor
) -> torch.Tensor

Apply attention and SwiGLU residuals without mutating inputs.

Parameters:

torch.Tensor

Floating packed [tokens, hidden] image features.

angles

torch.Tensor

FP32 [tokens, head_dim / 2] axial angles.

keep

torch.Tensor

Boolean [tokens, tokens] image membership mask.

Returns: torch.Tensor

Independent features [tokens, hidden].

class nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionConfig(
    hidden_size: int = 2048,
    intermediate_size: int = 6144,
    num_hidden_layers: int = 30,
    num_attention_heads: int = 16,
    num_channels: int = 3,
    patch_size: int = 14,
    image_size: int = 1540,
    rope_theta: float = 10000.0,
    spatial_merge_size: int = 2,
    adapter_bias: bool = False,
    image_token_id: int = 10,
    image_break_token_id: int = 12,
    image_end_token_id: int = 13
)

Dataclass

Typed native vision settings, parsed once at the checkpoint boundary.

adapter_bias

bool = False

hidden_size

int = 2048

image_break_token_id

int = 12

image_end_token_id

int = 13

image_size

int = 1540

image_token_id

int = 10

intermediate_size

int = 6144

num_attention_heads

int = 16

num_channels

int = 3

num_hidden_layers

int = 30

patch_size

int = 14

rope_theta

float = 10000.0

spatial_merge_size

int = 2

nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionConfig.__post_init__() -> None

nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionConfig.build(
    dtype: torch.dtype
) -> nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionEncoder

Construct the native-layout encoder without changing declarative settings.

nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionConfig.grid(
    pixels: torch.Tensor
) -> tuple[int, int]

Validate a complete preprocessed image and return its merged grid.

Parameters:

pixels

torch.Tensor

Floating tensor [channels, height, width], before patch convolution. Height/width must be positive multiples of patch_size * merge_size.

Returns: tuple[int, int]

Merged grid height and width in language image tokens.

class nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionEncoder(
    config: nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionConfig,
    dtype: torch.dtype
)

Bases: Module

Pack variable images into one set of per-layer calls with native weight names.

ln_pre

patch_conv

transformer

= nn.Module()

nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionEncoder.forward(
    pixels: list[torch.Tensor]
) -> tuple[torch.Tensor, ...]

Encode complete images; positions reset for each image.

Parameters:

pixels

list[torch.Tensor]

Nonempty list of floating [channels, height, width] tensors, on the model device. Inputs are neither resized nor modified.

Returns: torch.Tensor

One independent [patch_rows * patch_columns, vision_hidden] tensor per

class nemo_automodel.components.models.mistral_large4.vision._VisionAttention(
    config: nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionConfig,
    dtype: torch.dtype
)

Bases: Module

head_dim

= config.hidden_size // self.heads

heads

= config.num_attention_heads

nemo_automodel.components.models.mistral_large4.vision._VisionAttention.forward(
    x: torch.Tensor,
    angles: torch.Tensor,
    keep: torch.Tensor
) -> torch.Tensor

Attend bidirectionally within each image.

Parameters:

torch.Tensor

Floating packed [tokens, hidden] image features.

angles

torch.Tensor

FP32 [tokens, head_dim / 2] axial angles.

keep

torch.Tensor

Boolean [tokens, tokens] image membership mask, true permits attention.

Returns: torch.Tensor

Independent projected features [tokens, hidden].

class nemo_automodel.components.models.mistral_large4.vision._VisionMLP(
    config: nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionConfig,
    dtype: torch.dtype
)

Bases: Module

nemo_automodel.components.models.mistral_large4.vision._VisionMLP.forward(
    x: torch.Tensor
) -> torch.Tensor

Apply the native SwiGLU projections.

Parameters:

torch.Tensor

Floating [tokens, hidden] features.

Returns: torch.Tensor

Independent features [tokens, hidden].

nemo_automodel.components.models.mistral_large4.vision._rotate(
    x: torch.Tensor,
    angles: torch.Tensor
) -> torch.Tensor

Apply axial rotation without rounding trigonometric tables to BF16.

Parameters:

torch.Tensor

Floating [tokens, heads, head_dim], adjacent real/imaginary channel pairs.

angles

torch.Tensor

FP32 [tokens, head_dim / 2], height frequencies followed by width frequencies.

Returns: torch.Tensor

Independent rotated tensor with x’s shape and dtype.