From dbd6a35dd2e3d406bf0452d42d15f328b965b1e1 Mon Sep 17 00:00:00 2001 From: Haozhe Jiang <162801044+provoke210@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:28:00 +0800 Subject: [PATCH] [Bugfix] Reject partially bound diffusion LoRA adapters (#8008) Signed-off-by: Haozhe Jiang <162801044+provoke210@users.noreply.github.com> --- docs/user_guide/diffusion/lora.md | 15 +++ tests/diffusion/lora/test_lora_manager.py | 121 ++++++++++++++++-- .../minimax_h3/test_minimax_h3_native_lora.py | 2 + vllm_omni/diffusion/lora/manager.py | 34 +++-- 4 files changed, 146 insertions(+), 26 deletions(-) diff --git a/docs/user_guide/diffusion/lora.md b/docs/user_guide/diffusion/lora.md index c774cb288..1d9201796 100644 --- a/docs/user_guide/diffusion/lora.md +++ b/docs/user_guide/diffusion/lora.md @@ -27,6 +27,21 @@ The `adapter_config.json` file contains metadata about the LoRA adapter, includi - `lora_alpha`: LoRA alpha scaling factor - `target_modules`: List of module names to apply LoRA to +Every module supplied by a PEFT adapter must bind successfully when the adapter +is activated. If any module cannot bind, activation raises an error naming the +unbound modules and resets the LoRA layers to avoid leaving the adapter partly +active. An adapter may still target only some projections of a fused layer, such +as Q and V without K, as long as all supplied modules bind. + +Previously, some adapters activated even when unmatched modules were silently +dropped. Those adapters now fail at activation. For example, in-memory tensor +adapters for Qwen-Image, Boogu-Image, and Wan can contain checkpoint keys such +as `to_out.0` where the engine expects `to_out`. PEFT checkpoints with extra +unbound keys also fail at activation if those keys pass the loader's earlier +validation; unsupported module suffixes may already fail during loading. +Checkpoint-to-engine name mapping is tracked in [#8001](https://github.com/vllm-project/vllm-omni/issues/8001); +until it is available, adapter keys must resolve to supported engine modules. + ## Quick Start ### Offline Inference diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index f5d069b6a..7b68b2558 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -7,6 +7,7 @@ from types import SimpleNamespace import pytest import torch +from vllm.lora.lora_model import LoRAModel from vllm.lora.lora_weights import LoRALayerWeights from vllm.lora.utils import get_supported_lora_modules @@ -294,7 +295,9 @@ def test_lora_manager_activates_fused_lora_on_packed_layer(): assert torch.allclose(torch.cat([b0, b1, b2], dim=0), B * 0.5) -def test_lora_manager_activates_packed_lora_from_sublayers(): +@pytest.mark.parametrize("name_prefix", ["transformer.blocks.0.attn.", "blocks.0.attn.", ""]) +@pytest.mark.parametrize("target_modules", [("to_q",), ("to_q", "to_v"), ("to_q", "to_k", "to_v")]) +def test_lora_manager_activates_packed_lora_from_sublayers(name_prefix, target_modules): pipeline = torch.nn.Module() bound_names = None @@ -322,8 +325,10 @@ def test_lora_manager_activates_packed_lora_from_sublayers(): rank = 2 loras: dict[str, LoRALayerWeights] = {} for name, out_dim in zip(["to_q", "to_k", "to_v"], [2, 1, 1]): - loras[f"transformer.blocks.0.attn.{name}"] = LoRALayerWeights( - module_name=f"transformer.blocks.0.attn.{name}", + if name not in target_modules: + continue + loras[f"{name_prefix}{name}"] = LoRALayerWeights( + module_name=f"{name_prefix}{name}", rank=rank, lora_alpha=rank, lora_a=torch.ones((rank, 4)) * (1 if name == "to_q" else 2), @@ -344,9 +349,15 @@ def test_lora_manager_activates_packed_lora_from_sublayers(): assert len(lora_a_list) == 3 assert len(lora_b_list) == 3 # Scale should apply to B only. - assert torch.allclose(lora_b_list[0], torch.ones((2, rank)) * 3 * 2.0) - assert torch.allclose(lora_b_list[1], torch.ones((1, rank)) * 4 * 2.0) - assert torch.allclose(lora_b_list[2], torch.ones((1, rank)) * 4 * 2.0) + for index, name in enumerate(("to_q", "to_k", "to_v")): + if name in target_modules: + weights = loras[f"{name_prefix}{name}"] + assert torch.equal(lora_a_list[index], weights.lora_a) + assert torch.equal(lora_b_list[index], weights.lora_b * 2.0) + else: + assert lora_a_list[index] is None + assert lora_b_list[index] is None + assert manager._active_adapter_id == 1 assert bound_names == frozenset(loras) @@ -423,6 +434,7 @@ def test_lora_manager_rejects_adapter_that_binds_no_layer(): ) layer = _DummyLoRALayer(n_slices=1, output_slices=(2,)) manager._lora_modules = {"transformer.blocks.0.attn.to_q": layer} + manager._expected_lora_modules = {"to_q"} # Adapter-side name the engine does not expose, e.g. a diffusers-style # checkpoint against a differently named engine layout. @@ -445,11 +457,13 @@ def test_lora_manager_rejects_adapter_that_binds_no_layer(): )() } - with pytest.raises(ValueError, match="applies to no layer") as excinfo: + with pytest.raises(ValueError, match="binding is incomplete") as excinfo: manager._activate_adapter(3, scale=1.0) # The message must name what was received so the mismatch is diagnosable. + assert "bound=0/1" in str(excinfo.value) assert "unet.down_blocks.0.attn.to_q" in str(excinfo.value) + assert "expected target modules in ['to_q']" in str(excinfo.value) # Nothing was bound and the adapter must not be left marked active. assert manager._active_adapter_id is None @@ -457,6 +471,97 @@ def test_lora_manager_rejects_adapter_that_binds_no_layer(): assert layer.reset_calls >= 1 +def test_lora_manager_rejects_empty_adapter(): + """An empty adapter must not be considered successfully bound.""" + manager = DiffusionLoRAManager( + pipeline=torch.nn.Module(), + device=torch.device("cpu"), + dtype=torch.bfloat16, + max_cached_adapters=1, + ) + layer = _DummyLoRALayer(n_slices=1, output_slices=(2,)) + manager._lora_modules = {"transformer.attn.to_q": layer} + manager._registered_adapters = {1: LoRAModel(1, rank=2, loras={})} + + with pytest.raises(ValueError, match="bound=0/0"): + manager._activate_adapter(1, scale=1.0) + + assert manager._active_adapter_id is None + assert manager._suspended_adapter_id is None + assert 1 not in manager._adapter_scales + assert layer.set_calls == [] + assert layer.active_slices == () + assert layer.reset_calls >= 1 + + +@pytest.mark.parametrize("suspend_previous", [False, True]) +@pytest.mark.parametrize("unbound_name", ["transformer.attn.to_out.0", "transformer.attn.to_qkv"]) +def test_lora_manager_rejects_partial_binding_and_rolls_back(unbound_name, suspend_previous): + """A partial bind must clear uploaded weights and invalidate fast paths.""" + manager = DiffusionLoRAManager( + pipeline=torch.nn.Module(), + device=torch.device("cpu"), + dtype=torch.bfloat16, + max_cached_adapters=2, + ) + layer = _DummyLoRALayer(n_slices=1, output_slices=(2,)) + packed_layer = _DummyLoRALayer(n_slices=3, output_slices=(2, 1, 1)) + manager._lora_modules = { + "transformer.attn.to_out": layer, + "transformer.attn.to_qkv": packed_layer, + } + manager._expected_lora_modules = {"to_out", "to_qkv"} + + def weights(name, value): + return LoRALayerWeights( + module_name=name, + rank=2, + lora_alpha=2, + lora_a=torch.full((2, 2), value), + lora_b=torch.full((2, 2), value), + ) + + valid_name = "transformer.attn.to_out" + previous = LoRAModel(1, rank=2, loras={valid_name: weights(valid_name, 1.0)}) + # The extra weights either use an unmatched checkpoint name or have a B + # shape that cannot be split across the fused layer's four output rows. + partial = LoRAModel( + 2, + rank=2, + loras={valid_name: weights(valid_name, 2.0), unbound_name: weights(unbound_name, 2.0)}, + ) + manager._registered_adapters = {1: previous, 2: partial} + manager._activate_adapter(1, scale=0.5) + if suspend_previous: + manager._deactivate_all_adapters() + + with pytest.raises(ValueError, match="binding is incomplete") as excinfo: + manager._activate_adapter(2, scale=1.0) + + assert "LoRA adapter 2" in str(excinfo.value) + assert unbound_name in str(excinfo.value) + assert valid_name not in str(excinfo.value).replace(unbound_name, "") + if unbound_name.endswith("to_qkv"): + assert "lora_b.shape[0]=2" in str(excinfo.value) + assert "sum(output_slices)=4" in str(excinfo.value) + assert "output_slices=(2, 1, 1)" in str(excinfo.value) + else: + assert "expected target modules in ['to_out', 'to_qkv']" in str(excinfo.value) + assert len(layer.set_calls) == 2 + assert layer.reset_calls == 1 + assert layer.active_slices == () + assert packed_layer.active_slices == () + assert manager._active_adapter_id is None + assert manager._suspended_adapter_id is None + assert 2 not in manager._adapter_scales + + manager._activate_adapter(1, scale=0.5) + assert len(layer.set_calls) == 3 + assert layer.resume_calls == 0 + assert manager._active_adapter_id == 1 + assert torch.equal(layer.set_calls[-1][1], previous.loras[valid_name].lora_b * 0.5) + + def _dummy_lora_request(adapter_id: int) -> LoRARequest: return LoRARequest( lora_name=f"adapter_{adapter_id}", @@ -684,7 +789,7 @@ def test_lora_manager_rejects_hunyuan_image3_qkv_lora_with_unexpected_rows(): )() } - with pytest.raises(ValueError, match="applies to no layer"): + with pytest.raises(ValueError, match="cannot establish HunyuanImage-3 fused-QKV layout"): manager._activate_adapter(12, 0.5) assert packed_layer.set_calls == [] diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py index 28f074a19..aeef47c6a 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_native_lora.py @@ -710,6 +710,8 @@ def test_lora_manager_activates_native_packed_qkv(tmp_path): assert isinstance(packed, PackedLoRALayerWeights) assert len(packed.lora_b) == 3 assert all(b.shape[0] == _TINY_QKV_SLICE for b in packed.lora_b) + # This fixture only exposes one layer, so activate only its LoRA weights. + lora_model.loras = {"blocks.0.attn.qkv_proj": packed} class _DummyLoRALayer: def __init__(self): diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 9d6b2af92..5670bdacd 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -573,8 +573,7 @@ class DiffusionLoRAManager: def _bind_adapter_weights(self, lora_model: LoRAModel, scale: float) -> None: binding_validator = getattr(self.pipeline, "_validate_diffusion_lora_binding", None) - # Track bindings unconditionally. The zero-binding guard below needs this - # bookkeeping on every pipeline, not only the ones that supply a validator. + # Track successful bindings for generic and model-specific validation. lora_names_by_id = {id(weights): name for name, weights in lora_model.loras.items()} bound_lora_names: set[str] = set() @@ -680,24 +679,21 @@ class DiffusionLoRAManager: or len(output_sizes) != n_slices or deinterleaved.shape[0] != sum(output_sizes) ): - logger.warning( - "Skipping LoRA for %s: cannot establish HunyuanImage-3 fused-QKV layout", - full_module_name, + raise ValueError( + f"LoRA adapter {lora_model.id} binding is incomplete for {full_module_name}: " + "cannot establish HunyuanImage-3 fused-QKV layout " + f"(lora_b.shape[0]={lora_weights.lora_b.shape[0]}, " + f"expected output_sizes={output_sizes})" ) - lora_layer.reset_lora(0) - continue b_splits = list(torch.split(deinterleaved, list(output_sizes), dim=0)) else: total = sum(output_slices) if lora_weights.lora_b.shape[0] != total: - logger.warning( - "Skipping LoRA for %s due to shape mismatch: lora_b[0]=%d != sum(output_slices)=%d", - full_module_name, - lora_weights.lora_b.shape[0], - total, + raise ValueError( + f"LoRA adapter {lora_model.id} binding is incomplete for {full_module_name}: " + f"lora_b.shape[0]={lora_weights.lora_b.shape[0]} != " + f"sum(output_slices)={total} for output_slices={tuple(output_slices)}" ) - lora_layer.reset_lora(0) - continue b_splits = list(torch.split(lora_weights.lora_b, list(output_slices), dim=0)) lora_a_list = [lora_weights.lora_a] * n_slices @@ -722,11 +718,13 @@ class DiffusionLoRAManager: scale, ) - if not bound_lora_names: + unbound_lora_names = sorted(set(lora_model.loras) - bound_lora_names) + if not bound_lora_names or unbound_lora_names: raise ValueError( - f"LoRA adapter {lora_model.id} applies to no layer: expected target modules in " - f"{sorted(self._expected_lora_modules)} but received {sorted(lora_model.loras)}. " - "Activating it would leave the base model unchanged." + f"LoRA adapter {lora_model.id} binding is incomplete: " + f"bound={len(bound_lora_names)}/{len(lora_model.loras)}, " + f"unbound modules={unbound_lora_names}; " + f"expected target modules in {sorted(self._expected_lora_modules)}" ) if callable(binding_validator):