nemo_automodel.components.models.mistral3_vlm

FP8-native Mistral3 VLM custom model registration (dawn-ridge 128B).

Importing this module installs a resolver hook on _resolve_custom_model_cls_for_config (nemo_automodel/_transformers/ model_init.py) that routes FP8-native Mistral3 VLM configs to Mistral3FP8VLMForConditionalGeneration. The model registry is not overwritten, so non-FP8 Mistral3 VLMs keep the stock mistral4.model path and there is no regression for other users.

Submodules

Package Contents

Classes

Name	Description
`Mistral3FP8VLMForConditionalGeneration`	Full-VLM (vision + text) FP8 loader for Mistral3ForConditionalGeneration.

Functions

Name	Description
`_install_resolver_hook`	Prepend an FP8-Mistral3 VLM check to _resolve_custom_model_cls_for_config.

Data

logger

API

class nemo_automodel.components.models.mistral3_vlm.model.Mistral3FP8VLMForConditionalGeneration(
    config: transformers.PretrainedConfig
)

Bases: _HFMistral3ForConditionalGeneration

Full-VLM (vision + text) FP8 loader for Mistral3ForConditionalGeneration.

Used when the user instantiates through NeMoAutoModelForImageTextToText.from_pretrained on an FP8-native Mistral3 VLM checkpoint (e.g. dawn-ridge-128B).

_tied_weights_keys

state_dict_adapter

= Mistral3FP8StateDictAdapter.for_vlm_full(config)

tie_word_embeddings_support

TieSupport = TieSupport.BOTH

nemo_automodel.components.models.mistral3_vlm.model.Mistral3FP8VLMForConditionalGeneration.forward(
    input_ids: torch.LongTensor | None = None,
    pixel_values: torch.FloatTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values = None,
    inputs_embeds: torch.FloatTensor | None = None,
    labels: torch.LongTensor | None = None,
    use_cache: bool | None = None,
    logits_to_keep: typing.Union[int, torch.Tensor] = 0,
    image_sizes: torch.Tensor | None = None,
    output_hidden_states: bool | None = None,
    kwargs = {}
) -> transformers.models.mistral3.modeling_mistral3.Mistral3CausalLMOutputWithPast

Forward pass with memory-efficient fused cross-entropy (cut-CE) support.

Overrides HF’s Mistral3ForConditionalGeneration.forward so the train_ft recipe can enable FusedLinearCrossEntropy. The recipe only does so when (a) forward exposes a logits_to_keep parameter and (b) calling the model returns an output that carries the FINAL hidden states (full sequence) while logits cover only the kept positions.

HF’s stock forward gates hidden_states on a per-call output_hidden_states kwarg (which the recipe does not pass) and emits the full per-layer tuple. Here we instead resolve output_hidden_states from the text sub-config and surface the inner model’s last_hidden_state (the single [B, S, H] tensor fed to lm_head) directly, which is what get_final_hidden_states consumes.

Parameters:

input_ids

torch.LongTensor | NoneDefaults to None

Input token IDs [B, S].

pixel_values

torch.FloatTensor | NoneDefaults to None

Optional image pixel values for the vision tower.

attention_mask

torch.Tensor | NoneDefaults to None

Optional attention mask.

position_ids

torch.LongTensor | NoneDefaults to None

Optional position indices.

past_key_values

Defaults to None

Optional cached key/values.

inputs_embeds

torch.FloatTensor | NoneDefaults to None

Optional pre-computed embeddings.

labels

torch.LongTensor | NoneDefaults to None

Optional labels for loss computation.

use_cache

bool | NoneDefaults to None

Whether to use KV caching.

logits_to_keep

Union[int, torch.Tensor]Defaults to 0

Number of final logits to compute (0=all, N=last N tokens).

image_sizes

torch.Tensor | NoneDefaults to None

Optional image sizes for the vision tower.

output_hidden_states

bool | NoneDefaults to None

Whether to surface the final hidden states on the output (defaults to the text sub-config’s output_hidden_states).

**kwargs

Defaults to {}

Additional arguments forwarded to the base model.

Returns: Mistral3CausalLMOutputWithPast

class:~transformers.models.mistral3.modeling_mistral3.Mistral3CausalLMOutputWithPast

nemo_automodel.components.models.mistral3_vlm.model.Mistral3FP8VLMForConditionalGeneration.supports_config(
    config: transformers.PretrainedConfig
) -> bool

classmethod

Claim FP8-native and dequantized Mistral3 VLM configs.

Matches Mistral3Config (outer VLM) with a ministral3 text backbone and either no quantization config or quantization_config.quant_method == 'fp8'. Consolidated checkpoints are saved in BF16 without the source FP8 metadata and still require this class’s rotary-buffer reinit.

nemo_automodel.components.models.mistral3_vlm.model.Mistral3FP8VLMForConditionalGeneration.tie_weights(
    _args: object = (),
    _kwargs: object = {}
) -> None

Tie lm_head to the active text embedding when requested.

nemo_automodel.components.models.mistral3_vlm._install_resolver_hook() -> None

Prepend an FP8-Mistral3 VLM check to _resolve_custom_model_cls_for_config.

nemo_automodel.components.models.mistral3_vlm.logger = logging.getLogger(__name__)