diff --git a/nemo_automodel/components/models/qwen2_5_omni/state_dict_adapter.py b/nemo_automodel/components/models/qwen2_5_omni/state_dict_adapter.py index c26626ef0f..b428df5644 100644 --- a/nemo_automodel/components/models/qwen2_5_omni/state_dict_adapter.py +++ b/nemo_automodel/components/models/qwen2_5_omni/state_dict_adapter.py @@ -26,6 +26,9 @@ _THINKER_PREFIX = "thinker." _DROP_PREFIXES = ("talker.", "token2wav.") +# The outer prefix PEFT saves carry (added in ModelState.state_dict); the +# thinker namespace nests inside it on those keys. +_PEFT_PREFIX = "base_model.model." # Keys the NeMo Thinker class deletes at __init__ time (so the HF base # checkpoint must not be allowed to repopulate them, otherwise load_state_dict @@ -71,7 +74,7 @@ def to_hf( **kwargs, ) -> dict[str, Any]: if self._uses_thinker_prefix: - hf_state_dict = {_THINKER_PREFIX + k: v for k, v in state_dict.items()} + hf_state_dict = {self._add_thinker_prefix(k): v for k, v in state_dict.items()} else: hf_state_dict = dict(state_dict) @@ -89,13 +92,18 @@ def from_hf( out: dict[str, Any] = {} saw_thinker = False for key, value in hf_state_dict.items(): - if key.startswith(_DROP_PREFIXES): + # Drop talker/token2wav keys in either position: bare (full omni + # checkpoints) or nested inside the PEFT prefix (external + # full-omni adapters). + if key.startswith(_DROP_PREFIXES) or key.removeprefix(_PEFT_PREFIX).startswith(_DROP_PREFIXES): continue - stripped = key[len(_THINKER_PREFIX) :] if key.startswith(_THINKER_PREFIX) else key + stripped = self._strip_thinker_prefix(key) # Drop keys for parameters the NeMo Thinker deletes at __init__. if any(sub in stripped for sub in _DROP_THINKER_KEY_SUBSTRINGS): continue - if key.startswith(_THINKER_PREFIX): + # _strip_thinker_prefix changes the key iff it carried the + # thinker namespace in either position. + if stripped != key: saw_thinker = True out[stripped] = value self._uses_thinker_prefix = saw_thinker or self._uses_thinker_prefix @@ -103,8 +111,36 @@ def from_hf( def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[tuple[str, Any]]: exclude_key_regex = kwargs.get("exclude_key_regex", None) - key = _THINKER_PREFIX + fqn if self._uses_thinker_prefix else fqn + key = self._add_thinker_prefix(fqn) if self._uses_thinker_prefix else fqn if exclude_key_regex: if re.match(exclude_key_regex, key): return [] return [(key, tensor)] + + @staticmethod + def _add_thinker_prefix(key: str) -> str: + """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. + """ + if key.startswith(_PEFT_PREFIX): + return _PEFT_PREFIX + _THINKER_PREFIX + key.removeprefix(_PEFT_PREFIX) + return _THINKER_PREFIX + key + + @staticmethod + def _strip_thinker_prefix(key: str) -> str: + """Remove the omni checkpoint's ``thinker.`` namespace.""" + if key.startswith(_PEFT_PREFIX + _THINKER_PREFIX): + return _PEFT_PREFIX + key.removeprefix(_PEFT_PREFIX + _THINKER_PREFIX) + return key.removeprefix(_THINKER_PREFIX) + + def map_peft_target_module_to_hf(self, 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. + """ + return _THINKER_PREFIX + name diff --git a/tests/unit_tests/models/qwen2_5_omni/test_qwen2_5_omni_peft_prefix.py b/tests/unit_tests/models/qwen2_5_omni/test_qwen2_5_omni_peft_prefix.py new file mode 100644 index 0000000000..25f19b9869 --- /dev/null +++ b/tests/unit_tests/models/qwen2_5_omni/test_qwen2_5_omni_peft_prefix.py @@ -0,0 +1,126 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""PEFT save/resume prefix handling for the Qwen2.5 Omni adapter (CPU).""" + +from types import SimpleNamespace + +import torch + +from nemo_automodel.components.models.qwen2_5_omni.state_dict_adapter import Qwen2_5OmniStateDictAdapter + + +def _adapter(): + return Qwen2_5OmniStateDictAdapter(config=SimpleNamespace()) + + +def _peft_lora_state_dict(rank=4, dim=32): + base = "base_model.model.model.layers.0" + return { + f"{base}.self_attn.q_proj.lora_A.weight": torch.randn(rank, dim), + f"{base}.self_attn.q_proj.lora_B.weight": torch.randn(dim, rank), + f"{base}.mlp.gate_proj.lora_A.weight": torch.randn(rank, dim), + } + + +def test_peft_lora_keys_get_thinker_inside_the_peft_prefix(): + """The thinker namespace must land inside the PEFT prefix. + + Before the fix every key got "thinker." prepended on the outside, + producing ``thinker.base_model.model...`` names that HF PEFT cannot + attach. + """ + adapter = _adapter() + out = adapter.to_hf(_peft_lora_state_dict()) + + assert out, "no keys came back from to_hf" + for key in out: + assert key.startswith("base_model.model.thinker.model.layers."), key + + +def test_convert_single_tensor_moves_thinker_inside_the_peft_prefix(): + adapter = _adapter() + tensor = torch.randn(4, 32) + result = adapter.convert_single_tensor_to_hf( + "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight", tensor + ) + + assert result == [("base_model.model.thinker.model.layers.0.self_attn.q_proj.lora_A.weight", tensor)] + + +def test_full_weights_still_get_the_thinker_prefix(): + adapter = _adapter() + out = adapter.to_hf({"model.layers.0.self_attn.q_proj.weight": torch.randn(32, 32)}) + + assert list(out) == ["thinker.model.layers.0.self_attn.q_proj.weight"] + + +def test_peft_lora_save_round_trips_for_resume(): + """from_hf must rebuild the exact keys ModelState expects on resume.""" + adapter = _adapter() + sd = _peft_lora_state_dict() + back = adapter.from_hf(adapter.to_hf(sd)) + + assert set(back) == set(sd) + for key in sd: + torch.testing.assert_close(back[key], sd[key]) + + +def test_correctly_named_external_adapter_imports(): + """An adapter written in the proper HF PEFT layout must load. + + Before the fix from_hf only stripped a leading "thinker.", so these + keys passed through unchanged and never matched a model parameter. + """ + adapter = _adapter() + tensor = torch.randn(4, 32) + back = adapter.from_hf({"base_model.model.thinker.model.layers.0.self_attn.q_proj.lora_A.weight": tensor}) + + assert list(back) == ["base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight"] + + +def test_legacy_malformed_saves_still_resume(): + """Adapters saved before the fix carry the outer thinker prefix.""" + adapter = _adapter() + tensor = torch.randn(4, 32) + back = adapter.from_hf({"thinker.base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": tensor}) + + assert list(back) == ["base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight"] + + +def test_talker_keys_dropped_in_both_positions(): + """talker/token2wav weights are dropped whether bare (full checkpoints) + or nested inside the peft prefix (external full-omni adapters).""" + adapter = _adapter() + back = adapter.from_hf( + { + "thinker.model.layers.0.self_attn.q_proj.weight": torch.randn(32, 32), + "talker.model.layers.0.self_attn.q_proj.weight": torch.randn(32, 32), + "base_model.model.talker.model.layers.0.self_attn.q_proj.lora_A.weight": torch.randn(4, 32), + } + ) + + assert list(back) == ["model.layers.0.self_attn.q_proj.weight"] + + +def test_target_modules_get_the_thinker_namespace(): + """adapter_config.json target_modules must carry thinker. so PEFT's + suffix matching on the full omni model doesn't also hit the talker's + structurally identical submodules.""" + adapter = _adapter() + + assert ( + adapter.map_peft_target_module_to_hf("model.layers.0.self_attn.q_proj") + == "thinker.model.layers.0.self_attn.q_proj" + )