fix(lora): write exported adapters in PEFT's key layout

relax/utils/megatron_peft_utils.py

@@ -134,6 +134,27 @@ def build_hf_peft_config_dict(
     }
 
 
+# PEFT wraps the base model as ``base_model.model``, so every key in a standard
+# ``adapter_model.safetensors`` carries that prefix. Bridge's adapter export yields bare
+# HF parameter names, which PEFT then cannot match: ``PeftModel.from_pretrained`` reports
+# the keys as missing (a warning, not an error) and leaves every ``lora_B`` at zero, i.e.
+# it silently loads an adapter that does nothing.
+PEFT_STATE_DICT_PREFIX = "base_model.model."
+
+
+def to_peft_state_dict(adapter_weights: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Return ``adapter_weights`` keyed the way standard PEFT writes them.
+
+    Idempotent: keys that already carry the prefix are left alone, so callers that
+    hand over an exporter's output and callers that hand over an existing PEFT state
+    dict both get the same result.
+    """
+    return {
+        name if name.startswith(PEFT_STATE_DICT_PREFIX) else PEFT_STATE_DICT_PREFIX + name: tensor
+        for name, tensor in adapter_weights.items()
+    }
+
+
 def write_hf_peft_adapter(
     merged: dict[str, torch.Tensor],
     adapter_dir,
@@ -147,10 +168,12 @@ def write_hf_peft_adapter(
 
     Produces ``adapter_config.json`` + ``adapter_model.safetensors`` — the on-disk
     format used by the checkpoint save, which Megatron-Bridge's ``load_peft_adapter``
-    can read back.
+    can read back. Keys are normalized to PEFT's ``base_model.model.`` layout so the
+    directory also loads with ``peft.PeftModel.from_pretrained``.
 
     Args:
-        merged: Full (TP-gathered, PP-merged) adapter tensors keyed by HF name.
+        merged: Full (TP-gathered, PP-merged) adapter tensors keyed by HF name, with
+            or without the PEFT prefix (see ``to_peft_state_dict``).
         adapter_dir: Target directory (created if missing).
         lora_rank/lora_alpha/target_modules/lora_dropout: PEFT config written to
             ``adapter_config.json`` (HF-style target module names).
@@ -177,7 +200,7 @@ def write_hf_peft_adapter(
         json.dump(config_dict, f)
 
     # safetensors requires contiguous CPU tensors.
-    state = {name: t.contiguous() for name, t in merged.items()}
+    state = {name: t.contiguous() for name, t in to_peft_state_dict(merged).items()}
     save_file(state, str(adapter_dir / "adapter_model.safetensors"))
     return str(adapter_dir)
 
@@ -302,6 +325,8 @@ def build_lora_peft(args):
     "count_adapter_parameters",
     "convert_megatron_to_hf_target_modules",
     "MEGATRON_TO_HF_MODULES",
+    "PEFT_STATE_DICT_PREFIX",
+    "to_peft_state_dict",
     "write_hf_peft_adapter",
     "extract_lora_delta",
     "is_lora_enabled",

tests/utils/test_megatron_peft_utils.py

@@ -29,6 +29,7 @@
     is_lora_adapter_param,
     is_lora_enabled,
     is_lora_merge_mode,
+    to_peft_state_dict,
     write_hf_peft_adapter,
 )
 
@@ -145,6 +146,61 @@ def test_write_hf_peft_adapter_makes_missing_dirs(self, tmp_path):
         )
         assert (target / "adapter_config.json").is_file()
 
+    def test_write_hf_peft_adapter_adds_peft_prefix(self, tmp_path):
+        """Bare exporter names are written in PEFT's layout.
+
+        ``AutoBridge.export_adapter_weights`` yields bare HF parameter names.
+        Written as-is, ``peft.PeftModel.from_pretrained`` matches none of them:
+        it warns about missing adapter keys and leaves every ``lora_B`` at
+        zero, so the adapter loads as a no-op instead of failing.
+        """
+        from safetensors.torch import load_file
+
+        merged = {
+            "model.layers.0.self_attn.q_proj.lora_A.weight": torch.randn(8, 16),
+            "model.layers.0.self_attn.q_proj.lora_B.weight": torch.randn(16, 8),
+        }
+        adapter_dir = tmp_path / "lora_adapter"
+        write_hf_peft_adapter(
+            merged,
+            adapter_dir,
+            lora_rank=8,
+            lora_alpha=16,
+            target_modules=["q_proj"],
+            lora_dropout=0.0,
+        )
+
+        loaded = load_file(str(adapter_dir / "adapter_model.safetensors"))
+        assert set(loaded) == {
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight",
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight",
+        }
+        for name, tensor in merged.items():
+            assert torch.allclose(loaded[f"base_model.model.{name}"], tensor)
+
+
+class TestToPeftStateDict:
+    def test_bare_names_get_the_prefix(self):
+        state = to_peft_state_dict({"model.layers.0.self_attn.q_proj.lora_A.weight": torch.zeros(2, 2)})
+        assert list(state) == ["base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight"]
+
+    def test_multimodal_names_keep_their_own_prefix(self):
+        """A submodel prefix (e.g. Qwen3-Omni's ``thinker.``) is part of the
+        parameter path and stays inside the PEFT prefix."""
+        state = to_peft_state_dict({"thinker.model.layers.0.self_attn.o_proj.lora_B.weight": torch.zeros(2, 2)})
+        assert list(state) == ["base_model.model.thinker.model.layers.0.self_attn.o_proj.lora_B.weight"]
+
+    def test_already_prefixed_names_are_left_alone(self):
+        """Idempotent, so a state dict that is already in PEFT form survives a
+        second pass unchanged."""
+        name = "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight"
+        assert list(to_peft_state_dict({name: torch.zeros(2, 2)})) == [name]
+
+    def test_tensors_are_passed_through_unchanged(self):
+        tensor = torch.randn(4, 4)
+        state = to_peft_state_dict({"model.layers.0.self_attn.q_proj.lora_A.weight": tensor})
+        assert next(iter(state.values())) is tensor
+
 
 class TestLoraPredicates:
     def test_is_lora_enabled_true_when_rank_positive(self):