nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter

Module Contents

Classes

Name	Description
`MiMoV2FlashStateDictAdapter`	Convert MiMo-V2-Flash HF checkpoints to Automodel’s grouped MoE layout.

Functions

Name	Description
`_should_quantize_key`	-

Data

NON_QUANTIZED_KEY_PATTERNS

logger

API

class nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter(
    config: typing.Any,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype = torch.bfloat16
)

Bases: MoESplitExpertsStateDictMixin, StateDictAdapter

Convert MiMo-V2-Flash HF checkpoints to Automodel’s grouped MoE layout.

HF stores routed experts as split per-expert projections: mlp.experts.{E}.{gate,up,down}_proj.weight. Automodel groups those into gate_and_up_projs and down_projs so EP can shard experts without materializing every expert on every rank.

nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.MiMoV2FlashStateDictAdapter._create_scale_inv_for_hf_key(
    key: str,
    weight: torch.Tensor
) -> torch.Tensor

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

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

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

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

Convert Automodel state_dict to the HF MiMo-V2-Flash layout.

Note: The quantization parameter is accepted for interface compatibility but is ignored. MiMo-V2-Flash is distributed as an FP8 HF checkpoint, so this adapter always emits FP8 weights plus _scale_inv companions for keys that match _should_quantize_key, regardless of the caller’s preference.

nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter._should_quantize_key(
    key: str
) -> bool

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

nemo_automodel.components.models.mimo_v2_flash.state_dict_adapter.logger = logging.getLogger(__name__)