Skip to navigation

nemo_automodel.components.models.beam.state_dict_adapter

View as Markdown

State-dict adapter between Hugging Face Beam checkpoints and the native grouped-experts layout.

Module Contents

Classes

NameDescription
BeamStateDictAdapterConverts between Hugging Face Beam checkpoints and the grouped-experts native format.

API

class nemo_automodel.components.models.beam.state_dict_adapter.BeamStateDictAdapter()

Bases: Qwen3MoeStateDictAdapter

Converts between Hugging Face Beam checkpoints and the grouped-experts native format.

Beam stores its experts exactly like Qwen3-MoE, so the Qwen3-MoE conversion (including fused PEFT v5 ParamWrapper LoRA export/import and low-memory DCP loading) applies unchanged. The native module tree mirrors the HF one, so every tensor outside the experts keeps its HF name and passes through unchanged (attention projections and gate_bias, q/k norms, the four per-layer norms, the dense MLP, mlp.gate.weight and the mlp.gate.e_score_correction_bias buffer, embeddings, final norm, lm_head). Experts differ:

HF checkpoint (per expert, the qwen2_moe conversion mapping)::

model.layers.{L}.mlp.experts.{E}.gate_proj.weight [moe_inter_dim, hidden] model.layers.{L}.mlp.experts.{E}.up_proj.weight [moe_inter_dim, hidden] model.layers.{L}.mlp.experts.{E}.down_proj.weight [hidden, moe_inter_dim]

Native (grouped)::

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

to_hf therefore writes the per-expert layout that transformers itself produces with save_pretrained, not the stacked in-memory layout.