nemo_automodel.components.models.qwen3_omni_moe.state_dict_adapter

Module Contents

Classes

Name	Description
`Qwen3OmniMoeStateDictAdapter`	Converts between HF Qwen3OmniMoe checkpoints and grouped-experts native format.

Data

logger

API

class nemo_automodel.components.models.qwen3_omni_moe.state_dict_adapter.Qwen3OmniMoeStateDictAdapter(
    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 Qwen3OmniMoe checkpoints and grouped-experts native format.

nemo_automodel.components.models.qwen3_omni_moe.state_dict_adapter.Qwen3OmniMoeStateDictAdapter._add_thinker_prefix(
    key: str
) -> str

staticmethod

Namespace a native key the way the HF omni checkpoint expects.

PEFT adapter keys keep their base_model.model. outer prefix, so for those the thinker. namespace goes inside it — matching how PEFT names modules on the actual HF omni model.

nemo_automodel.components.models.qwen3_omni_moe.state_dict_adapter.Qwen3OmniMoeStateDictAdapter._strip_thinker_prefix(
    key: str
) -> str

staticmethod

Remove the omni checkpoint’s thinker. namespace.

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

nemo_automodel.components.models.qwen3_omni_moe.state_dict_adapter.Qwen3OmniMoeStateDictAdapter.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_omni_moe.state_dict_adapter.Qwen3OmniMoeStateDictAdapter.map_peft_target_module_to_hf(
    name: str,
    v4_compatible: bool = False
) -> str

Give target_modules the same namespace the exported tensors get.

from_hf records whether the base checkpoint nests the text tower under thinker.; to_hf and convert_single_tensor_to_hf both honor that flag, so target_modules has to honor it too. Otherwise the two halves of one checkpoint name different modules and PEFT cannot match them up.

On the full omni layout the namespace is required: PEFT suffix-matches target_modules against the receiving model, so an entry without it also hits the talker’s identically named modules, which then get adapters the checkpoint has no weights for and stay randomly initialized. On a standalone-thinker base the namespace does not exist on the receiving model at all, and adding it makes PEFT reject the adapter with “Target modules … not found”.

Parameters:

name

str

A target-module name in native layout.

v4_compatible

boolDefaults to False

Legacy export selection; omni uses one module layout for both.

Returns: str

Target-module name in the layout the exported tensors use.

nemo_automodel.components.models.qwen3_omni_moe.state_dict_adapter.Qwen3OmniMoeStateDictAdapter.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_omni_moe.state_dict_adapter.logger = logging.getLogger(__name__)