Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 44 additions & 0 deletions tests/diffusion/lora/test_lora_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -411,6 +411,50 @@ def adapter(adapter_id: int):
assert len(first.set_calls) == first_calls_after_success + 2


def test_lora_manager_rejects_adapter_that_binds_no_layer():
"""An adapter whose target modules match nothing must fail, not silently no-op."""
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.blocks.0.attn.to_q": layer}

# Adapter-side name the engine does not expose, e.g. a diffusers-style
# checkpoint against a differently named engine layout.
mismatched = LoRALayerWeights(
module_name="unet.down_blocks.0.attn.to_q",
rank=2,
lora_alpha=2,
lora_a=torch.ones((2, 2)),
lora_b=torch.ones((2, 2)),
)
manager._registered_adapters = {
3: type(
"LM",
(),
{
"id": 3,
"loras": {"unet.down_blocks.0.attn.to_q": mismatched},
"get_lora": lambda self, key: self.loras.get(key),
},
)()
}

with pytest.raises(ValueError, match="applies to no layer") as excinfo:
manager._activate_adapter(3, scale=1.0)

# The message must name what was received so the mismatch is diagnosable.
assert "unet.down_blocks.0.attn.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
assert layer.set_calls == []
assert layer.reset_calls >= 1


def _dummy_lora_request(adapter_id: int) -> LoRARequest:
return LoRARequest(
lora_name=f"adapter_{adapter_id}",
Expand Down
13 changes: 10 additions & 3 deletions vllm_omni/diffusion/lora/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -556,9 +556,9 @@ def _is_active_at_scale(self, adapter_id: int, scale: float) -> bool:

def _bind_adapter_weights(self, lora_model: LoRAModel, scale: float) -> None:
binding_validator = getattr(self.pipeline, "_validate_diffusion_lora_binding", None)
lora_names_by_id = (
{id(weights): name for name, weights in lora_model.loras.items()} if callable(binding_validator) else {}
)
# Track bindings unconditionally. The zero-binding guard below needs this
# bookkeeping on every pipeline, not only the ones that supply a validator.
lora_names_by_id = {id(weights): name for name, weights in lora_model.loras.items()}
bound_lora_names: set[str] = set()

def _record_bound(weights: LoRALayerWeights | PackedLoRALayerWeights) -> None:
Expand Down Expand Up @@ -677,6 +677,13 @@ def _record_bound(weights: LoRALayerWeights | PackedLoRALayerWeights) -> None:
scale,
)

if not bound_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."
)

if callable(binding_validator):
binding_validator(
lora_model=lora_model,
Expand Down
Loading