nemo_automodel.components.models.mistral_large4.state_dict_adapter

Native names and bounded E4M3 reconstruction using existing DCP load parts.

Matrices retain separate gate/up projections and individual expert identities. Grouped expert storage exports the same released identities through transposed views.

Module Contents

Classes

Name	Description
`MistralLarge4StateDictAdapter`	Lossless name mapping and bounded source-FP8 initialization.

Functions

Name	Description
`_decode_into`	Reconstruct into existing floating storage, using global scale coordinates.
`_finish_decodes`	Decode local FP8 expert matrices sequentially into existing model views.
`_finish_load_group`	Decode each FP8 matrix [out, in] sequentially into final model storage.
`_local_offsets`	Find global coordinates of a contiguous two-dimensional DTensor shard.
`_native_key`	Map a released text tensor name to its model-owned native parameter.
`_released_key`	Reverse the text mapping without allocating or reading tensor values.

Data

_AUXILIARY_PREFIXES

_LAYER_NAMES

_NAMES

_QUANTIZED

API

class nemo_automodel.components.models.mistral_large4.state_dict_adapter.MistralLarge4StateDictAdapter(
    dtype: torch.dtype = torch.bfloat16,
    grouped_experts: bool = False,
    n_routed_experts: int = 0
)

Bases: StateDictAdapter

Lossless name mapping and bounded source-FP8 initialization.

Individual and grouped experts export the same released gate/up/down matrices. Vision tensors retain their native names/layouts and load directly into the conditional model. DSpark draft tensors keep their released names and are not loaded. Native floating save/resume never attaches original FP8 scales. FSDP and balanced expert-subset loading are supported. Pipeline peers group local destinations into equal numbers of nonempty load parts, retaining bounded FP8 staging.

_grouped_load_views

dict[str, tuple[Tensor, dict[str, Tensor], dict[str, Tensor]]] = {}

nemo_automodel.components.models.mistral_large4.state_dict_adapter.MistralLarge4StateDictAdapter._group_experts(
    state: dict[str, torch.Tensor],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None
) -> dict[str, torch.Tensor]

Fuse complete per-layer expert banks for the in-memory import path.

Parameters:

state

dict[str, torch.Tensor]

Floating native matrices [out, in]. Expert inputs must be complete unsharded matrices; the streaming loader supports sharded destinations without making this in-memory copy.

device_mesh

DeviceMesh | None

Optional MoE mesh; only the EP axis participates here.

Returns: dict[str, torch.Tensor]

New mapping with grouped [E, hidden, 2*intermediate] gate/up and

nemo_automodel.components.models.mistral_large4.state_dict_adapter.MistralLarge4StateDictAdapter._iter_individual_load_parts(
    model_state_dict: dict[str, torch.Tensor]
) -> collections.abc.Iterator[nemo_automodel.components.checkpoint.state_dict_adapter.CheckpointLoadPart]

Yield one parameter per part, staging at most one local expert bank.

Parameters:

model_state_dict

dict[str, torch.Tensor]

Native tensors; linear matrices [out, in], grouped gate/up [experts, hidden, 2*intermediate] and down [experts, intermediate, hidden]. EP shards axis 0 and expert FSDP shards axis 1. Other tensors retain checkpoint axes.

Returns: Iterator[CheckpointLoadPart]

Parts whose FP8 staging covers only the local parameter shard.

nemo_automodel.components.models.mistral_large4.state_dict_adapter.MistralLarge4StateDictAdapter._local_grouped_checkpoint_views(
    name: str,
    tensor: torch.Tensor
) -> dict[str, torch.Tensor]

Keep original local storage views even when distributed export changes device.

Parameters:

name

str

Native grouped gate/up or down parameter name.

tensor

torch.Tensor

Grouped [E, hidden, 2*intermediate] or [E, intermediate, hidden] tensor, optionally EP/FSDP sharded and CPU offloaded.

Returns: dict[str, torch.Tensor]

Released per-expert names to plain local [out, in] views. Their

nemo_automodel.components.models.mistral_large4.state_dict_adapter.MistralLarge4StateDictAdapter.convert_single_tensor_to_hf(
    fqn: str,
    tensor: torch.Tensor,
    quantization: bool = False
) -> list[tuple[str, torch.Tensor]]

Convert one floating tensor for checkpoint import or export.

Parameters:

fqn

str

Model parameter name.

tensor

torch.Tensor

Floating tensor with native axes, [output, input] for matrices; grouped gate/up [E, hidden, 2*intermediate], down [E, intermediate, hidden]. EP shards axis 0; expert FSDP shards axis 1. Splits/transposes alias storage. Floating exports preserve DTensor layouts.

quantization

boolDefaults to False

Unsupported; exports preserve the floating training dtype.

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

Native (name, tensor) pairs whose values alias the input.

nemo_automodel.components.models.mistral_large4.state_dict_adapter.MistralLarge4StateDictAdapter.forced_hf_dtype_mapping(
    state_dict: dict[str, torch.Tensor]
) -> dict[str, str]

Prevent export from recasting decoded weights to the source FP8 dtype.

Parameters:

state_dict

dict[str, torch.Tensor]

Already-exported tensors retaining checkpoint axes, matrices [output, input], optionally DTensors. Only dtype is inspected.

Returns: dict[str, str]

Released-key dtype overrides for every floating tensor. FP32 master

nemo_automodel.components.models.mistral_large4.state_dict_adapter.MistralLarge4StateDictAdapter.from_hf(
    hf_state_dict: dict[str, torch.Tensor],
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None
) -> dict[str, torch.Tensor]

Convert a dependency-complete source mapping without consuming it.

Parameters:

hf_state_dict

dict[str, torch.Tensor]

Native checkpoint tensors, with weights [output, input] and FP32 scale grids [ceil(output/128), ceil(input/128)]. BF16 tensors retain checkpoint axes and may alias the returned values. FP8 matrices may be regular DTensor shards with global shapes.

device_mesh

DeviceMesh | NoneDefaults to None

Existing loader argument; DTensors own their placements.

Returns: dict[str, torch.Tensor]

Model names with matching axes, or fused grouped gate/up

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

Yield bounded destination groups through the existing checkpoint loader hook.

Parameters:

model_state_dict

dict[str, torch.Tensor]

Model names mapped to final tensors; linear weights use [output, input]; grouped gate/up uses [E, hidden, 2*intermediate] and down [E, intermediate, hidden]. Other tensors keep checkpoint axes. DTensors may shard either matrix axis; every rank must iterate the same number of parts. Nonquantized values alias final model storage. Quantized values use local FP8 staging.

device_mesh

DeviceMesh | NoneDefaults to None

Existing loader argument; each DTensor carries its mesh.

Returns: Iterator[CheckpointLoadPart]

Parts covering every model key exactly once. PP groups at most

nemo_automodel.components.models.mistral_large4.state_dict_adapter.MistralLarge4StateDictAdapter.map_peft_target_module_to_hf(
    name: str,
    v4_compatible: bool = False
) -> str

Use the same projection mapping in adapter metadata and tensor export.

nemo_automodel.components.models.mistral_large4.state_dict_adapter.MistralLarge4StateDictAdapter.to_hf(
    state_dict: dict[str, torch.Tensor],
    exclude_key_regex: str | None = None,
    quantization: bool = False,
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    v4_compatible: bool = False,
    legacy_paramwrapper_layout: bool = False,
    for_checkpoint_load: bool = False
) -> dict[str, torch.Tensor]

Export floating training weights with released checkpoint names.

Parameters:

state_dict

dict[str, torch.Tensor]

Model tensors with checkpoint axes; matrix [output, input] or grouped experts [E, hidden, 2*intermediate] / [E, intermediate, hidden], optionally EP Shard(0) and expert-FSDP Shard(1) DTensors.

exclude_key_regex

str | NoneDefaults to None

Existing checkpoint API filter on model names.

quantization

boolDefaults to False

Unsupported; source FP8 is decoded on import, training saves floating weights.

device_mesh

DeviceMesh | NoneDefaults to None

Existing export hook; DTensors carry their EP ownership.

v4_compatible

boolDefaults to False

Existing hook; released names have one layout.

legacy_paramwrapper_layout

boolDefaults to False

Existing hook; this model has no grouped ParamWrapper.

for_checkpoint_load

boolDefaults to False

Existing hook; floating destinations always alias model storage.

Returns: dict[str, torch.Tensor]

Floating tensors alias inputs and retain their training dtype.

nemo_automodel.components.models.mistral_large4.state_dict_adapter._decode_into(
    weight: torch.Tensor,
    scale: torch.Tensor,
    destination: torch.Tensor
) -> None

Reconstruct into existing floating storage, using global scale coordinates.

FP32 conversion and elementwise scale temporaries cover at most 1024x1024 values each. FP8 source values are cast to FP32, scaled, then cast once into destination storage. Neither source input is modified.

Parameters:

weight

torch.Tensor

E4M3 matrix [global_output, global_input], local or a regular Shard/Replicate DTensor. No projection transpose has occurred.

scale

torch.Tensor

Replicated FP32 grid [ceil(global_output/128), ceil(global_input/128)] on the weight’s device, holding positive finite numeric multipliers.

destination

torch.Tensor

BF16/FP32 matrix with the same global shape and placements as weight. Written in place without changing parameter identity. The local tensor may be a noncontiguous view.

nemo_automodel.components.models.mistral_large4.state_dict_adapter._finish_decodes(
    decodes: list[collections.abc.Callable[[], None]]
) -> None

Decode local FP8 expert matrices sequentially into existing model views.

nemo_automodel.components.models.mistral_large4.state_dict_adapter._finish_load_group(
    parts: list[nemo_automodel.components.checkpoint.state_dict_adapter.CheckpointLoadPart]
) -> None

Decode each FP8 matrix [out, in] sequentially into final model storage.

nemo_automodel.components.models.mistral_large4.state_dict_adapter._local_offsets(
    tensor: torch.distributed.tensor.DTensor
) -> tuple[int, int]

Find global coordinates of a contiguous two-dimensional DTensor shard.

Parameters:

tensor

DTensor

Matrix [global_output, global_input], with regular Shard or Replicate placements on any mesh axes. Partial and strided shards are unsupported. Repeated sharding of either axis is allowed.

Returns: tuple[int, int]

Global output and input offsets of the rank-local matrix.

nemo_automodel.components.models.mistral_large4.state_dict_adapter._native_key(
    key: str
) -> str

Map a released text tensor name to its model-owned native parameter.

nemo_automodel.components.models.mistral_large4.state_dict_adapter._released_key(
    key: str
) -> str

Reverse the text mapping without allocating or reading tensor values.

nemo_automodel.components.models.mistral_large4.state_dict_adapter._AUXILIARY_PREFIXES = ('vision_encoder.', 'vision_language_adapter.', 'patch_merger.', 'pre_mm_project...

nemo_automodel.components.models.mistral_large4.state_dict_adapter._LAYER_NAMES = (('attention_norm.', 'input_layernorm.'), ('ffn_norm.', 'post_attention_layernor...

nemo_automodel.components.models.mistral_large4.state_dict_adapter._NAMES = (('tok_embeddings.', 'model.embed_tokens.'), ('output.', 'lm_head.'), ('lerope.'...

nemo_automodel.components.models.mistral_large4.state_dict_adapter._QUANTIZED = re.compile('layers\\.\\d+\\.(?:attention\\.(?:wq|wo|indexer\\.wq_b)|(?:feed_forw...