nemo_automodel.components.models.qwen3_moe.state_dict_adapter
#
Module Contents#
Classes#
Converts between HF Qwen3-MoE checkpoints and our grouped-experts native format. |
Data#
API#
- nemo_automodel.components.models.qwen3_moe.state_dict_adapter.logger#
‘getLogger(…)’
- class nemo_automodel.components.models.qwen3_moe.state_dict_adapter.Qwen3MoeStateDictAdapter(
- config: Any,
- moe_config: nemo_automodel.components.moe.layers.MoEConfig,
- backend: nemo_automodel.components.moe.utils.BackendConfig,
- dtype: torch.dtype = torch.float32,
Bases:
nemo_automodel.components.moe.state_dict_mixin.MoESplitExpertsStateDictMixin
,nemo_automodel.components.checkpoint.state_dict_adapter.StateDictAdapter
Converts between HF Qwen3-MoE checkpoints and our grouped-experts native format.
Qwen3-MoE HF experts use keys: model.layers.{L}.mlp.experts.{E}.gate_proj.weight model.layers.{L}.mlp.experts.{E}.up_proj.weight model.layers.{L}.mlp.experts.{E}.down_proj.weight
Our native format groups them into: model.layers.{L}.mlp.experts.gate_and_up_projs # [n_experts, dim, 2*moe_inter_dim] model.layers.{L}.mlp.experts.down_projs # [n_experts, moe_inter_dim, dim]
Initialization
- to_hf(
- state_dict: dict[str, Any],
- exclude_key_regex: Optional[str] = None,
- quantization: bool = False,
- **kwargs,
- from_hf(
- hf_state_dict: dict[str, Any],
- device_mesh: Optional[torch.distributed.device_mesh.DeviceMesh] = None,
- **kwargs,