> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

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

### Functions

| Name                                                                                         | Description                                                                     |
| -------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- |
| [`get_world_size_safe`](#nemo_automodel-components-moe-state_dict_mixin-get_world_size_safe) | Return the distributed world size before or after process-group initialization. |

### Data

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

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

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

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

### 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_checkpoint_tensors_use_model_storage`** `bool`

Whether grouped expert checkpoint tensors use the model's weight memory.

This covers only the shared expert conversion. A concrete adapter must also verify that its non-expert
checkpoint tensors require at most small temporary tensors before enabling low-memory DCP loading.

---

**`_expert_path_segment`** `str`

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

---

**`_grouped_expert_storage_is_model_weight`** `bool`

Whether the runtime grouped expert tensors use the model's parameter storage.

This mirrors the expert implementation selected by `MoE.__init__`. EP dispatchers use ordinary
`GroupedExperts` at world size one, `GroupedExpertsDeepEP` for the grouped backends at larger world sizes,
and `GroupedExpertsTE` otherwise. Non-EP dispatchers construct ordinary grouped experts.

---

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

---

**`supports_low_memory_dcp_load`** `bool`

Whether DCP needs at most small temporary tensors for this MoE checkpoint.

---

**`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,
    legacy_layout: bool = False
) -> 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.
peft flipped the in/out interpretation of these parameters in 0.19.1
(huggingface/peft#3165); the corrected layout is emitted by default,
and *legacy\_layout* selects the pre-flip layout for consumers pinned
to peft \<= 0.19.0.

Shape mapping (automodel native -> ParamWrapper), default (peft >= 0.19.1):

down\_proj (outer wrapper, NO `base_layer` prefix — processed first alphabetically):

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

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

* `lora_gate_and_up_A`  (E, H, r) -> `base_layer.lora_A.weight`  (r\*E, H)  permute+reshape
* `lora_gate_and_up_B`  (E, r, U) -> `base_layer.lora_B.weight`  (U, r\*E)  permute+reshape

and with `legacy_layout=True` (peft \<= 0.19.0, the pre-flip layout):

* `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
* `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. Because peft flipped the
ParamWrapper layout in 0.19.1 (huggingface/peft#3165), the file's
layout is resolved once per call: from the checkpoint's metadata stamp
when the loader provided one (`_paramwrapper_layout_hint`), otherwise
from the tensors whose shapes distinguish the two generations, so
adapters exported under either peft load correctly. When neither
settles it, a clear error is raised instead of guessing. The stamp
must be one of the two known layout ids and must agree with whatever
the shapes say; a stamp that fails either check is rejected, not trusted.

Reverse transforms (down\_proj is outer, the input projection is inner),
peft >= 0.19.1 layout:

* `experts.lora_A.weight`            (r\*E, I)  -> (E, I, r)    = lora\_down\_A
* `experts.lora_B.weight`            (H, r\*E)  -> (E, r, H)    = lora\_down\_B
* `experts.base_layer.lora_A.weight` (r\*E, H)  -> (E, H, r)    = lora\_gate\_and\_up\_A
* `experts.base_layer.lora_B.weight` (U, r\*E)  -> (E, r, U)    = lora\_gate\_and\_up\_B

peft \<= 0.19.0 layout:

* `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, U)  -> (E, r, U)    = 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,
    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`

Native grouped-expert tensor. Gate/up tensors have shape
`[experts, hidden, 2 * expert_hidden]` (or
`[experts, hidden, expert_hidden]` for non-gated MoE), and
down tensors have shape `[experts, expert_hidden, hidden]`.

---

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

---

**`for_checkpoint_load`** `bool` — default: 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`** — 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 Hugging Face format, or `None`

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

```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:**

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

---

**Returns:** `dict[str, Any]`

Native state mapping. Gated input projections have shape

```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.get_world_size_safe() -> int
```

Return the distributed world size before or after process-group initialization.

**Returns:** `int`

The initialized process-group size, or `WORLD_SIZE` before initialization.

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

```python
nemo_automodel.components.moe.state_dict_mixin._PARAMWRAPPER_LAYOUT_LEGACY = 'peft-0.18'
```

```python
nemo_automodel.components.moe.state_dict_mixin._PARAMWRAPPER_LAYOUT_METADATA_KEY = 'paramwrapper_layout'
```

```python
nemo_automodel.components.moe.state_dict_mixin._PARAMWRAPPER_LAYOUT_MODERN = 'peft-0.19.1'
```