nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter

State-dict adapter for the MiniMax M3 (text) backbone.

Converts between the released HF checkpoint layout and the native AutoModel layout:

block_sparse_moe.{gate,e_score_correction_bias} -> mlp.gate.*
block_sparse_moe.experts.{e}.{w1,w3,w2} -> grouped mlp.experts.* (gate/up/down) via MoESplitExpertsStateDictMixin
block_sparse_moe.shared_experts.* -> shared_experts.* (a sibling of mlp on the decoder block)
dense (non-MoE) layers keep mlp.{gate,up,down}_proj.* unchanged

MXFP8 weights (FP8 e4m3 with e8m0/uint8 block scales along the input dim) are dequantized to the training dtype on load. M3.1 adds NVFP4 packed routed experts and MXFP8 indexer projections. Its native load-part iterator writes directly into final sharded storage while bounding temporary memory to one grouped expert parameter or ordinary projection.

Module Contents

Classes

Name	Description
`MiniMaxM3StateDictAdapter`	Convert MiniMax M3 HF checkpoints to/from the native grouped-expert format.
`MiniMaxM3VLStateDictAdapter`	VLM adapter: splits the M3 VL checkpoint into text / vision / projector parts.
`_QuantizedProjection`	Temporary checkpoint tensors plus the final projection view they fill.

Functions

Name	Description
`_dequantize_mxfp8_local`	-
`_dequantize_nvfp4`	Decode the compressed-tensors NVFP4 projection layout used by M3.1.
`_finish_quantized_projections`	Decode each projection separately, bounding unpacked temporary storage.
`_quantized_destination`	Allocate a compressed checkpoint destination matching a weight shard.
`_should_quantize_mxfp8_key`	True for HF-format weight keys stored as MXFP8 in the checkpoint.
`_slice_mxfp8_scale_for_dtensor`	Slice a global scale_inv to a DTensor weight’s local shard.
`_uses_nvfp4_experts`	Recognize the released M3.1 mixed MXFP8/NVFP4 checkpoint layout.
`create_mxfp8_scale_inv`	Load-time placeholder scale_inv (e8m0/uint8, GLOBAL shape `[out, ceil(in/block)]`).
`dequantize_mxfp8`	Dequantize an MXFP8 weight (FP8 e4m3 + e8m0/uint8 block scales) to `dtype`.

Data

MXFP8_BLOCK_SIZE

_MXFP8_INDEXER_RE

_MXFP8_QUANT_KEY_RE

_MXFP8_SCALE_INV_IDENTITY

_NVFP4_BLOCK_SIZE

_NVFP4_EXPERT_RE

_NVFP4_QUANTIZATION_CONFIG

_NVFP4_VALUES

API

class nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3StateDictAdapter(
    config: nemo_automodel.components.models.minimax_m3_vl.config.MiniMaxM3VLTextConfig,
    moe_config: nemo_automodel.components.moe.layers.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype = torch.bfloat16
)

Bases: MoESplitExpertsStateDictMixin, StateDictAdapter

Convert MiniMax M3 HF checkpoints to/from the native grouped-expert format.

_expert_path_segment

str

_mtp_enabled

bool

_nvfp4_experts

= _uses_nvfp4_experts(config)

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3StateDictAdapter._checkpoint_projection_views(
    fqn: str,
    tensor: torch.Tensor
) -> list[tuple[str, torch.Tensor, torch.Tensor]]

Separate checkpoint metadata from final local projection views.

Parameters:

fqn

str

Native model state-dict name.

tensor

torch.Tensor

Final model tensor. Grouped experts use global [experts, input, 2 * intermediate] gate/up or [experts, intermediate, output] down layouts. EP shards the expert axis; FSDP may additionally shard the input axis.

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

Tuples of HF name, metadata template, and final local destination.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3StateDictAdapter._dequantize(
    state_dict: dict[str, torch.Tensor]
) -> dict[str, torch.Tensor]

Decode quantized projections in place before grouped-expert merging.

Parameters:

state_dict

dict[str, torch.Tensor]

HF projection weights and quantization companions. MXFP8 uses [output, input] FP8 weights and [output, input / 32] exponent scales. NVFP4 uses [output, input / 2] packed bytes, [output, input / 16] E4M3 scales, and an FP32 global scalar. Distributed tensors retain their matching mesh and placements.

Returns: dict[str, torch.Tensor]

The same mapping, with floating-point [output, input] weights.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3StateDictAdapter._hf_key_to_native(
    key: str
) -> str

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3StateDictAdapter._iter_mixed_checkpoint_parts(
    model_state_dict: dict[str, torch.Tensor],
    hf_prefix: str = ''
) -> collections.abc.Iterator[nemo_automodel.components.checkpoint.state_dict_adapter.CheckpointLoadPart]

Yield bounded M3.1 load parts over final native model storage.

Parameters:

model_state_dict

dict[str, torch.Tensor]

Native names and final CPU/CUDA/DTensor storage with layouts documented by iter_checkpoint_load_parts.

hf_prefix

strDefaults to ''

HF wrapper namespace prepended to text checkpoint keys.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3StateDictAdapter._mtp_from_hf(
    mtp_keys: dict[str, typing.Any],
    device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None
) -> dict[str, typing.Any]

Convert MTP tensors through a temporary one-layer text namespace.

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

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3StateDictAdapter._native_key_to_hf(
    key: str
) -> str

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3StateDictAdapter._prepare_hf_for_merge(
    state_dict: dict[str, torch.Tensor]
) -> None

Prepare Hugging Face projection tensors for grouped-expert merging.

Parameters:

state_dict

dict[str, torch.Tensor]

Mapping mutated in place. Expert projection tensors have shape [expert_hidden, hidden] for gate/up and [hidden, expert_hidden] for down; MXFP8 projection tensors are dequantized before their keys are renamed.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3StateDictAdapter.convert_single_tensor_to_hf(
    fqn: str,
    tensor: torch.Tensor,
    kwargs = {}
) -> list[tuple[str, torch.Tensor]]

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3StateDictAdapter.from_hf(
    hf_state_dict: dict[str, torch.Tensor],
    device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None,
    kwargs = {}
) -> dict[str, torch.Tensor]

Convert an HF checkpoint to native format (operates in-place to limit peak memory).

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3StateDictAdapter.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] | None

Load M3.1 mixed quantized weights into bounded model-owned parts.

Parameters:

model_state_dict

dict[str, torch.Tensor]

Final native storage, including global grouped expert tensors [experts, input, output]. DTensor expert-axis and input-axis shards, and CPU-offloaded local storage, are preserved. The iterator never creates a full model copy.

device_mesh

DeviceMesh | NoneDefaults to None

Existing expert mesh; placements are read from each destination DTensor.

Returns: Iterator[CheckpointLoadPart] | None

One direct-load part followed by one part per quantized native

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3StateDictAdapter.to_hf(
    state_dict: dict[str, torch.Tensor],
    exclude_key_regex: str | None = None,
    quantization: bool = False,
    kwargs = {}
) -> dict[str, torch.Tensor]

class nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3VLStateDictAdapter(
    config: nemo_automodel.components.models.minimax_m3_vl.config.MiniMaxM3VLConfig,
    moe_config: nemo_automodel.components.moe.layers.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype = torch.bfloat16
)

Bases: StateDictAdapter

VLM adapter: splits the M3 VL checkpoint into text / vision / projector parts.

The released checkpoint stores the language backbone under language_model.model.* / language_model.lm_head and the vision side under vision_tower.vision_model.* with the projector / patch-merger at top level (multi_modal_projector.* / patch_merge_mlp.*). The native VLM keeps the text model at model.* / lm_head and nests the projector / merger under vision_tower.*. Text tensors are delegated to MiniMaxM3StateDictAdapter (block_sparse_moe -> mlp, index branch, MXFP8 dequant, grouped experts); vision tensors are BF16 and pass through.

text_adapter

view_loaded_native_keys

set[str]

Native text keys loaded through the inner adapter’s checkpoint views.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3VLStateDictAdapter._iter_mixed_checkpoint_parts(
    model_state_dict: dict[str, torch.Tensor]
) -> collections.abc.Iterator[nemo_automodel.components.checkpoint.state_dict_adapter.CheckpointLoadPart]

Yield direct vision storage and bounded text parameter load parts.

Parameters:

model_state_dict

dict[str, torch.Tensor]

Final native tensors with the layout and alias contract documented by iter_checkpoint_load_parts.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3VLStateDictAdapter._map_non_text_from_hf(
    key: str
) -> str | None

staticmethod

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3VLStateDictAdapter._map_non_text_to_hf(
    key: str
) -> str

staticmethod

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3VLStateDictAdapter.convert_single_tensor_to_hf(
    fqn: str,
    tensor: torch.Tensor,
    kwargs = {}
) -> list[tuple[str, torch.Tensor]]

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3VLStateDictAdapter.from_hf(
    hf_state_dict: dict[str, torch.Tensor],
    device_mesh = None,
    kwargs = {}
) -> dict[str, typing.Any]

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3VLStateDictAdapter.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] | None

Stream text quantization while loading native vision storage directly.

Parameters:

model_state_dict

dict[str, torch.Tensor]

Final model tensors. Text layouts follow the text adapter’s iter_checkpoint_load_parts contract; vision tensors retain their arbitrary native shapes and DTensor placements. DCP writes every destination in place.

device_mesh

DeviceMesh | NoneDefaults to None

Existing expert mesh; tensor placements are retained.

Returns: Iterator[CheckpointLoadPart] | None

Complete, bounded M3.1 load parts, or None for earlier checkpoints.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MiniMaxM3VLStateDictAdapter.to_hf(
    state_dict: dict[str, typing.Any],
    exclude_key_regex = None,
    quantization: bool = False,
    kwargs = {}
)

class nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._QuantizedProjection(
    destination: torch.Tensor,
    weight: torch.Tensor,
    scale: torch.Tensor,
    global_scale: torch.Tensor | None = None
)

Dataclass

Temporary checkpoint tensors plus the final projection view they fill.

destination

Tensor

Final floating-point [output, input] projection view. For experts this aliases the original rank-local CPU/CUDA storage; ordinary projections may retain their global DTensor placements.

global_scale

Tensor | None = None

Optional replicated FP32 scalar for NVFP4. It divides the block scales; None denotes MXFP8.

scale

Tensor

Uint8 [output, input / 32] MXFP8 exponents or E4M3 [output, input / 16] NVFP4 scales, matching weight’s placements.

weight

Tensor

FP8 [output, input] or packed uint8 [output, input / 2] checkpoint staging tensor, with the source DTensor placements.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._QuantizedProjection.finish() -> None

Decode one [output, input] projection into existing local storage.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._dequantize_mxfp8_local(
    w_local: torch.Tensor,
    scale_local: torch.Tensor,
    block_size: int,
    dtype
) -> torch.Tensor

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._dequantize_nvfp4(
    packed: torch.Tensor,
    scales: torch.Tensor,
    global_scale: torch.Tensor,
    dtype: torch.dtype
) -> torch.Tensor

Decode the compressed-tensors NVFP4 projection layout used by M3.1.

Reference

https://github.com/vllm-project/compressed-tensors/blob/e69c8dc58aa152e5f5e36e85800ec1d2e7de5271/src/compressed_tensors/quantization/lifecycle/forward_helpers.py

Parameters:

packed

torch.Tensor

Uint8 projection of shape [output, input / 2]. Even input columns use the low nibble, odd columns the high nibble.

scales

torch.Tensor

E4M3 block scales of shape [output, input / 16].

global_scale

torch.Tensor

Positive FP32 scalar in the reciprocal convention.

dtype

torch.dtype

Floating-point output dtype. Computation takes place in FP32.

Returns: torch.Tensor

Independent tensor of shape [output, input], on packed’s device.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._finish_quantized_projections(
    projections: tuple[nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._QuantizedProjection, ...]
) -> None

Decode each projection separately, bounding unpacked temporary storage.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._quantized_destination(
    weight: torch.Tensor,
    divisor: int,
    dtype: torch.dtype,
    device: torch.device | None = None
) -> torch.Tensor

Allocate a compressed checkpoint destination matching a weight shard.

Parameters:

weight

torch.Tensor

Final floating-point projection of global shape [output, input]. A DTensor may shard either axis; every input-axis shard must align to divisor. Its local storage may be CPU or CUDA and noncontiguous.

divisor

int

Compression factor for the input axis, such as 2, 16, or 32.

dtype

torch.dtype

Checkpoint storage dtype.

device

torch.device | NoneDefaults to None

Materialization device for a meta template. A CUDA mesh uses CUDA staging even when final model storage is CPU offloaded.

Returns: torch.Tensor

Independent destination of global shape [output, input / divisor],

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._should_quantize_mxfp8_key(
    key: str
) -> bool

True for HF-format weight keys stored as MXFP8 in the checkpoint.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._slice_mxfp8_scale_for_dtensor(
    scale_inv: torch.Tensor,
    weight_dtensor: torch.Tensor,
    weight_local: torch.Tensor,
    block_size: int
) -> torch.Tensor

Slice a global scale_inv to a DTensor weight’s local shard.

MXFP8 block is [1, block_size]: dim 0 (out) is full-resolution (block 1, so a row range maps 1:1) and dim 1 (in) is grouped by block_size. Custom MoE is always tp=1, so sharding is on dim 0 (FSDP / ep_shard); dim 1 handled for safety.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._uses_nvfp4_experts(
    config: nemo_automodel.components.models.minimax_m3_vl.config.MiniMaxM3VLConfig | nemo_automodel.components.models.minimax_m3_vl.config.MiniMaxM3VLTextConfig
) -> bool

Recognize the released M3.1 mixed MXFP8/NVFP4 checkpoint layout.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.create_mxfp8_scale_inv(
    weight: torch.Tensor,
    block_size: int = MXFP8_BLOCK_SIZE
) -> torch.Tensor

Load-time placeholder scale_inv (e8m0/uint8, GLOBAL shape [out, ceil(in/block)]).

Emitted by to_hf(quantization=True) so the DCP planner requests the checkpoint’s *_scale_inv tensors; the values here are overwritten by the load. Kept a regular (non-DTensor) tensor with global shape — the per-shard slice happens in dequantize_mxfp8 (mirrors deepseek_v3).

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.dequantize_mxfp8(
    weight: torch.Tensor,
    scale_inv: torch.Tensor,
    block_size: int = MXFP8_BLOCK_SIZE,
    dtype: torch.dtype = torch.bfloat16
) -> torch.Tensor

Dequantize an MXFP8 weight (FP8 e4m3 + e8m0/uint8 block scales) to dtype.

weight is FP8 e4m3 [out, in]; scale_inv holds e8m0 (uint8) exponents [out, ceil(in/block_size)] with dequant scale for input-block b = 2 ** (scale_inv[:, b] - 127) (MX e8m0; confirmed vs sglang). Handles DTensor weights: the local shard is dequantized against the matching slice of a global scale_inv and rewrapped with the weight’s placements.

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter.MXFP8_BLOCK_SIZE = 32

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._MXFP8_INDEXER_RE = re.compile('\\.self_attn\\.index_[qk]_proj\\.weight$')

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._MXFP8_QUANT_KEY_RE = re.compile('\\.layers\\.\\d+\\.(?:self_attn\\.[qkvo]_proj|mlp\\.(?:gate|up|down)...

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._MXFP8_SCALE_INV_IDENTITY = 127

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._NVFP4_BLOCK_SIZE = 16

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._NVFP4_EXPERT_RE = re.compile('\\.block_sparse_moe\\.experts\\.\\d+\\.w[123]\\.weight$')

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._NVFP4_QUANTIZATION_CONFIG = MiniMaxM3QuantizationConfig(quant_method='mxfp8', weight_block_size=[1, MXFP8_BL...

nemo_automodel.components.models.minimax_m3_vl.state_dict_adapter._NVFP4_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4....