nemo_automodel.components.moe.state_dict_mixin

Module Contents

Classes

Name	Description
`MoESplitExpertsStateDictMixin`	Mixin 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²).

_supports_write_through_expert_checkpoint_load

bool

Whether grouped expert checkpoint tensors load directly into model weight memory.

This covers only the shared expert conversion. A concrete adapter must also verify that all non-expert checkpoint tensors load directly before it enables the full-checkpoint fast path.

_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.

supports_write_through_checkpoint_load

bool

Whether every checkpoint tensor, including expert tensors, loads directly into model weights.

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
) -> torch.Tensor | None

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: torch.Tensor | None

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

input projection (gate_up_proj or up_proj; inner wrapper, HAS base_layer. prefix):

lora_gate_and_up_B (E, r, U) -> base_layer.lora_A.weight (r*E, U) 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, the input projection 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,
    for_checkpoint_load: bool = False,
    kwargs = {}
) -> list[tuple[str, torch.Tensor]]

Convert one grouped expert tensor to Hugging Face’s per-expert layout.

During checkpoint loading, DCP can write into contiguous or non-contiguous views. A view into model weight memory updates the model directly. A view into temporary TE or MoK grouped storage is rebuilt into the model by from_hf after the read. Save/export conversion still creates contiguous tensors for serialization.

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.

for_checkpoint_load

boolDefaults to False

Return views that DCP will completely overwrite. Save/export callers leave this disabled so converted tensors preserve their current values in contiguous storage.

**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._direct_fill_grouped_expert_tensor(
    expert_parts: list[tuple[torch.Tensor, ...]]
) -> torch.Tensor

Merge experts directly into one final grouped tensor.

Parameters:

expert_parts

list[tuple[torch.Tensor, ...]]

Projection tensors for each local expert. Multiple tensors in a tuple are joined along the last dimension.

Returns: torch.Tensor

One contiguous grouped tensor in self.dtype.

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:

hf_state_dict

dict[str, Any]

State mapping consumed by this method. Per-expert gate and up tensors have shape [expert_hidden, hidden], while down tensors have shape [hidden, expert_hidden]. DTensor values use the same global layouts and are localized before merging.

device_mesh

Optional[DeviceMesh]Defaults to None

Optional device mesh whose expert-parallel dimension selects the local experts. The returned grouped expert tensors use the placements created by create_dtensor_from_local.

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).

Returns: dict[str, Any]

Native state mapping. Gated input projections have shape

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._v5_peft_hf_expert_path_segment() -> str

Return the common HF module path that owns the fused expert parameters.

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')