ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsMoe

nemo_automodel.components.moe.state_dict_mixin

View as Markdown

Module Contents

Classes

NameDescription
MoESplitExpertsStateDictMixinMixin class providing MoE state dict conversion utilities.

Data

_LORA_EXPERT_SUFFIXES

API

class nemo_automodel.components.moe.state_dict_mixin.MoESplitExpertsStateDictMixin()

Mixin class providing MoE state dict conversion utilities.

This mixin provides methods for:

  • Expert parallelism calculations (ranges, assignment)
  • Format conversion between HuggingFace and native formats
  • Both GroupedExperts and DeepEP format support
  • DTensor-aware expert loading and conversion

Can be used by any MoE model that needs expert parallelism and format conversion.

_expert_path_segment
str

Path segment for experts (e.g., ‘mlp.experts’ or ‘mixer.experts’). Override in subclass.

_hf_prefix
str

Prefix for HuggingFace format keys. Override in subclass.

_is_gated_moe
bool

Check if the MoE uses gated activation (e.g., SwiGLU) or non-gated (e.g., ReLU²).

_v5_peft_target_parameters
tuple[str, ...]

Fused expert parameters validated for PEFT v5 ParamWrapper export.

Adapters opt in by overriding this property. Keeping the default empty preserves the legacy per-expert export for model families whose HF naming, activation layout, or checkpoint post-processing has not been validated against ParamWrapper yet.

view_loaded_native_keys
set[str]

Native keys loaded in-place via strided views during the most recent from_hf.

MoE experts with a plain local split are loaded by DCP writing the checkpoint tensors straight through non-contiguous strided views into the model’s grouped expert storage. Such keys are intentionally absent from the dict from_hf returns (the data is already in the model) but are NOT missing. _from_hf_w_merged_experts records them here so the checkpoint loader can exclude them from false “missing” key-diff warnings. The record is reset at the start of each load by _from_hf_w_merged_experts(reset_view_loaded_keys=True).

nemo_automodel.components.moe.state_dict_mixin.MoESplitExpertsStateDictMixin._concatenate_expert_weights(
    expert_weights_by_layer: dict[str, typing.Any],
    n_experts: int
) -> typing.Optional[torch.Tensor]

Concatenate the weights of separate experts into GroupedExpert weights.

Parameters:

expert_weights_by_layer
dict[str, Any]

Nested dict structure containing expert weights

n_experts
int

Total number of experts expected

Returns: Optional[torch.Tensor]

Stacked tensor if all experts are available for a layer, None otherwise

nemo_automodel.components.moe.state_dict_mixin.MoESplitExpertsStateDictMixin._convert_lora_expert_to_hf(
    fqn: str,
    tensor: torch.Tensor,
    n_experts: int,
    inter_dim: int,
    expert_segment: str
) -> list[tuple[str, torch.Tensor]]

Convert a grouped MoE expert LoRA tensor to per-expert HF PEFT format.

Handles the four LoRA parameter types produced by GroupedExpertsLoRA / GroupedExpertsDeepEPLoRA and converts them to per-expert lora_A.weight / lora_B.weight keys that HF PEFT understands.

The prefix (e.g. base_model.model.model.) is preserved from the incoming fqn so that both PEFT and FFT save paths work correctly.

nemo_automodel.components.moe.state_dict_mixin.MoESplitExpertsStateDictMixin._convert_lora_to_paramwrapper(
    fqn: str,
    tensor: torch.Tensor
) -> list[tuple[str, torch.Tensor]]

Convert a single grouped MoE LoRA tensor to PEFT ParamWrapper format.

ParamWrapper format stores fused 3-D expert LoRA parameters as 2-D tensors with the expert dimension folded into the rank dimension.

Shape mapping (automodel native -> ParamWrapper):

down_proj (outer wrapper, NO base_layer prefix — processed first alphabetically):

  • lora_down_B (E, r, H) -> lora_A.weight (r*E, H) reshape
  • lora_down_A (E, I, r) -> lora_B.weight (I, r*E) permute+reshape

gate_up_proj (inner wrapper, HAS base_layer. prefix):

  • lora_gate_and_up_B (E, r, 2I) -> base_layer.lora_A.weight (rE, 2*I) reshape
  • lora_gate_and_up_A (E, H, r) -> base_layer.lora_B.weight (H, r*E) permute+reshape

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

List containing one (fqn, tensor) tuple in ParamWrapper format.

nemo_automodel.components.moe.state_dict_mixin.MoESplitExpertsStateDictMixin._convert_paramwrapper_to_native(
    state_dict: dict[str, typing.Any]
) -> dict[str, typing.Any]

Convert PEFT ParamWrapper LoRA keys to native grouped MoE LoRA format.

This is the reverse of _convert_lora_to_paramwrapper. It detects ParamWrapper-format keys and converts them back to the 3-D grouped tensors expected by GroupedExpertsLoRA.

Reverse transforms (down_proj is outer, gate_up_proj is inner):

  • experts.lora_A.weight (r*E, H) -> (E, r, H) = lora_down_B
  • experts.lora_B.weight (I, r*E) -> (E, I, r) = lora_down_A
  • experts.base_layer.lora_A.weight (rE, 2I) -> (E, r, 2*I) = lora_gate_and_up_B
  • experts.base_layer.lora_B.weight (H, r*E) -> (E, H, r) = lora_gate_and_up_A
nemo_automodel.components.moe.state_dict_mixin.MoESplitExpertsStateDictMixin._convert_single_merged_expert_to_hf_split_experts(
    fqn: str,
    tensor: torch.Tensor,
    prefix_override: str | None = None,
    kwargs = {}
) -> list[tuple[str, torch.Tensor]]

Convert a single merged expert tensor from native format to split HuggingFace format.

When tensor is a model DTensor with a plain (non-DTensor) local split — i.e. ep_shard == 1 — the per-expert outputs are returned as non-contiguous strided views into the local storage of the model’s grouped DTensor instead of newly-allocated contiguous copies. DCP’s target.copy_(source) then writes safetensors data directly through the views into the model’s storage, and _from_hf_w_merged_experts skips the rebuild for the corresponding native key (tracked in _inplace_loaded_native_keys). For loads of large MoE checkpoints this avoids tens of GB of per-expert scratch on top of the already-materialized model.

Save callers must materialize the views before serializing — safetensors.torch.save rejects non-contiguous tensors. See _materialize_to_hf_views_for_save in checkpointing.py.

Parameters:

fqn
str

Fully qualified name of the tensor in native format.

tensor
torch.Tensor

The tensor to convert.

prefix_override
str | NoneDefaults to None

When provided, replaces self._hf_prefix in emitted HF keys. Used to route conversions through namespaces outside the main backbone, e.g. "mtp." for the MTP head.

**kwargs
Defaults to {}

Absorbed for forward-compatibility with base callers that forward arbitrary state-dict kwargs (e.g. exclude_key_regex).

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

List of (fqn, tensor) tuples in HuggingFace format, or None if not an expert tensor.

nemo_automodel.components.moe.state_dict_mixin.MoESplitExpertsStateDictMixin._from_hf_w_merged_experts(
    hf_state_dict: dict[str, typing.Any],
    device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None,
    reset_view_loaded_keys: bool = True
) -> dict[str, typing.Any]

Convert HF checkpoint to native format.

For gated activations (SwiGLU, Quick-GEGLU): Creates combined gate_and_up_projs [n_experts, dim, 2*inter_dim] and transposed down_projs tensors.

For non-gated activations (ReLU²): Creates gate_and_up_projs [n_experts, dim, inter_dim] and transposed down_projs tensors.

Parameters:

reset_view_loaded_keys
boolDefaults to True

Clear the in-place (strided-view) loaded-key record at the start of this call. A single from_hf may invoke this method more than once (e.g. backbone then MTP merge); the later call(s) pass False so the view-loaded keys accumulate across one logical load. Resetting here (rather than in the loader) keeps the whole view-key lifecycle inside the adapter and ensures each load starts clean (no leak from a prior load such as an init-time partial load).

nemo_automodel.components.moe.state_dict_mixin.MoESplitExpertsStateDictMixin._recombine_lora_expert_keys(
    state_dict: dict[str, typing.Any]
) -> dict[str, typing.Any]

Recombine per-expert HF LoRA keys back to grouped MoE LoRA format.

This is the reverse of _convert_lora_expert_to_hf. It detects per-expert LoRA keys (e.g. layers.0.mlp.experts.0.gate_proj.lora_A.weight) and recombines them into the grouped tensors expected by GroupedExpertsLoRA / GroupedExpertsDeepEPLoRA (e.g. layers.0.mlp.experts.lora_gate_and_up_A).

nemo_automodel.components.moe.state_dict_mixin.MoESplitExpertsStateDictMixin._register_inplace_loaded_key(
    fqn: str,
    prefix_override: str | None
) -> None

Mark fqn as loaded via in-place views so _from_hf_w_merged_experts skips its rebuild.

The tracked key must match the native_key that the from_hf merge loop reconstructs from the HF per-expert keys. For backbone tensors the native_key equals fqn; for MTP tensors (prefix_override="mtp.") the HF keys live under the mtp. namespace and from_hf processes them with that prefix stripped, so the tracked key is also the mtp.-less form. The user of this set (_from_hf_w_merged_experts) receives the matching stripped key when called via the adapter’s per-namespace dispatch.

nemo_automodel.components.moe.state_dict_mixin.MoESplitExpertsStateDictMixin._split_experts_weights(
    weight: torch.Tensor,
    n_experts: int
) -> list[torch.Tensor]

Split grouped expert weights into per-expert tensors.

Parameters:

weight
torch.Tensor

Tensor of shape [experts, …], with arbitrary trailing dimensions. An EP DTensor uses an ep mesh dimension. A non-EP DTensor may use any FSDP placement on a mesh without ep.

n_experts
int

Global number of experts in weight.

Returns: list[torch.Tensor]

Per-expert tensors of shape […]. A DTensor sharded on the expert axis returns only the experts local

nemo_automodel.components.moe.state_dict_mixin.MoESplitExpertsStateDictMixin._to_hf_w_split_experts(
    state_dict: dict[str, typing.Any],
    kwargs: typing.Any = {}
) -> dict[str, typing.Any]

Convert DeepEP format to HuggingFace format.

Handles gate_and_up_projs / down_projs -> individual expert weights. Forwards **kwargs to _convert_single_merged_expert_to_hf_split_experts for adapter compatibility (e.g. exclude_key_regex).

nemo_automodel.components.moe.state_dict_mixin.MoESplitExpertsStateDictMixin._validate_expert_availability(
    hf_state_dict: dict[str, typing.Any],
    n_experts: int,
    device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None
) -> None

Validate that all required experts are available in the HF state dict before loading. Only validates experts needed for the current rank and layers present in the state dict. Expert groups already loaded through registered in-place views validate the rank-local IDs recorded by _split_experts_weights, including when EP is disabled and no MoE mesh is passed to this method.

Parameters:

hf_state_dict
dict[str, Any]

HuggingFace state mapping. Expert gate/up values are tensors of shape [expert_hidden, hidden], and down values have shape [hidden, expert_hidden]. This method validates their keys only.

n_experts
int

Total number of experts.

device_mesh
Optional[DeviceMesh]Defaults to None

Optional device mesh whose ep dimension partitions the experts axis.

Raises:

  • RuntimeError: If required expert weights are missing from the checkpoint.
nemo_automodel.components.moe.state_dict_mixin._LORA_EXPERT_SUFFIXES = ('lora_gate_and_up_A', 'lora_gate_and_up_B', 'lora_down_A', 'lora_down_B')