nemo_automodel.components.models.mistral_large4.multimodal

Image/text execution over the existing Large 4 training backbone.

Module Contents

Classes

Name	Description
`MistralLarge4ForConditionalGeneration`	Native multimodal SFT/LoRA with complete images and packed document isolation.

Functions

Name	Description
`image_token_positions`	Validate complete image grids independently inside each document.

API

class nemo_automodel.components.models.mistral_large4.multimodal.MistralLarge4ForConditionalGeneration(
    config: nemo_automodel.components.models.mistral_large4.config.MistralLarge4Config,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None
)

Bases: MistralLarge4ForCausalLM

Native multimodal SFT/LoRA with complete images and packed document isolation.

Shares the text backbone, loss, and checkpoint adapter without wrapping a second causal-model root. All vision checkpoint names stay native.

_no_split_modules

= ['MistralLarge4Block', 'MistralLarge4VisionBlock']

parallel_layer_groups

patch_merger

pre_mm_projector_norm

vision_encoder

= self.vision_settings.build(dtype=dtype)

vision_language_adapter

vision_settings

nemo_automodel.components.models.mistral_large4.multimodal.MistralLarge4ForConditionalGeneration.customize_pipeline_stage_modules(
    module_names_per_stage: list[list[str]],
    layers_prefix: str,
    text_model: torch.nn.Module | None = None
) -> list[list[str]]

Keep the native vision tower and all image projections on stage zero.

nemo_automodel.components.models.mistral_large4.multimodal.MistralLarge4ForConditionalGeneration.forward(
    input_ids: torch.Tensor | None = None,
    pipeline_tensors: torch.Tensor = (),
    pixel_values: list[torch.Tensor] | torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    labels: torch.Tensor | None = None,
    logits_to_keep: int | torch.Tensor = 0,
    output_hidden_states: bool = False,
    cu_seqlens: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    qkv_format: str = 'bshd',
    use_cache: bool = False,
    _mistral_packed_document_ids: torch.Tensor | None = None
) -> transformers.modeling_outputs.CausalLMOutputWithPast | torch.Tensor | tuple[torch.Tensor, ...]

Replace complete image grids and run the existing causal training path.

Parameters:

input_ids

torch.Tensor | NoneDefaults to None

Global integer [batch, sequence] or packed [tokens]. Required for images; under CP, complete images are spliced before the first stage shards the sequence. Later PP stages take local residuals.

pipeline_tensors

torch.TensorDefaults to ()

Propagated state in _pipeline_boundary_metas layouts; non-first stages receive [batch, sequence, hidden] residuals in input_ids and skip the image encoder.

pixel_values

list[torch.Tensor] | torch.Tensor | NoneDefaults to None

Optional list of floating [channels, height, width] images in encounter order.

inputs_embeds

torch.Tensor | NoneDefaults to None

Optional floating [batch, sequence, hidden] or [tokens, hidden], replacing the embedding lookup. IDs remain necessary for image alignment.

position_ids

torch.Tensor | NoneDefaults to None

Optional nonnegative integers matching token axes; CP sharder supplies [batch, local_sequence] global positions.

attention_mask

torch.Tensor | NoneDefaults to None

Binary keep mask or packed document IDs matching global token axes. CP combines both user padding masks here.

padding_mask

torch.Tensor | NoneDefaults to None

Optional boolean padding mask matching token axes; CP sharder supplies [batch, local_sequence], including CP padding.

labels

torch.Tensor | NoneDefaults to None

Optional integer targets matching token axes, -100 ignored; [batch, local_sequence] under CP. Image tokens and delimiters are masked on a copy before the existing shifted CE.

logits_to_keep

int | torch.TensorDefaults to 0

Zero for all, positive tail length or int64 [kept_positions].

output_hidden_states

boolDefaults to False

Return final hidden states in addition to logits.

cu_seqlens

torch.Tensor | NoneDefaults to None

Optional packed integer [documents + 1] boundaries.

seq_lens

torch.Tensor | NoneDefaults to None

Optional real document lengths [batch, documents], with -1000 padding, retained by the packed sharder for document-count validation.

qkv_format

strDefaults to 'bshd'

bshd or thd; THD requires cu_seqlens.

use_cache

boolDefaults to False

Must be false, as in the causal training path.

_mistral_packed_document_ids

torch.Tensor | NoneDefaults to None

Prepared integer [batch, local_sequence] document identity for prepared THD execution.

Returns: CausalLMOutputWithPast | torch.Tensor | tuple[torch.Tensor, ...]

CausalLMOutputWithPast with FP32 [batch, kept_sequence, vocab] capped logits,

nemo_automodel.components.models.mistral_large4.multimodal.MistralLarge4ForConditionalGeneration.get_image_features(
    pixel_values: list[torch.Tensor]
) -> tuple[torch.Tensor, ...]

Encode and project variable-size complete images in encounter order.

Parameters:

pixel_values

list[torch.Tensor]

Nonempty list of floating [channels, height, width] images on the model device.

Returns: tuple[torch.Tensor, ...]

Per-image floating [merged_rows * merged_columns, text_hidden] features.

nemo_automodel.components.models.mistral_large4.multimodal.MistralLarge4ForConditionalGeneration.prepare_model_inputs_for_cp(
    batch: nemo_automodel.components.models.mistral_large4.context_parallel.CPBatch,
    num_chunks: int = 1
) -> dict[str, nemo_automodel.components.distributed.context_parallel.sharder.ContextParallelSharder]

Keep complete image/token streams until stage-zero embedding.

Parameters:

batch

CPBatch

Global IDs/masks [batch, sequence], optional embeddings [batch, sequence, hidden], and complete images [channels, height, width].

num_chunks

intDefaults to 1

Number of independently scheduled PP microbatches.

Returns: dict[str, ContextParallelSharder]

Sharder mapping using prepare_cp_inputs’ aux-only contract: primary

nemo_automodel.components.models.mistral_large4.multimodal.image_token_positions(
    input_ids: torch.Tensor,
    pixels: list[torch.Tensor],
    config: nemo_automodel.components.models.mistral_large4.vision.MistralLarge4VisionConfig,
    cu_seqlens: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None
) -> torch.Tensor

Validate complete image grids independently inside each document.

Parameters:

input_ids

torch.Tensor

Integer [batch, sequence] or packed [tokens] storage.

pixels

list[torch.Tensor]

Images [channels, height, width], in document/encounter order.

config

MistralLarge4VisionConfig

Native vision dimensions and image delimiter IDs.

cu_seqlens

torch.Tensor | NoneDefaults to None

Optional integer [documents + 1] partition of packed storage.

attention_mask

torch.Tensor | NoneDefaults to None

Optional binary keep mask with input_ids’ shape.

padding_mask

torch.Tensor | NoneDefaults to None

Optional boolean padding mask with input_ids’ shape.

Returns: torch.Tensor

Int64 [total_merged_patches] indices into flattened tokens, on input_ids’