> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.moe.state_dict_mixin

## Module Contents

### Classes

| Name                                                                                                             | Description                                                |
| ---------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------- |
| [`MoESplitExpertsStateDictMixin`](#nemo_automodel-components-moe-state_dict_mixin-MoESplitExpertsStateDictMixin) | Mixin class providing MoE state dict conversion utilities. |

### Data

[`_LORA_EXPERT_SUFFIXES`](#nemo_automodel-components-moe-state_dict_mixin-_LORA_EXPERT_SUFFIXES)

### API

```python
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 all grouped expert destinations alias final model storage.

This only describes the shared expert conversion. Concrete adapters
must separately verify that their non-expert conversions also alias
model storage before opting into a write-through full-checkpoint load.

---

**`_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 non-expert and grouped-expert destinations both alias model storage.

---

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

---

```python
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

```python
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.

```python
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.

```python
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` (r*E, 2*I) -> (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

```python
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 | None` — default: 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`** — default: \{}

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.

```python
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`** `bool` — default: 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).

---

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

```python
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.

```python
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

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

```python
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.

```python
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]` — default: None

Optional device mesh whose `ep` dimension partitions the experts axis.

---

**Raises:**

* `RuntimeError`: If required expert weights are missing from the checkpoint.

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