fix(qwen2_5_omni): put the thinker prefix inside the peft prefix on lora saves

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,22 +92,55 @@ 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
         return out
 
     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

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"
+    )