> 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.models.gemma4_moe.state_dict_adapter

State-dict adapter for Gemma4 MoE.

HF Gemma4 MoE (eevee-4 26B-A4B) stores expert weights as 3-D tensors:

layers.\{L}.moe.gate\_up\_proj       # \[n\_experts, 2\*expert\_inter\_size, hidden\_size]
layers.\{L}.moe.down\_proj          # \[n\_experts, hidden\_size, expert\_inter\_size]
layers.\{L}.moe.per\_expert\_scale   # \[n\_experts]

NeMo uses transposed layout with concatenated gate+up:

layers.\{L}.moe.experts.gate\_and\_up\_projs  # \[n\_experts, hidden\_size, 2\*expert\_inter\_size]
layers.\{L}.moe.experts.down\_projs         # \[n\_experts, expert\_inter\_size, hidden\_size]

Additionally, the Gemma4 router is mapped to the NeMo Gemma4Gate:

HF:   .router.proj.weight / .router.scale
NeMo: .moe.gate.proj.weight / .moe.gate.scale

The per\_expert\_scale is absorbed into down\_projs during from\_hf.  When
saving back to HF, per\_expert\_scale is emitted as ones (scale already baked
into the weights).

## Module Contents

### Classes

| Name                                                                                                                     | Description                                                     |
| ------------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------- |
| [`Gemma4MoEStateDictAdapter`](#nemo_automodel-components-models-gemma4_moe-state_dict_adapter-Gemma4MoEStateDictAdapter) | Converts between HF Gemma4 MoE checkpoints and the NeMo format. |

### API

```python
class nemo_automodel.components.models.gemma4_moe.state_dict_adapter.Gemma4MoEStateDictAdapter(
    config: typing.Any,
    moe_config: nemo_automodel.components.moe.layers.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype = torch.float32
)
```

**Bases:** [StateDictAdapter](/nemo-automodel/nemo_automodel/components/checkpoint/state_dict_adapter#nemo_automodel-components-checkpoint-state_dict_adapter-StateDictAdapter)

Converts between HF Gemma4 MoE checkpoints and the NeMo format.

Handles:

1. Expert weight concatenation (gate\_proj + up\_proj -> gate\_and\_up\_projs)
2. per\_expert\_scale absorption into down\_projs
3. Router key remapping (router.\* -> moe.gate.\*)
4. Expert-parallel sharding when a device mesh is provided

```python
nemo_automodel.components.models.gemma4_moe.state_dict_adapter.Gemma4MoEStateDictAdapter._gather_expert_tensor(
    tensor: torch.Tensor,
    device_mesh: torch.distributed.device_mesh.DeviceMesh | None,
    n_experts: int
) -> torch.Tensor
```

Gather EP-sharded expert tensor across ranks into a full tensor.

```python
nemo_automodel.components.models.gemma4_moe.state_dict_adapter.Gemma4MoEStateDictAdapter._supports_ep_load_destination(
    tensor: typing.Any,
    n_experts: int
) -> bool
```

staticmethod

Return whether a grouped expert DTensor can receive its HF checkpoint slice in place.

**Parameters:**

**`tensor`** `Any`

Native expert DTensor with global shape `[experts, ...]` and local shape
`[local_experts, ...]`. The `ep` mesh dimension must use `Shard(0)`; every other mesh dimension
must replicate the tensor. Inner-axis expert sharding is intentionally unsupported by this path.

---

**`n_experts`** `int`

Total number of routed experts in the checkpoint.

---

**Returns:** `bool`

`True` when transposing the final local tensor preserves an HF-layout `Shard(0)` destination that

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

Convert a single native tensor back to HF format.

Handles per-tensor conversion for weight streaming (IPC refit) required in RL training:

* Router keys: moe.gate.\{proj.weight,scale} -> router.\{proj.weight,scale}
* Expert gate\_and\_up\_projs: transpose \[E, hidden, 2*inter] -> \[E, 2*inter, hidden]
  and rename to experts.gate\_up\_proj
* Expert down\_projs: transpose \[E, inter, hidden] -> \[E, hidden, inter],
  rename to experts.down\_proj, and emit router.per\_expert\_scale as ones

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

Convert Hugging Face Gemma4 weights into native model layout.

**Parameters:**

**`hf_state_dict`** `dict[str, Any]`

Hugging Face state mapping. Expert gate/up tensors have shape
`[experts, 2 * expert_hidden, hidden]` and down tensors have shape
`[experts, hidden, expert_hidden]`. During a direct checkpoint load, those tensors use the model's
existing weight memory with the last two dimensions transposed.

---

**`device_mesh`** `DeviceMesh | None` — default: None

Optional expert-parallel mesh. Distributed conversion slices the global expert axis and may
shard the native feature axis according to the mesh.

---

**`**kwargs`** — default: \{}

Additional adapter-interface arguments.

---

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

Native state mapping. Expert gate/up tensors have shape `[local_experts, hidden, 2 * expert_hidden]`

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

Convert native Gemma4 weights to Hugging Face keys and layouts.

**Parameters:**

**`state_dict`** `dict[str, Any]`

Native state mapping. Expert gate/up tensors have shape
`[local_experts, hidden, 2 * expert_hidden]` and down tensors have shape
`[local_experts, expert_hidden, hidden]`.

---

**`exclude_key_regex`** `str | None` — default: None

Optional pattern selecting keys to omit.

---

**`quantization`** `bool` — default: False

Whether checkpoint initialization requires a precision conversion. Quantized loads do not
load directly into the model's existing weight memory.

---

**`**kwargs`** — default: \{}

Adapter-interface arguments. `device_mesh` describes expert sharding.
`for_checkpoint_load=True` requests destinations for DCP to load.

---

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

Hugging Face state mapping. Expert gate/up tensors have shape