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,
    v4_compatible: bool = False
) -> str

Namespace adapter_config.json target_modules the way the tensors are namespaced.

PEFT suffix-matches target_modules against the receiving model, so on the full omni layout an entry without thinker. also hits the talker’s structurally identical submodules and injects adapters the saved file has no weights for. This follows _uses_thinker_prefix the same way to_hf and convert_single_tensor_to_hf do, so the hook and the tensor exporters stay aligned.

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.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__)