nemo_automodel.components.models.minimax_m2.state_dict_adapter

Module Contents

Classes

Name	Description
`MiniMaxM2StateDictAdapter`	Convert between MiniMax-M2.1 HF checkpoints and native grouped-expert format.

Functions

Name	Description
`should_quantize_key`	-

Data

NON_QUANTIZED_KEY_PATTERNS

API

class nemo_automodel.components.models.minimax_m2.state_dict_adapter.MiniMaxM2StateDictAdapter(
    config: typing.Any,
    moe_config: nemo_automodel.components.moe.layers.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype = torch.float32
)

Bases: MoESplitExpertsStateDictMixin, StateDictAdapter

Convert between MiniMax-M2.1 HF checkpoints and native grouped-expert format.

_expert_path_segment

str

nemo_automodel.components.models.minimax_m2.state_dict_adapter.MiniMaxM2StateDictAdapter._dequantize(
    state_dict: dict[str, typing.Any]
) -> dict[str, typing.Any]

nemo_automodel.components.models.minimax_m2.state_dict_adapter.MiniMaxM2StateDictAdapter._hf_key_to_native(
    key: str
) -> str

nemo_automodel.components.models.minimax_m2.state_dict_adapter.MiniMaxM2StateDictAdapter._native_key_to_hf(
    key: str
) -> str

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

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

Convert HF checkpoint to native format in-place.

Operates in-place on the input dict to avoid allocating a full copy, reducing peak memory from 2x to ~1x model size.

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

nemo_automodel.components.models.minimax_m2.state_dict_adapter.should_quantize_key(
    key: str
) -> bool

nemo_automodel.components.models.minimax_m2.state_dict_adapter.NON_QUANTIZED_KEY_PATTERNS = ['input_layernorm.weight', 'post_attention_layernorm.weight', 'norm.weight', 'lm...