nemo_automodel.components.models.mistral3_vlm.state_dict_adapter

Compatibility export for the shared Mistral3 FP8 state-dict adapter.

Module Contents

Classes

Name	Description
`Mistral3FP8StateDictAdapter`	Per-tensor FP8 dequant adapter for the Mistral3 model family.

API

class nemo_automodel.components.models.mistral3.state_dict_adapter.Mistral3FP8StateDictAdapter(
    native_to_hf: typing.Callable[[str], str] = _identity,
    hf_to_native: typing.Callable[[str], str] = _identity,
    layout_name: str = 'vlm_full',
    not_fp8_prefixes: tuple[str, ...] = (),
    num_hidden_layers: int | None = None
)

Bases: StateDictAdapter

Per-tensor FP8 dequant adapter for the Mistral3 model family.

Text-only causal-LM checkpoint keys already match the model state dict. VLM checkpoints additionally select the appropriate body-key layout and exclude their BF16 vision and projector modules from FP8 conversion.

_not_fp8_prefixes

= tuple(not_fp8_prefixes)

nemo_automodel.components.models.mistral3.state_dict_adapter.Mistral3FP8StateDictAdapter._iter_checkpoint_load_parts(
    model_state_dict: dict[str, torch.Tensor],
    decoder_layer_key: re.Pattern[str]
) -> collections.abc.Iterator[nemo_automodel.components.checkpoint.state_dict_adapter.CheckpointLoadPart]

Build load parts for a complete causal-LM or VLM decoder.

Parameters:

model_state_dict

dict[str, torch.Tensor]

Native model names mapped to final tensors of arbitrary model-defined rank and shape. Quantized linear weights must use BF16 or FP32 final storage; non-quantized tensors remain direct DCP destinations and are mutated in place during the load.

decoder_layer_key

re.Pattern[str]

Pattern that identifies decoder-layer names and captures the zero-based layer index in group 2. It must exclude vision-tower layers from the bounded FP8 groups.

Returns: Iterator[CheckpointLoadPart]

Dependency-complete load parts. Each quantized destination has the same shape, strides, device, and DTensor

nemo_automodel.components.models.mistral3.state_dict_adapter.Mistral3FP8StateDictAdapter.convert_single_tensor_to_hf(
    fqn: str,
    tensor: typing.Any,
    kwargs = {}
) -> list[tuple[str, typing.Any]]

Per-tensor model → HF used by Checkpointer.save_model.

nemo_automodel.components.models.mistral3.state_dict_adapter.Mistral3FP8StateDictAdapter.for_causal_lm(
    config: typing.Any | None = None
) -> 'Mistral3FP8StateDictAdapter'

classmethod

Text-only path for per-tensor FP8 Ministral3ForCausalLM checkpoints.

Devstral-2 stores text-model keys in the same layout exposed by the local Ministral3ForCausalLM implementation. Linear weights are FP8, while embeddings, norms, and the untied LM head remain BF16 and are excluded by _NON_QUANTIZED_SUFFIXES.

nemo_automodel.components.models.mistral3.state_dict_adapter.Mistral3FP8StateDictAdapter.for_vlm_full(
    config: typing.Any | None = None
) -> 'Mistral3FP8StateDictAdapter'

classmethod

Full-VLM path for Mistral3ForConditionalGeneration checkpoints.

Mistral3 FP8 VLM checkpoints have two observed body-key layouts. The Mistral-Medium-3.5 128B checkpoint already stores keys in the same layout as HF’s VLM state_dict() (model.language_model.* / model.vision_tower.* / model.multi_modal_projector.*). Newer Ministral/Devstral-style checkpoints store text weights under language_model.model.* and non-text component names at top level.

The LM head has one extra quirk in the nested layout: the model exposes it at the top level (lm_head.weight) while the checkpoint nests it (language_model.lm_head.weight). Tied checkpoints (Ministral-3) never serialize the head, so the head translation is a harmless no-op there; untied checkpoints (Devstral-24B) rely on it to find the head during the DCP load.

Only the language_model layer weights are FP8; vision / mm_projector / lm_head are BF16 on disk and must be passed through without a scale_inv placeholder — otherwise DCP would fail trying to fetch a non-existent _scale_inv key.

nemo_automodel.components.models.mistral3.state_dict_adapter.Mistral3FP8StateDictAdapter.from_hf(
    hf_state_dict: dict[str, typing.Any],
    device_mesh: 'DeviceMesh' | None = None,
    kwargs = {}
) -> dict[str, typing.Any]

Convert an HF-format (possibly FP8) state dict to model-native format.

nemo_automodel.components.models.mistral3.state_dict_adapter.Mistral3FP8StateDictAdapter.iter_checkpoint_load_parts(
    model_state_dict: dict[str, torch.Tensor],
    device_mesh: 'DeviceMesh' | None = None
) -> collections.abc.Iterator[nemo_automodel.components.checkpoint.state_dict_adapter.CheckpointLoadPart] | None

Load Mistral3 FP8 weights in bounded decoder-layer groups.

Each quantized model tensor has BF16 or FP32 model shape and storage. Its load part creates an FP8 destination with the same shape, device, strides, and distributed placements, plus the scalar BF16 _scale_inv destination stored by the checkpoint. After DCP fills both tensors, the part copies and scales the FP8 value directly into the final model tensor. Non-quantized tensors, including VLM vision and projector weights, use their final model storage as the DCP destination.

This path requires the complete decoder on every rank. Pipeline-parallel ranks own different layer subsets and therefore retain the existing rank-local DCP path until part scheduling can be coordinated across stages. Tied VLM checkpoints also retain that path because they omit the LM-head tensor expected by grouped loading.

Parameters:

model_state_dict

dict[str, torch.Tensor]

Native model names mapped to final parameter and persistent-buffer tensors. Each tensor has arbitrary model-defined rank and shape. Decoder tensors may be local DTensor shards; their global shapes, local shards, and placements are preserved by torch.empty_like. Non-quantized destinations alias and are populated through final model storage.

device_mesh

'DeviceMesh' | NoneDefaults to None

Optional distributed mesh. The tensors already carry their final placements, so this value is not otherwise needed.

Returns: Iterator[CheckpointLoadPart] | None

One direct-load part for tensors outside decoder layers plus bounded temporary-load parts for decoder

nemo_automodel.components.models.mistral3.state_dict_adapter.Mistral3FP8StateDictAdapter.to_hf(
    state_dict: dict[str, typing.Any],
    exclude_key_regex: str | None = None,
    quantization: bool = False,
    kwargs = {}
) -> dict[str, typing.Any]

Convert a model-native state dict to HF (on-disk) layout.

When quantization=True the weight placeholder is also cast to torch.float8_e4m3fn so the DCP storage reader fetches FP8 bytes verbatim from safetensors (a bf16 target would silently cast-on-read and lose the scale multiply — see deepseek_v3/state_dict_adapter.py:220). A scalar _scale_inv placeholder is also emitted so DCP pulls it alongside the weight.