nemo_automodel.components.models.qwen3_next.state_dict_adapter

Module Contents

Classes

Name	Description
`Qwen3NextStateDictAdapter`	Converts between HF Qwen3Next checkpoints and our grouped-experts native format.

Data

logger

API

class nemo_automodel.components.models.qwen3_next.state_dict_adapter.Qwen3NextStateDictAdapter(
    config: typing.Any,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype = torch.float32
)

Bases: MoESplitExpertsStateDictMixin, StateDictAdapter

Converts between HF Qwen3Next checkpoints and our grouped-experts native format.

Qwen3Next 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]

Qwen3Next HF shared experts use keys: model.layers.{L}.mlp.shared_expert.gate_proj.weight model.layers.{L}.mlp.shared_expert.up_proj.weight model.layers.{L}.mlp.shared_expert.down_proj.weight

Our native format uses: model.layers.{L}.mlp.shared_experts.gate_proj.weight # Note: plural “shared_experts” model.layers.{L}.mlp.shared_experts.up_proj.weight model.layers.{L}.mlp.shared_experts.down_proj.weight

hf_to_internal_map

= {'.mlp.shared_expert.': '.mlp.shared_experts.'}

internal_to_hf_map

nemo_automodel.components.models.qwen3_next.state_dict_adapter.Qwen3NextStateDictAdapter._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.

Parameters:

state_dict

dict[str, Any]

State dict to apply mappings to

mapping

dict[str, str]

Dictionary mapping substrings to replace them with

Returns: dict[str, Any]

New state dict with mapped keys

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

Convert a single tensor from native format to HuggingFace format.

Parameters:

fqn

str

Fully qualified name of the tensor in native format

tensor

Any

The tensor to convert

**kwargs

Defaults to {}

Additional arguments for conversion

Returns: list[tuple[str, Any]]

List of (fqn, tensor) tuples in HuggingFace format

nemo_automodel.components.models.qwen3_next.state_dict_adapter.Qwen3NextStateDictAdapter.forced_hf_dtype_mapping(
    state_dict: dict[str, typing.Any]
) -> dict[str, str]

Return HF export dtype overrides for intrinsically-fp32 GDN tensors.

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

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

nemo_automodel.components.models.qwen3_next.state_dict_adapter.logger = logging.getLogger(__name__)