diff --git a/tests/lora/test_lora_manager.py b/tests/lora/test_lora_manager.py index 85d079673df5..49436d662431 100644 --- a/tests/lora/test_lora_manager.py +++ b/tests/lora/test_lora_manager.py @@ -179,6 +179,100 @@ def test_wrap_gate_linear(default_vllm_config, dist_init, dummy_model): ) +def test_dedup_shared_module_across_paths(default_vllm_config, dist_init, dummy_model): + """A module reachable from two attribute paths (e.g. a MoE gate held + both directly on the block and inside its inner runner) must produce a + single LoRA wrapper. Both paths must end up pointing to that same + wrapper instance, and only the canonical path should live in + `manager.modules` — otherwise activate_adapter would call `reset_lora` + on the alias and clobber weights set under the canonical name. + """ + from vllm.model_executor.layers.linear import ReplicatedLinear + + class AliasContainer(nn.Module): + def __init__(self, gate: nn.Module): + super().__init__() + self.gate = gate # canonical path: "moe.gate" + + # Inner submodule holding the SAME gate instance under another + # path. This mirrors how FusedMoE.runner.gate references the + # block's gate in qwen3_moe. + class _Runner(nn.Module): + def __init__(self, g): + super().__init__() + self.gate = g # alias path: "moe.runner.gate" + + self.runner = _Runner(gate) + + gate = ReplicatedLinear(10, 4, bias=False) + model = dummy_model + model.add_module("moe", AliasContainer(gate)) + + assert model.moe.gate is model.moe.runner.gate + + manager = LoRAModelManager( + model, + 1, + 1, + 1, + LoRAConfig( + max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE + ), + torch.device(DEVICES[0]), + ) + + canonical = manager.model.get_submodule("moe.gate") + alias = manager.model.get_submodule("moe.runner.gate") + + # Same wrapper instance on both paths so forward through either side + # sees the LoRA-augmented module. + assert isinstance(canonical, ReplicatedLinearWithLoRA) + assert alias is canonical + + # Only the canonical path is tracked as a LoRA target. Tracking the + # alias would cause activate_adapter to reset_lora on it after the + # canonical entry already populated the weights. + assert "moe.gate" in manager.modules + assert "moe.runner.gate" not in manager.modules + + +def test_lm_head_exempt_from_dedup(default_vllm_config, dist_init, dummy_model): + """The dedup logic must NOT collapse `lm_head` even when it is reachable + from another attribute path (tied-embedding models do + `self.lm_head = self.model.embed_tokens`, sharing the same nn.Module + instance). The lm_head branch additionally rewires `logits_processor` + into a `LogitsProcessorWithLoRA`, so skipping it would silently break + LoRA on lm_head. + """ + from vllm.lora.layers import LogitsProcessorWithLoRA + + # Add a non-lm_head alias to the same module instance as lm_head. The + # dedup keys on id(module); without the lm_head exemption the alias + # would consume the wrapped_by_id slot first and lm_head would be + # silently skipped, so logits_processor would never be wrapped. + model = dummy_model + model.add_module("embed_tokens", model.lm_head) + assert model.embed_tokens is model.lm_head + + manager = LoRAModelManager( + model, + 1, + 1, + 1, + LoRAConfig( + max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE + ), + torch.device(DEVICES[0]), + ) + + # lm_head's special handling still ran: logits_processor got wrapped + # and the lm_head entry is tracked under self.modules. + assert isinstance( + manager.model.get_submodule("logits_processor"), LogitsProcessorWithLoRA + ) + assert "lm_head" in manager.modules + + def test_skip_unsupported_matched_modules(default_vllm_config, dist_init, dummy_model): class UnsupportedContainer(nn.Module): def __init__(self): diff --git a/vllm/lora/model_manager.py b/vllm/lora/model_manager.py index d68d0c4fbb74..93f57f8fe8ad 100644 --- a/vllm/lora/model_manager.py +++ b/vllm/lora/model_manager.py @@ -361,6 +361,8 @@ def _parent_module(module_name: str) -> str: # - given an input 'x' return '' return module_name.rpartition(".")[0] + wrapped_by_id: dict[int, BaseLayerWithLoRA] = {} + for module_name, module in self.model.named_modules(remove_duplicate=False): if isinstance(module, PPMissingLayer): continue @@ -391,6 +393,23 @@ def _parent_module(module_name: str) -> str: ) continue + existing_wrapper = wrapped_by_id.get(id(module)) + if existing_wrapper is not None and "lm_head" not in module_name: + # Same underlying module was already wrapped under another + # path (e.g. a MoE gate held both directly on the block and + # inside the MoE runner). The adapter targets the canonical + # path (`mlp.gate`); rewire the alias attribute + # (`runner.gate`) to the SAME wrapper so the forward path + # through the alias still applies LoRA, but do NOT add a + # second entry to self.modules — otherwise `activate_adapter` + # would call `reset_lora` on the alias and wipe the weights + # just set under the canonical name, because the alias can't + # load LoRA weights due to name mismatch. + parent = self.model.get_submodule(_parent_module(module_name)) + # reference + setattr(parent, module_name.rpartition(".")[-1], existing_wrapper) + continue + parts = module_name.split(".")[-1] packed_moduled_lst = self.packed_modules_mapping.get(parts, []) if isinstance(module, FusedMoE): @@ -411,6 +430,8 @@ def _parent_module(module_name: str) -> str: self.model.config, ), ) + if isinstance(new_module, BaseLayerWithLoRA): + wrapped_by_id[id(module)] = new_module # (yard1): TODO make this more robust if "lm_head" in module_name: