nemo_automodel.components.models.qwen2_5_omni.state_dict_adapter

Module Contents

Classes

Name	Description
`Qwen2_5OmniStateDictAdapter`	HF Qwen2.5-Omni checkpoint adapter (thinker-only path).

Data

_DROP_PREFIXES

_DROP_THINKER_KEY_SUBSTRINGS

_PEFT_PREFIX

_THINKER_PREFIX

logger

API

class nemo_automodel.components.models.qwen2_5_omni.state_dict_adapter.Qwen2_5OmniStateDictAdapter(
    config: typing.Any,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    dtype: torch.dtype = torch.float32
)

Bases: StateDictAdapter

HF Qwen2.5-Omni checkpoint adapter (thinker-only path).

HF Qwen/Qwen2.5-Omni-* checkpoints store keys under three top-level prefixes: thinker.*, talker.*, token2wav.*. For ASR/text fine-tuning we only train the Thinker, so this adapter:

on from_hf: drops talker.* and token2wav.* keys and strips the thinker. prefix so keys align with our NeMo Thinker class.
on to_hf: re-adds the thinker. prefix so the saved checkpoint can be merged back with the original talker/token2wav shards.

Qwen2.5-Omni-3B is dense (no MoE), so no expert grouping logic is needed — this is a thin key-renaming adapter.

backend

= backend or BackendConfig()

nemo_automodel.components.models.qwen2_5_omni.state_dict_adapter.Qwen2_5OmniStateDictAdapter._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.qwen2_5_omni.state_dict_adapter.Qwen2_5OmniStateDictAdapter._strip_thinker_prefix(
    key: str
) -> str

staticmethod

Remove the omni checkpoint’s thinker. namespace.

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

nemo_automodel.components.models.qwen2_5_omni.state_dict_adapter.Qwen2_5OmniStateDictAdapter.from_hf(
    hf_state_dict: dict[str, typing.Any],
    device_mesh: typing.Any | None = None,
    kwargs = {}
) -> dict[str, typing.Any]

nemo_automodel.components.models.qwen2_5_omni.state_dict_adapter.Qwen2_5OmniStateDictAdapter.map_peft_target_module_to_hf(
    name: str
) -> str

Namespace adapter_config.json target_modules under thinker..

Without it, PEFT’s suffix matching on the full omni model also hits the talker’s structurally identical submodules and injects adapter modules the saved file has no weights for.

nemo_automodel.components.models.qwen2_5_omni.state_dict_adapter.Qwen2_5OmniStateDictAdapter.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.qwen2_5_omni.state_dict_adapter._DROP_PREFIXES = ('talker.', 'token2wav.')

nemo_automodel.components.models.qwen2_5_omni.state_dict_adapter._DROP_THINKER_KEY_SUBSTRINGS = ('audio_tower.audio_bos_eos_token',)

nemo_automodel.components.models.qwen2_5_omni.state_dict_adapter._PEFT_PREFIX = 'base_model.model.'

nemo_automodel.components.models.qwen2_5_omni.state_dict_adapter._THINKER_PREFIX = 'thinker.'

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