From 547648fa4af12f1d7d5ad0aa3ce12dd747e2c161 Mon Sep 17 00:00:00 2001 From: SakaiXue6666 <2441789115@qq.com> Date: Tue, 11 Aug 2026 22:38:54 +0800 Subject: [PATCH] fix(lora): write exported adapters in PEFT's key layout _save_lora_to_checkpoint documents lora_adapter/ as a portable artifact for external use, loadable with peft.PeftModel.from_pretrained, but write_hf_peft_adapter saved AutoBridge's bare parameter names. PEFT keys carry the base_model.model. prefix of the wrapper module, so from_pretrained matches none of them: it warns about missing adapter keys and leaves every lora_B at zero, loading an adapter that does nothing. Normalize the keys on write. The transform is idempotent, so a state dict that is already in PEFT form passes through unchanged, and the SGLang transports are unaffected because they match on suffixes and a layer-index regex. --- relax/utils/megatron_peft_utils.py | 31 ++++++++++++-- tests/utils/test_megatron_peft_utils.py | 56 +++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 3 deletions(-) diff --git a/relax/utils/megatron_peft_utils.py b/relax/utils/megatron_peft_utils.py index 7fd78a9ab..c09270e6b 100644 --- a/relax/utils/megatron_peft_utils.py +++ b/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", diff --git a/tests/utils/test_megatron_peft_utils.py b/tests/utils/test_megatron_peft_utils.py index f16219e21..efff2e74a 100644 --- a/tests/utils/test_megatron_peft_utils.py +++ b/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):