nemo_automodel.components.models.qwen3_5_moe.state_dict_adapter

View as Markdown

State-dict adapter for Qwen3.5-MoE.

HF Qwen3.5-MoE stores expert weights as aggregated 3-D tensors:

model.language_model.layers.{L}.mlp.experts.gate_up_proj # [n_experts, 2*moe_inter, hidden] model.language_model.layers.{L}.mlp.experts.down_proj # [n_experts, hidden, moe_inter]

NeMo uses a different naming convention and transposed layout (x @ weight):

model.language_model.layers.{L}.mlp.experts.gate_and_up_projs # [n_experts, hidden, 2*moe_inter] model.language_model.layers.{L}.mlp.experts.down_projs # [n_experts, moe_inter, hidden]

Both expert tensors require .transpose(1, 2) when converting between formats.

Additionally, the shared expert uses singular in HF and plural in NeMo:

HF: .mlp.shared_expert.{gate,up,down}_proj.weight NeMo: .mlp.shared_experts.{gate,up,down}_proj.weight

All other keys (attention, linear_attn/GatedDeltaNet, norms, embeddings, vision encoder) pass through unchanged. The HF VLM checkpoint stores the language model head as model.lm_head while Automodel registers it on the outer model as lm_head.

Module Contents

Classes

NameDescription
Qwen3_5MoeStateDictAdapterConverts between HF Qwen3.5-MoE checkpoints and the NeMo native format.

Functions

NameDescription
_route_fp32_paramsRoute bare GDN fp32 params into the holder used by the native module.
_strip_fp32_paramsStrip the fp32 holder segment from GDN state-dict keys.

API

class nemo_automodel.components.models.qwen3_5_moe.state_dict_adapter.Qwen3_5MoeStateDictAdapter(
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 Qwen3.5-MoE checkpoints and the NeMo native format.

HF Qwen3.5-MoE stores expert weights as aggregated 3-D tensors:

model.language_model.layers.{L}.mlp.experts.gate_up_proj # [n_experts, 2*moe_inter, hidden] model.language_model.layers.{L}.mlp.experts.down_proj # [n_experts, hidden, moe_inter]

NeMo uses a different naming convention and transposed layout (x @ weight):

model.language_model.layers.{L}.mlp.experts.gate_and_up_projs # [n_experts, hidden, 2*moe_inter] model.language_model.layers.{L}.mlp.experts.down_projs # [n_experts, moe_inter, hidden]

Both expert tensors require .transpose(1, 2) when converting between formats.

Additionally, the shared expert uses singular in HF and plural in NeMo:

HF: .mlp.shared_expert.{gate,up,down}_proj.weight NeMo: .mlp.shared_experts.{gate,up,down}_proj.weight

hf_to_internal_map
= {'.mlp.shared_expert.': '.mlp.shared_experts.'}
internal_to_hf_map
nemo_automodel.components.models.qwen3_5_moe.state_dict_adapter.Qwen3_5MoeStateDictAdapter._apply_key_mapping(
state_dict: dict[str, typing.Any],
mapping: dict[str, str]
) -> dict[str, typing.Any]

Apply key substring mappings to state dict keys.

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

Rename a single native key to HF format and transpose expert tensors.

nemo_automodel.components.models.qwen3_5_moe.state_dict_adapter.Qwen3_5MoeStateDictAdapter.from_hf(
hf_state_dict: dict[str, typing.Any],
device_mesh: typing.Optional[torch.distributed.device_mesh.DeviceMesh] = None,
kwargs = {}
) -> dict[str, typing.Any]

Rename HF keys to native keys and transpose expert tensors.

DTensors (DCP path): rename + transpose, no slicing — DCP handles sharding. Plain tensors (init path): slice to local EP shard, transpose, create DTensor.

nemo_automodel.components.models.qwen3_5_moe.state_dict_adapter.Qwen3_5MoeStateDictAdapter.to_hf(
state_dict: dict[str, typing.Any],
exclude_key_regex: typing.Optional[str] = None,
quantization: bool = False,
kwargs = {}
) -> dict[str, typing.Any]

Rename native keys to HF keys and transpose expert tensors. No comms needed.

nemo_automodel.components.models.qwen3_5_moe.state_dict_adapter._route_fp32_params(
key: str
) -> str

Route bare GDN fp32 params into the holder used by the native module.

nemo_automodel.components.models.qwen3_5_moe.state_dict_adapter._strip_fp32_params(
key: str
) -> str

Strip the fp32 holder segment from GDN state-dict keys.