diff --git a/nemo_automodel/components/models/qwen3_omni_moe/state_dict_adapter.py b/nemo_automodel/components/models/qwen3_omni_moe/state_dict_adapter.py index 20fcf60bc8..7e6e674f61 100644 --- a/nemo_automodel/components/models/qwen3_omni_moe/state_dict_adapter.py +++ b/nemo_automodel/components/models/qwen3_omni_moe/state_dict_adapter.py @@ -51,11 +51,7 @@ def to_hf( hf_state_dict = self._to_hf_w_split_experts(state_dict) if self._uses_thinker_prefix: - hf_state_dict_with_prefix = {} - for key, value in hf_state_dict.items(): - new_key = "thinker." + key - hf_state_dict_with_prefix[new_key] = value - hf_state_dict = hf_state_dict_with_prefix + hf_state_dict = {self._add_thinker_prefix(key): value for key, value in hf_state_dict.items()} if exclude_key_regex: import re @@ -69,25 +65,44 @@ def from_hf( device_mesh: Optional["DeviceMesh"] = None, **kwargs, ) -> dict[str, Any]: - for key in hf_state_dict.keys(): - if ".mlp.experts." in key and key.endswith(".weight"): - self._uses_thinker_prefix = key.startswith("thinker.") - self._uses_model_prefix = "model." in key - break + # Detect the checkpoint's layout from its expert weight keys. PEFT + # saves nest the thinker namespace inside the "base_model.model." + # outer prefix, so check both positions, and consider every matching + # key rather than the first: a full omni dict also carries talker + # expert weights with no thinker namespace, and adapter dicts may mix + # lora keys with modules_to_save-style full weights. + expert_weight_keys = [key for key in hf_state_dict if ".mlp.experts." in key and key.endswith(".weight")] + if expert_weight_keys: + self._uses_thinker_prefix = any( + key.startswith(("thinker.", "base_model.model.thinker.")) for key in expert_weight_keys + ) + self._uses_model_prefix = any("model." in key for key in expert_weight_keys) # Remove thinker prefix if present to match our internal format if self._uses_thinker_prefix: - hf_state_dict_no_prefix = {} - for key, value in hf_state_dict.items(): - if key.startswith("thinker."): - new_key = key[len("thinker.") :] - hf_state_dict_no_prefix[new_key] = value - else: - hf_state_dict_no_prefix[key] = value - hf_state_dict = hf_state_dict_no_prefix + hf_state_dict = {self._strip_thinker_prefix(key): value for key, value in hf_state_dict.items()} return self._from_hf_w_merged_experts(hf_state_dict, device_mesh) + @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("base_model.model."): + return "base_model.model.thinker." + key.removeprefix("base_model.model.") + return "thinker." + key + + @staticmethod + def _strip_thinker_prefix(key: str) -> str: + """Remove the omni checkpoint's ``thinker.`` namespace.""" + if key.startswith("base_model.model.thinker."): + return "base_model.model." + key.removeprefix("base_model.model.thinker.") + return key.removeprefix("thinker.") + 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) @@ -96,7 +111,7 @@ def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[t converted = [(fqn, tensor)] if self._uses_thinker_prefix: - converted = [(f"thinker.{key}", value) for key, value in converted] + converted = [(self._add_thinker_prefix(key), value) for key, value in converted] if exclude_key_regex: import re diff --git a/tests/unit_tests/models/qwen3_omni_moe/test_qwen3_omni_moe_peft_prefix.py b/tests/unit_tests/models/qwen3_omni_moe/test_qwen3_omni_moe_peft_prefix.py new file mode 100644 index 0000000000..3d48981317 --- /dev/null +++ b/tests/unit_tests/models/qwen3_omni_moe/test_qwen3_omni_moe_peft_prefix.py @@ -0,0 +1,178 @@ +# 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 Qwen3 Omni MoE adapter (CPU).""" + +from types import SimpleNamespace + +import torch + +from nemo_automodel.components.models.common import BackendConfig +from nemo_automodel.components.models.qwen3_omni_moe.state_dict_adapter import Qwen3OmniMoeStateDictAdapter +from nemo_automodel.components.moe.config import MoEConfig + + +def _tiny_adapter(): + moe = MoEConfig( + dim=32, + inter_dim=64, + moe_inter_dim=16, + n_routed_experts=2, + n_shared_experts=0, + n_activated_experts=1, + n_expert_groups=1, + n_limited_groups=1, + train_gate=True, + gate_bias_update_factor=0.0, + score_func="softmax", + route_scale=1.0, + aux_loss_coeff=0.0, + norm_topk_prob=False, + expert_bias=False, + router_bias=False, + expert_activation="swiglu", + softmax_before_topk=True, + ) + backend = BackendConfig(linear="torch", rms_norm="torch", attn="sdpa") + return Qwen3OmniMoeStateDictAdapter(SimpleNamespace(), moe, backend) + + +def _peft_lora_state_dict(rank=4, n_experts=2, dim=32, inter=16): + # grouped expert LoRA params exactly as ModelState.state_dict() emits them on + # a PEFT save, shapes matching GroupedExpertsLoRA (lora_experts.py): A is + # [experts, in_features, rank], B is [experts, rank, out_features]. Plus one + # attention LoRA key for contrast. + base = "base_model.model.model.layers.0.mlp.experts" + attn = "base_model.model.model.layers.0.self_attn.q_proj" + return { + f"{base}.lora_gate_and_up_A": torch.randn(n_experts, dim, rank), + f"{base}.lora_gate_and_up_B": torch.randn(n_experts, rank, 2 * inter), + f"{base}.lora_down_A": torch.randn(n_experts, inter, rank), + f"{base}.lora_down_B": torch.randn(n_experts, rank, dim), + f"{attn}.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. On the actual HF omni model PEFT names the text modules + ``base_model.model.thinker.model.layers...``. + """ + adapter = _tiny_adapter() + out = adapter.to_hf(_peft_lora_state_dict()) + + assert out, "no keys came back from to_hf" + for key in out: + assert ".lora_" in key, key + assert key.startswith("base_model.model.thinker.model.layers."), key + # grouped expert tensors still split into per-expert projections + assert any(".mlp.experts.0.gate_proj.lora_A.weight" in k for k in out) + assert any(".mlp.experts.1.down_proj.lora_B.weight" in k for k in out) + + +def test_convert_single_tensor_moves_thinker_inside_the_peft_prefix(): + adapter = _tiny_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 = _tiny_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 grouped keys ModelState expects on resume.""" + adapter = _tiny_adapter() + sd = _peft_lora_state_dict() + out = adapter.to_hf({k: v.clone() for k, v in sd.items()}) + back = adapter.from_hf(dict(out)) + + assert set(back) == set(sd) + for key in sd: + torch.testing.assert_close(back[key], sd[key]) + + +def test_peft_resume_does_not_flip_the_prefix_flags(): + """A LoRA-only adapter dict must not disable the thinker prefix for later + full-checkpoint saves; the flags describe the base checkpoint layout.""" + adapter = _tiny_adapter() + adapter.from_hf(adapter.to_hf(_peft_lora_state_dict())) + + assert adapter._uses_thinker_prefix is True + assert adapter._uses_model_prefix is True + + +def test_full_weight_adapter_keys_keep_the_thinker_flag(): + """modules_to_save-style full weights carry the thinker namespace inside + the peft prefix; detection must recognize it there instead of flipping + the layout flags off.""" + from unittest.mock import patch + + adapter = _tiny_adapter() + hf_state = { + "base_model.model.thinker.model.layers.0.mlp.experts.0.gate_proj.weight": torch.randn(16, 32), + "base_model.model.thinker.model.layers.0.self_attn.q_proj.lora_A.weight": torch.randn(4, 32), + } + + with patch.object(adapter, "_from_hf_w_merged_experts", side_effect=lambda sd, mesh=None: sd): + out = adapter.from_hf(hf_state) + + assert adapter._uses_thinker_prefix is True + assert adapter._uses_model_prefix is True + assert "base_model.model.model.layers.0.mlp.experts.0.gate_proj.weight" in out + assert not any(key.startswith("base_model.model.thinker.") for key in out) + + +def test_talker_expert_keys_do_not_disable_the_thinker_prefix(): + """A full omni dict carries talker experts too; any thinker evidence wins.""" + from unittest.mock import patch + + adapter = _tiny_adapter() + hf_state = { + "talker.model.layers.0.mlp.experts.0.gate_proj.weight": torch.randn(16, 32), + "thinker.model.layers.0.mlp.experts.0.gate_proj.weight": torch.randn(16, 32), + } + + with patch.object(adapter, "_from_hf_w_merged_experts", side_effect=lambda sd, mesh=None: sd): + out = adapter.from_hf(hf_state) + + assert adapter._uses_thinker_prefix is True + assert "model.layers.0.mlp.experts.0.gate_proj.weight" in out + assert "talker.model.layers.0.mlp.experts.0.gate_proj.weight" in out + + +def test_thinker_less_adapter_dict_updates_the_flags(): + """An adapter trained against a thinker-less base sets the flags from its + own keys, so later saves match that base's layout.""" + from unittest.mock import patch + + adapter = _tiny_adapter() + hf_state = { + "base_model.model.model.layers.0.mlp.experts.0.gate_proj.lora_A.weight": torch.randn(4, 32), + } + + with patch.object(adapter, "_from_hf_w_merged_experts", side_effect=lambda sd, mesh=None: sd): + adapter.from_hf(hf_state) + + assert adapter._uses_thinker_prefix is False