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`	Converts between HF Gemma4 MoE checkpoints and the NeMo format.

API

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

Converts between HF Gemma4 MoE checkpoints and the NeMo format.

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.

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

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, 2inter] -> [E, 2inter, 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

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 | NoneDefaults to None

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

**kwargs

Defaults to {}

Additional adapter-interface arguments.

Returns: dict[str, Any]

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

nemo_automodel.components.models.gemma4_moe.state_dict_adapter.Gemma4MoEStateDictAdapter.get_hf_state_dict_keys(
    state_dict: dict[str, typing.Any]
) -> list[str]

Return converted keys without gathering real expert weights.

Parameters:

state_dict

dict[str, Any]

Native Gemma4 state mapping. Expert tensors have shape [local_experts, hidden, 2 * expert_hidden] for fused gate-up weights or [local_experts, expert_hidden, hidden] for down weights. Other tensor values retain their model-owned layouts.

Returns: list[str]

Hugging Face state-dict keys in adapter iteration order.

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 | NoneDefaults to None

Optional pattern selecting keys to omit.

quantization

boolDefaults to False

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

**kwargs

Defaults to {}

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