1237 lines
45 KiB
Python
1237 lines
45 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
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
|
|
|
|
from tests.diffusion.lora.helpers import (
|
|
DummyBaseLayerWithLoRA,
|
|
FakeLinearBase,
|
|
fake_replace_submodule,
|
|
)
|
|
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
|
|
from vllm_omni.lora.request import LoRARequest
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class _DummyLoRALayer:
|
|
base_layer: torch.nn.Module
|
|
|
|
def __init__(self, n_slices: int, output_slices: tuple[int, ...], tp_size: int = 1):
|
|
self.n_slices = n_slices
|
|
self.output_slices = output_slices
|
|
self.tp_size = tp_size
|
|
self.set_calls: list[
|
|
tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
|
|
] = []
|
|
self.reset_calls: int = 0
|
|
self.suspend_calls: int = 0
|
|
self.resume_calls: int = 0
|
|
self.active_slices: tuple[bool, ...] = ()
|
|
self.suspended_slices: tuple[bool, ...] | None = None
|
|
|
|
def set_lora(self, index: int, lora_a, lora_b):
|
|
assert index == 0
|
|
self.set_calls.append((lora_a, lora_b))
|
|
if isinstance(lora_b, list):
|
|
self.active_slices = tuple(b is not None for b in lora_b)
|
|
else:
|
|
self.active_slices = (True,)
|
|
self.suspended_slices = None
|
|
|
|
def reset_lora(self, index: int):
|
|
assert index == 0
|
|
self.reset_calls += 1
|
|
self.active_slices = ()
|
|
self.suspended_slices = None
|
|
|
|
def suspend_lora(self) -> None:
|
|
if self.suspended_slices is not None:
|
|
return
|
|
self.suspended_slices = self.active_slices
|
|
self.active_slices = (False,) * len(self.active_slices)
|
|
self.suspend_calls += 1
|
|
|
|
def resume_lora(self) -> None:
|
|
if self.suspended_slices is not None:
|
|
self.active_slices = self.suspended_slices
|
|
self.suspended_slices = None
|
|
self.resume_calls += 1
|
|
|
|
|
|
# Aliases for backward compatibility within this file
|
|
_FakeLinearBase = FakeLinearBase
|
|
_DummyBaseLayerWithLoRA = DummyBaseLayerWithLoRA
|
|
|
|
|
|
class _DummyPipeline(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.transformer = torch.nn.Module()
|
|
self.transformer.foo = _FakeLinearBase()
|
|
|
|
|
|
class _DummyLM(torch.nn.Module):
|
|
"""LoRA enabled wrapper for _DummyPipeline."""
|
|
|
|
def __init__(self, rank: int):
|
|
super().__init__()
|
|
self.transformer = torch.nn.Module()
|
|
self.transformer.foo = _DummyBaseLayerWithLoRA(_FakeLinearBase())
|
|
self.rank = rank
|
|
self.loras = self.get_lora_modules()
|
|
|
|
def get_lora_modules(self):
|
|
return {"transformer.foo": self._get_initial_lora(self.rank)}
|
|
|
|
def get_lora(self, k: str) -> LoRALayerWeights:
|
|
"""Get the unscaled LoRA weights for transformer.foo"""
|
|
return self.loras[k]
|
|
|
|
def _get_initial_lora(self, rank: int) -> LoRALayerWeights:
|
|
"""Initializes a dummy LoRA for the current rank."""
|
|
A = torch.ones((rank, 4))
|
|
B = torch.ones((4, rank))
|
|
return LoRALayerWeights(
|
|
module_name="foo",
|
|
rank=rank,
|
|
lora_alpha=rank,
|
|
lora_a=A,
|
|
lora_b=B,
|
|
)
|
|
|
|
|
|
def test_lora_manager_supported_modules_are_stable_with_wrapped_layers(monkeypatch):
|
|
# Simulate a pipeline that already contains LoRA wrappers where the original
|
|
# LinearBase is nested under ".base_layer".
|
|
import vllm_omni.diffusion.lora.manager as manager_mod
|
|
|
|
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
|
|
|
|
pipeline = _DummyLM(rank=2)
|
|
|
|
# vLLM helper would see only the nested LinearBase and yield "base_layer".
|
|
assert get_supported_lora_modules(pipeline) == ["base_layer"]
|
|
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=pipeline,
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
max_cached_adapters=1,
|
|
)
|
|
|
|
assert "foo" in manager._supported_lora_modules
|
|
assert "base_layer" not in manager._supported_lora_modules
|
|
|
|
|
|
def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
|
|
import vllm_omni.diffusion.lora.manager as manager_mod
|
|
|
|
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
|
|
|
|
def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
|
|
if isinstance(layer, _FakeLinearBase):
|
|
return _DummyBaseLayerWithLoRA(layer)
|
|
return layer
|
|
|
|
replace_calls: list[str] = []
|
|
|
|
def _fake_replace_submodule(root: torch.nn.Module, module_name: str, submodule: torch.nn.Module):
|
|
replace_calls.append(module_name)
|
|
setattr(root, module_name, submodule)
|
|
|
|
monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
|
|
monkeypatch.setattr(manager_mod, "replace_submodule", _fake_replace_submodule)
|
|
|
|
pipeline = torch.nn.Module()
|
|
pipeline.transformer = torch.nn.Module()
|
|
pipeline.transformer.foo = _FakeLinearBase()
|
|
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=pipeline,
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
max_cached_adapters=1,
|
|
)
|
|
|
|
peft_helper = type("_PH", (), {"r": 1})()
|
|
|
|
manager._replace_layers_with_lora(peft_helper)
|
|
manager._replace_layers_with_lora(peft_helper)
|
|
|
|
# Only the top-level layer should have been replaced; nested ".base_layer"
|
|
# must be skipped to avoid nesting LoRA wrappers.
|
|
assert replace_calls == ["foo"]
|
|
|
|
|
|
def test_lora_manager_keeps_dlo_lora_buffers_on_compute_device(monkeypatch):
|
|
import vllm_omni.diffusion.lora.manager as manager_mod
|
|
|
|
class _DLOLoRALayer(DummyBaseLayerWithLoRA):
|
|
lora_buffer_device = None
|
|
|
|
def _set_diffusion_lora_buffer_device(self, device):
|
|
self.lora_buffer_device = device
|
|
|
|
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DLOLoRALayer)
|
|
monkeypatch.setattr(
|
|
manager_mod,
|
|
"from_layer_diffusion",
|
|
lambda *, layer, **_kwargs: _DLOLoRALayer(layer),
|
|
)
|
|
|
|
pipeline = torch.nn.Module()
|
|
pipeline.od_config = SimpleNamespace(enable_distributed_layerwise_offload=True)
|
|
pipeline.transformer = torch.nn.Module()
|
|
pipeline.transformer.foo = _FakeLinearBase()
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=pipeline,
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
)
|
|
|
|
manager._replace_layers_with_lora(type("_PH", (), {"r": 1})())
|
|
|
|
assert pipeline.transformer.foo.lora_buffer_device == torch.device("cpu")
|
|
|
|
|
|
def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch):
|
|
import vllm_omni.diffusion.lora.manager as manager_mod
|
|
|
|
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
|
|
|
|
def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
|
|
return _DummyBaseLayerWithLoRA(layer)
|
|
|
|
replace_calls: list[str] = []
|
|
|
|
def _fake_replace_submodule(root: torch.nn.Module, module_name: str, submodule: torch.nn.Module):
|
|
replace_calls.append(module_name)
|
|
setattr(root, module_name, submodule)
|
|
|
|
monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
|
|
monkeypatch.setattr(manager_mod, "replace_submodule", _fake_replace_submodule)
|
|
|
|
pipeline = torch.nn.Module()
|
|
pipeline.stacked_params_mapping = [
|
|
(".to_qkv.", ".to_q.", "q"),
|
|
(".to_qkv.", ".to_k.", "k"),
|
|
(".to_qkv.", ".to_v.", "v"),
|
|
]
|
|
pipeline.transformer = torch.nn.Module()
|
|
pipeline.transformer.to_qkv = _FakeLinearBase()
|
|
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=pipeline,
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
max_cached_adapters=1,
|
|
)
|
|
|
|
# Treat the dummy layer as a packed 3-slice projection so the manager uses
|
|
# `stacked_params_mapping` to decide replacement based on target_modules.
|
|
monkeypatch.setattr(manager, "_get_packed_modules_list", lambda _module: ["q", "k", "v"])
|
|
|
|
peft_helper = type("_PH", (), {"r": 1, "target_modules": ["to_q"]})()
|
|
manager._replace_layers_with_lora(peft_helper)
|
|
|
|
assert replace_calls == ["to_qkv"]
|
|
|
|
|
|
def test_lora_manager_activates_fused_lora_on_packed_layer():
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=torch.nn.Module(),
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
max_cached_adapters=1,
|
|
)
|
|
|
|
packed_layer = _DummyLoRALayer(n_slices=3, output_slices=(2, 1, 1))
|
|
manager._lora_modules = {"transformer.blocks.0.attn.to_qkv": packed_layer}
|
|
|
|
rank = 2
|
|
A = torch.ones((rank, 4))
|
|
B = torch.arange(0, sum(packed_layer.output_slices) * rank, dtype=torch.bfloat16).view(-1, rank)
|
|
lora = LoRALayerWeights(
|
|
module_name="transformer.blocks.0.attn.to_qkv",
|
|
rank=rank,
|
|
lora_alpha=rank,
|
|
lora_a=A,
|
|
lora_b=B,
|
|
)
|
|
manager._registered_adapters = {
|
|
7: type(
|
|
"LM",
|
|
(),
|
|
{
|
|
"id": 7,
|
|
"loras": {"transformer.blocks.0.attn.to_qkv": lora},
|
|
"get_lora": lambda self, k: self.loras.get(k),
|
|
},
|
|
)()
|
|
}
|
|
|
|
manager._activate_adapter(7, 0.5)
|
|
|
|
assert packed_layer.reset_calls == 0
|
|
assert len(packed_layer.set_calls) == 1
|
|
lora_a_list, lora_b_list = packed_layer.set_calls[0]
|
|
assert isinstance(lora_a_list, list)
|
|
assert isinstance(lora_b_list, list)
|
|
assert len(lora_a_list) == 3
|
|
assert len(lora_b_list) == 3
|
|
assert all(torch.allclose(a, A) for a in lora_a_list)
|
|
# B should be split into 3 slices and scaled.
|
|
b0, b1, b2 = lora_b_list
|
|
assert b0.shape[0] == 2 and b1.shape[0] == 1 and b2.shape[0] == 1
|
|
assert torch.allclose(torch.cat([b0, b1, b2], dim=0), B * 0.5)
|
|
|
|
|
|
@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
|
|
|
|
def validate_binding(*, lora_model, bound_lora_names):
|
|
nonlocal bound_names
|
|
assert lora_model.id == 1
|
|
bound_names = bound_lora_names
|
|
|
|
pipeline._validate_diffusion_lora_binding = validate_binding
|
|
pipeline.stacked_params_mapping = [
|
|
(".to_qkv", ".to_q", "q"),
|
|
(".to_qkv", ".to_k", "k"),
|
|
(".to_qkv", ".to_v", "v"),
|
|
]
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=pipeline,
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
max_cached_adapters=1,
|
|
)
|
|
|
|
packed_layer = _DummyLoRALayer(n_slices=3, output_slices=(2, 1, 1))
|
|
manager._lora_modules = {"transformer.blocks.0.attn.to_qkv": packed_layer}
|
|
|
|
rank = 2
|
|
loras: dict[str, LoRALayerWeights] = {}
|
|
for name, out_dim in zip(["to_q", "to_k", "to_v"], [2, 1, 1]):
|
|
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),
|
|
lora_b=torch.ones((out_dim, rank)) * (3 if name == "to_q" else 4),
|
|
)
|
|
|
|
manager._registered_adapters = {
|
|
1: type("LM", (), {"id": 1, "loras": loras, "get_lora": lambda self, k: self.loras.get(k)})()
|
|
}
|
|
|
|
manager._activate_adapter(1, scale=2.0)
|
|
|
|
assert packed_layer.reset_calls == 0
|
|
assert len(packed_layer.set_calls) == 1
|
|
lora_a_list, lora_b_list = packed_layer.set_calls[0]
|
|
assert isinstance(lora_a_list, list)
|
|
assert isinstance(lora_b_list, list)
|
|
assert len(lora_a_list) == 3
|
|
assert len(lora_b_list) == 3
|
|
# Scale should apply to B only.
|
|
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)
|
|
|
|
|
|
def test_lora_manager_rolls_back_all_layers_when_activation_fails():
|
|
class _FailingLoRALayer(_DummyLoRALayer):
|
|
fail = False
|
|
|
|
def set_lora(self, index: int, lora_a, lora_b):
|
|
if self.fail:
|
|
raise RuntimeError("bind failed")
|
|
super().set_lora(index, lora_a, lora_b)
|
|
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=torch.nn.Module(),
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
max_cached_adapters=2,
|
|
)
|
|
first = _DummyLoRALayer(n_slices=1, output_slices=(2,))
|
|
second = _FailingLoRALayer(n_slices=1, output_slices=(2,))
|
|
manager._lora_modules = {
|
|
"transformer.first": first,
|
|
"transformer.second": second,
|
|
}
|
|
|
|
def adapter(adapter_id: int):
|
|
loras = {
|
|
name: LoRALayerWeights(
|
|
module_name=name,
|
|
rank=2,
|
|
lora_alpha=2,
|
|
lora_a=torch.full((2, 2), float(adapter_id)),
|
|
lora_b=torch.full((2, 2), float(adapter_id)),
|
|
)
|
|
for name in manager._lora_modules
|
|
}
|
|
return type(
|
|
"LM",
|
|
(),
|
|
{
|
|
"id": adapter_id,
|
|
"loras": loras,
|
|
"get_lora": lambda self, key: self.loras.get(key),
|
|
},
|
|
)()
|
|
|
|
manager._registered_adapters = {1: adapter(1), 2: adapter(2)}
|
|
manager._activate_adapter(1, scale=1.0)
|
|
first_calls_after_success = len(first.set_calls)
|
|
|
|
second.fail = True
|
|
with pytest.raises(RuntimeError, match="bind failed"):
|
|
manager._activate_adapter(2, scale=1.0)
|
|
|
|
assert manager._active_adapter_id is None
|
|
assert first.reset_calls == 1
|
|
assert second.reset_calls == 1
|
|
|
|
# The failed activation must not leave the old adapter eligible for the
|
|
# fast path: reactivating it binds every layer again.
|
|
second.fail = False
|
|
manager._activate_adapter(1, scale=1.0)
|
|
assert manager._active_adapter_id == 1
|
|
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}
|
|
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.
|
|
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="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
|
|
assert layer.set_calls == []
|
|
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}",
|
|
lora_int_id=adapter_id,
|
|
lora_path=f"/tmp/adapter_{adapter_id}",
|
|
)
|
|
|
|
|
|
def test_lora_manager_evicts_lru_adapter_when_cache_full(monkeypatch):
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=torch.nn.Module(),
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
max_cached_adapters=2,
|
|
)
|
|
|
|
def _fake_load(_req: LoRARequest):
|
|
lora_model = type("LM", (), {"id": _req.lora_int_id})()
|
|
peft_helper = type("PH", (), {})()
|
|
return lora_model, peft_helper
|
|
|
|
monkeypatch.setattr(manager, "_load_adapter", _fake_load)
|
|
monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
|
|
monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id, scale: None)
|
|
|
|
req1 = _dummy_lora_request(1)
|
|
req2 = _dummy_lora_request(2)
|
|
req3 = _dummy_lora_request(3)
|
|
|
|
manager.set_active_adapter(req1, lora_scale=1.0)
|
|
manager.set_active_adapter(req2, lora_scale=1.0)
|
|
|
|
# Touch adapter 1 so adapter 2 becomes LRU.
|
|
manager.set_active_adapter(req1, lora_scale=1.0)
|
|
|
|
manager.set_active_adapter(req3, lora_scale=1.0)
|
|
|
|
assert set(manager.list_adapters()) == {1, 3}
|
|
|
|
|
|
def test_lora_manager_does_not_evict_pinned_adapter(monkeypatch):
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=torch.nn.Module(),
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
max_cached_adapters=2,
|
|
)
|
|
|
|
def _fake_load(_req: LoRARequest):
|
|
lora_model = type("LM", (), {"id": _req.lora_int_id})()
|
|
peft_helper = type("PH", (), {})()
|
|
return lora_model, peft_helper
|
|
|
|
monkeypatch.setattr(manager, "_load_adapter", _fake_load)
|
|
monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
|
|
monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id, scale: None)
|
|
|
|
manager.set_active_adapter(_dummy_lora_request(1), lora_scale=1.0)
|
|
assert manager.pin_adapter(1)
|
|
|
|
manager.set_active_adapter(_dummy_lora_request(2), lora_scale=1.0)
|
|
manager.set_active_adapter(_dummy_lora_request(3), lora_scale=1.0)
|
|
|
|
assert set(manager.list_adapters()) == {1, 3}
|
|
|
|
|
|
def test_lora_manager_warns_when_all_adapters_pinned(monkeypatch):
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=torch.nn.Module(),
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
max_cached_adapters=2,
|
|
)
|
|
|
|
def _fake_load(_req: LoRARequest):
|
|
lora_model = type("LM", (), {"id": _req.lora_int_id})()
|
|
peft_helper = type("PH", (), {})()
|
|
return lora_model, peft_helper
|
|
|
|
monkeypatch.setattr(manager, "_load_adapter", _fake_load)
|
|
monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
|
|
monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id, scale: None)
|
|
|
|
manager.set_active_adapter(_dummy_lora_request(1), lora_scale=1.0)
|
|
manager.set_active_adapter(_dummy_lora_request(2), lora_scale=1.0)
|
|
|
|
assert manager.pin_adapter(1)
|
|
assert manager.pin_adapter(2)
|
|
|
|
manager.max_cached_adapters = 1
|
|
manager._evict_for_new_adapter()
|
|
|
|
assert set(manager.list_adapters()) == {1, 2}
|
|
|
|
|
|
def _make_hunyuan_image3_pipeline(num_heads=4, num_kv_heads=2, head_dim=2, hidden_size=8):
|
|
"""Use production hooks and the production base-weight splitter without
|
|
allocating the full checkpoint or initializing distributed workers.
|
|
"""
|
|
from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import HunyuanImage3Model
|
|
from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import HunyuanImage3Pipeline
|
|
|
|
model = HunyuanImage3Model.__new__(HunyuanImage3Model)
|
|
torch.nn.Module.__init__(model)
|
|
model.config = SimpleNamespace(
|
|
num_attention_heads=num_heads,
|
|
num_key_value_heads=num_kv_heads,
|
|
attention_head_dim=head_dim,
|
|
hidden_size=hidden_size,
|
|
)
|
|
pipeline = HunyuanImage3Pipeline.__new__(HunyuanImage3Pipeline)
|
|
torch.nn.Module.__init__(pipeline)
|
|
pipeline.model = model
|
|
pipeline.transformer = model
|
|
return pipeline
|
|
|
|
|
|
def _make_hunyuan_image3_qkv_base():
|
|
"""An un-initialized QKVParallelLinear instance for isinstance() checks.
|
|
|
|
Supply the type and checkpoint head metadata for binding-only tests.
|
|
The PEFT forward test below constructs the real distributed projection.
|
|
"""
|
|
from vllm.model_executor.layers.linear import QKVParallelLinear
|
|
|
|
layer = object.__new__(QKVParallelLinear)
|
|
layer.total_num_heads = 4
|
|
layer.total_num_kv_heads = 2
|
|
layer.head_size = layer.v_head_size = 2
|
|
return layer
|
|
|
|
|
|
def test_lora_manager_activates_hunyuan_image3_fused_qkv_lora():
|
|
"""A fused HI3 ``qkv_proj`` LoRA-B (GQA-interleaved in the adapter) is
|
|
de-interleaved to the block [Q; K; V] layout the QKV output slices expect,
|
|
with the namespace alias resolving the PEFT adapter tensors."""
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=_make_hunyuan_image3_pipeline(),
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
max_cached_adapters=1,
|
|
)
|
|
|
|
wrapper = "transformer.layers.0.self_attn.qkv_proj"
|
|
packed_layer = _DummyLoRALayer(
|
|
n_slices=3,
|
|
output_slices=(8, 4, 4),
|
|
tp_size=1,
|
|
)
|
|
# Use QKV type identity and checkpoint head metadata for the manager's
|
|
# de-interleave branch; the forward test below constructs real layers.
|
|
packed_layer.base_layer = _make_hunyuan_image3_qkv_base()
|
|
manager._lora_modules = {wrapper: packed_layer}
|
|
|
|
rank = 3
|
|
# GQA-interleaved fused QKV rows: [Q-group0, K0, V0, Q-group1, K1, V1].
|
|
interleaved = torch.arange(16, dtype=torch.bfloat16).unsqueeze(1).repeat(1, rank)
|
|
lora = LoRALayerWeights(
|
|
module_name="model.layers.0.self_attn.qkv_proj",
|
|
rank=rank,
|
|
lora_alpha=rank,
|
|
lora_a=torch.ones((rank, 8)),
|
|
lora_b=interleaved,
|
|
)
|
|
manager._registered_adapters = {
|
|
11: type(
|
|
"LM",
|
|
(),
|
|
{
|
|
"id": 11,
|
|
"loras": {"model.layers.0.self_attn.qkv_proj": lora},
|
|
"get_lora": lambda self, k: self.loras.get(k),
|
|
},
|
|
)()
|
|
}
|
|
|
|
manager._activate_adapter(11, 0.5)
|
|
|
|
assert packed_layer.reset_calls == 0
|
|
assert len(packed_layer.set_calls) == 1
|
|
lora_a_list, lora_b_list = packed_layer.set_calls[0]
|
|
assert len(lora_a_list) == 3 and len(lora_b_list) == 3
|
|
# De-interleaved full [Q; K; V] slices, scaled by the lora scale.
|
|
expected_q = torch.cat([interleaved[0:4], interleaved[8:12]], dim=0) * 0.5
|
|
expected_k = torch.cat([interleaved[4:6], interleaved[12:14]], dim=0) * 0.5
|
|
expected_v = torch.cat([interleaved[6:8], interleaved[14:16]], dim=0) * 0.5
|
|
assert torch.equal(lora_b_list[0], expected_q)
|
|
assert torch.equal(lora_b_list[1], expected_k)
|
|
assert torch.equal(lora_b_list[2], expected_v)
|
|
|
|
|
|
def test_lora_manager_rejects_hunyuan_image3_qkv_lora_with_unexpected_rows():
|
|
"""If the fused HI3 QKV layout cannot be established (wrong row count), the
|
|
manager must fail closed instead of applying the LoRA to wrong slices."""
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=_make_hunyuan_image3_pipeline(),
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
max_cached_adapters=1,
|
|
)
|
|
|
|
wrapper = "transformer.layers.0.self_attn.qkv_proj"
|
|
packed_layer = _DummyLoRALayer(n_slices=3, output_slices=(8, 4, 4))
|
|
packed_layer.base_layer = _make_hunyuan_image3_qkv_base()
|
|
manager._lora_modules = {wrapper: packed_layer}
|
|
|
|
rank = 3
|
|
# Row count (12) does not match the HI3 QKV layout (16): fail closed.
|
|
lora = LoRALayerWeights(
|
|
module_name="model.layers.0.self_attn.qkv_proj",
|
|
rank=rank,
|
|
lora_alpha=rank,
|
|
lora_a=torch.ones((rank, 8)),
|
|
lora_b=torch.arange(12 * rank, dtype=torch.bfloat16).view(-1, rank),
|
|
)
|
|
manager._registered_adapters = {
|
|
12: type(
|
|
"LM",
|
|
(),
|
|
{
|
|
"id": 12,
|
|
"loras": {"model.layers.0.self_attn.qkv_proj": lora},
|
|
"get_lora": lambda self, k: self.loras.get(k),
|
|
},
|
|
)()
|
|
}
|
|
|
|
with pytest.raises(ValueError, match="cannot establish HunyuanImage-3 fused-QKV layout"):
|
|
manager._activate_adapter(12, 0.5)
|
|
|
|
assert packed_layer.set_calls == []
|
|
assert packed_layer.reset_calls >= 1
|
|
assert manager._active_adapter_id is None
|
|
|
|
|
|
def test_lora_manager_applies_multiple_scales_correctly(monkeypatch):
|
|
"""Ensure that the LoRA manager applies scales correctly when the
|
|
active adapter receives a different scale, i.e., the rank is unchanged.
|
|
"""
|
|
import vllm_omni.diffusion.lora.manager as manager_mod
|
|
|
|
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyLoRALayer)
|
|
|
|
rank = 2
|
|
adapter_id = 7
|
|
req1 = _dummy_lora_request(adapter_id)
|
|
scale_1 = 0.25
|
|
scale_2 = 0.5
|
|
|
|
lora_model = _DummyLM(rank=rank)
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=_DummyPipeline(),
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
)
|
|
|
|
def _fake_load(_req: LoRARequest):
|
|
peft_helper = type("PH", (), {"r": rank})()
|
|
return lora_model, peft_helper
|
|
|
|
monkeypatch.setattr(manager, "_load_adapter", _fake_load)
|
|
manager._registered_adapters = {
|
|
adapter_id: lora_model,
|
|
}
|
|
manager._lora_modules = {"transformer.foo": lora_model.transformer.foo}
|
|
|
|
# After the first scale, all B values should go from 1 -> scale_1
|
|
manager.set_active_adapter(req1, lora_scale=scale_1)
|
|
assert len(lora_model.transformer.foo.set_calls) == 1
|
|
lora_a, lora_b = lora_model.transformer.foo.set_calls[0]
|
|
assert torch.all(lora_a == 1)
|
|
assert torch.all(lora_b == scale_1)
|
|
|
|
# After the second scale, all B values should go from 1 -> scale_2
|
|
manager.set_active_adapter(req1, lora_scale=scale_2)
|
|
assert len(lora_model.transformer.foo.set_calls) == 2
|
|
|
|
lora_a, lora_b = lora_model.transformer.foo.set_calls[1]
|
|
assert torch.all(lora_a == 1)
|
|
assert torch.all(lora_b == scale_2)
|
|
|
|
|
|
def test_lora_manager_scales_correctly_with_rank_changes(monkeypatch):
|
|
"""Ensure that the LoRA manager correctly handles scaling when the rank
|
|
is changed and the buffers are reset + we reactivate.
|
|
"""
|
|
import vllm_omni.diffusion.lora.manager as manager_mod
|
|
|
|
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
|
|
|
|
rank = 2
|
|
adapter_id = 7
|
|
req1 = _dummy_lora_request(adapter_id)
|
|
initial_scale = 0.5
|
|
|
|
lora_model = _DummyLM(rank=rank)
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=_DummyPipeline(),
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
)
|
|
|
|
def _fake_load(_req: LoRARequest):
|
|
peft_helper = type("PH", (), {"r": rank})()
|
|
return lora_model, peft_helper
|
|
|
|
monkeypatch.setattr(manager, "_load_adapter", _fake_load)
|
|
manager._registered_adapters = {
|
|
adapter_id: lora_model,
|
|
}
|
|
manager._lora_modules = {"transformer.foo": lora_model.transformer.foo}
|
|
|
|
# Activate adapter with initial scale
|
|
manager.set_active_adapter(req1, lora_scale=initial_scale)
|
|
assert lora_model.transformer.foo.create_calls == 0
|
|
assert len(lora_model.transformer.foo.set_calls) == 1
|
|
lora_a, lora_b = lora_model.transformer.foo.set_calls[0]
|
|
assert torch.all(lora_a == 1)
|
|
assert torch.all(lora_b == initial_scale)
|
|
|
|
# Increase the rank; this resets the buffers, so the adapter is activated again
|
|
manager._ensure_max_lora_rank(8)
|
|
|
|
# Ensure we actually took the rank expansion path, which recreates
|
|
# and sets the weight buffets, but that the scale didn't change
|
|
assert lora_model.transformer.foo.create_calls == 1
|
|
assert len(lora_model.transformer.foo.set_calls) == 2
|
|
lora_a, lora_b = lora_model.transformer.foo.set_calls[1]
|
|
assert torch.all(lora_a == 1)
|
|
assert torch.all(lora_b == initial_scale)
|
|
|
|
|
|
def test_scale_keys_are_rounded():
|
|
"""Ensure that added adapter scales are rounded to avoid lookup
|
|
issues due to precision differences, e.g., computed scales.
|
|
"""
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=_DummyPipeline(),
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
)
|
|
adapter_id = 1
|
|
# Currently we round keys to 3 decimal places
|
|
manager._update_adapter_scale(adapter_id, 0.0031)
|
|
assert manager._adapter_scales[adapter_id] == 0.003
|
|
|
|
|
|
def test_lora_manager_uses_valid_max_rank(monkeypatch):
|
|
"""Ensure that the LoRA manager uses a valid max rank for vLLM."""
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=torch.nn.Module(),
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
)
|
|
|
|
# Ensure that the rank is correctly adjusted to the smallest valid max rank
|
|
supported_max_rank = 64
|
|
unsupported_max_rank = 63
|
|
assert supported_max_rank in DiffusionLoRAManager._VALID_MAX_RANKS
|
|
assert unsupported_max_rank not in DiffusionLoRAManager._VALID_MAX_RANKS
|
|
|
|
def _fake_load(_req: LoRARequest):
|
|
lora_model = type("LM", (), {"id": _req.lora_int_id})()
|
|
peft_helper = type("PH", (), {"r": unsupported_max_rank})()
|
|
return lora_model, peft_helper
|
|
|
|
monkeypatch.setattr(manager, "_load_adapter", _fake_load)
|
|
req1 = _dummy_lora_request(1)
|
|
manager.add_adapter(req1)
|
|
assert manager._max_lora_rank == supported_max_rank
|
|
|
|
|
|
@pytest.mark.parametrize("rank", [-1, 0, DiffusionLoRAManager._VALID_MAX_RANKS[-1] + 1])
|
|
def test_lora_manager_max_rank_validation(monkeypatch, rank):
|
|
"""Check that invalid max ranks are handled correctly."""
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=torch.nn.Module(),
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
)
|
|
|
|
lora_rank = rank
|
|
|
|
def _fake_load(_req: LoRARequest):
|
|
lora_model = type("LM", (), {"id": _req.lora_int_id})()
|
|
peft_helper = type("PH", (), {"r": lora_rank})()
|
|
return lora_model, peft_helper
|
|
|
|
monkeypatch.setattr(manager, "_load_adapter", _fake_load)
|
|
req1 = _dummy_lora_request(1)
|
|
with pytest.raises(ValueError):
|
|
manager.add_adapter(req1)
|
|
|
|
|
|
def test_lora_manager_discovers_bagel_component(monkeypatch):
|
|
"""Verify that _replace_layers_with_lora finds layers under 'bagel'."""
|
|
import vllm_omni.diffusion.lora.manager as manager_mod
|
|
|
|
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
|
|
|
|
def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
|
|
if isinstance(layer, _FakeLinearBase):
|
|
return _DummyBaseLayerWithLoRA(layer)
|
|
return layer
|
|
|
|
replace_calls: list[str] = []
|
|
|
|
monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
|
|
monkeypatch.setattr(
|
|
manager_mod,
|
|
"replace_submodule",
|
|
lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls),
|
|
)
|
|
|
|
# Pipeline with a 'bagel' component (no 'transformer') declared via
|
|
# _lora_components, matching how BagelPipeline opts in.
|
|
pipeline = torch.nn.Module()
|
|
pipeline._lora_components = ["bagel"]
|
|
pipeline.bagel = torch.nn.Module()
|
|
pipeline.bagel.language_model = torch.nn.Module()
|
|
pipeline.bagel.language_model.qkv_proj = _FakeLinearBase()
|
|
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=pipeline,
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
max_cached_adapters=1,
|
|
)
|
|
|
|
peft_helper = type("_PH", (), {"r": 1})()
|
|
manager._replace_layers_with_lora(peft_helper)
|
|
|
|
assert "language_model.qkv_proj" in replace_calls
|
|
assert "bagel.language_model.qkv_proj" in manager._lora_modules
|
|
# Verify the module was actually replaced in the tree (not just recorded)
|
|
assert isinstance(pipeline.bagel.language_model.qkv_proj, _DummyBaseLayerWithLoRA)
|
|
|
|
|
|
def test_lora_manager_discovers_unet_component(monkeypatch):
|
|
"""Verify that _replace_layers_with_lora finds layers under 'unet'."""
|
|
import vllm_omni.diffusion.lora.manager as manager_mod
|
|
|
|
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
|
|
|
|
def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
|
|
if isinstance(layer, _FakeLinearBase):
|
|
return _DummyBaseLayerWithLoRA(layer)
|
|
return layer
|
|
|
|
replace_calls: list[str] = []
|
|
|
|
monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
|
|
monkeypatch.setattr(
|
|
manager_mod,
|
|
"replace_submodule",
|
|
lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls),
|
|
)
|
|
|
|
# Pipeline with a 'unet' component (no 'transformer')
|
|
pipeline = torch.nn.Module()
|
|
pipeline.unet = torch.nn.Module()
|
|
pipeline.unet.down_block = torch.nn.Module()
|
|
pipeline.unet.down_block.proj = _FakeLinearBase()
|
|
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=pipeline,
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
max_cached_adapters=1,
|
|
)
|
|
|
|
peft_helper = type("_PH", (), {"r": 1})()
|
|
manager._replace_layers_with_lora(peft_helper)
|
|
|
|
assert "down_block.proj" in replace_calls
|
|
assert "unet.down_block.proj" in manager._lora_modules
|
|
# Verify the module was actually replaced in the tree (not just recorded)
|
|
assert isinstance(pipeline.unet.down_block.proj, _DummyBaseLayerWithLoRA)
|
|
|
|
|
|
def _suspend_harness(monkeypatch, adapter_id: int = 7, rank: int = 2):
|
|
"""Manager wired to a single dummy layer, ready to activate `adapter_id`."""
|
|
import vllm_omni.diffusion.lora.manager as manager_mod
|
|
|
|
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
|
|
|
|
lora_model = _DummyLM(rank=rank)
|
|
manager = DiffusionLoRAManager(
|
|
pipeline=_DummyPipeline(),
|
|
device=torch.device("cpu"),
|
|
dtype=torch.bfloat16,
|
|
)
|
|
monkeypatch.setattr(manager, "_load_adapter", lambda _req: (lora_model, type("PH", (), {"r": rank})()))
|
|
manager._registered_adapters = {adapter_id: lora_model}
|
|
manager._lora_modules = {"transformer.foo": lora_model.transformer.foo}
|
|
return manager, lora_model.transformer.foo
|
|
|
|
|
|
def test_reactivating_same_adapter_resumes_without_rebinding(monkeypatch):
|
|
"""activate -> deactivate -> activate at the same scale must re-arm the
|
|
saved mask instead of re-uploading every layer."""
|
|
adapter_id = 7
|
|
req = _dummy_lora_request(adapter_id)
|
|
manager, layer = _suspend_harness(monkeypatch, adapter_id)
|
|
|
|
manager.set_active_adapter(req, lora_scale=0.5)
|
|
assert len(layer.set_calls) == 1
|
|
assert layer.active_slices == (True,)
|
|
|
|
manager.set_active_adapter(None)
|
|
assert manager._suspended_adapter_id == adapter_id
|
|
assert layer.suspend_calls == 1
|
|
assert layer.active_slices == (False,)
|
|
# Suspending must not tear the upload down.
|
|
assert layer.reset_calls == 0
|
|
|
|
manager.set_active_adapter(req, lora_scale=0.5)
|
|
assert layer.resume_calls == 1
|
|
assert len(layer.set_calls) == 1, "resume must not rebind"
|
|
assert layer.active_slices == (True,)
|
|
assert manager._active_adapter_id == adapter_id
|
|
assert manager._suspended_adapter_id is None
|
|
|
|
|
|
@pytest.mark.parametrize("second_scale", [0.5, 0.25])
|
|
def test_scale_or_adapter_change_forces_rebind(monkeypatch, second_scale):
|
|
"""A different scale, or a different adapter, must take the full bind
|
|
path rather than re-arming a stale mask."""
|
|
adapter_id = 7
|
|
manager, layer = _suspend_harness(monkeypatch, adapter_id)
|
|
|
|
manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5)
|
|
manager.set_active_adapter(None)
|
|
assert manager._suspended_adapter_id == adapter_id
|
|
|
|
# Same id at a new scale, or a different id entirely.
|
|
other_id = adapter_id if second_scale != 0.5 else adapter_id + 1
|
|
if other_id != adapter_id:
|
|
manager._registered_adapters[other_id] = manager._registered_adapters[adapter_id]
|
|
|
|
manager.set_active_adapter(_dummy_lora_request(other_id), lora_scale=second_scale)
|
|
assert len(layer.set_calls) == 2, "must rebind"
|
|
assert manager._suspended_adapter_id is None
|
|
assert manager._active_adapter_id == other_id
|
|
|
|
|
|
def test_partial_packed_mask_survives_suspend_resume(monkeypatch):
|
|
"""A packed layer where only some slices carry LoRA must come back with
|
|
the same mask, not an all-True one."""
|
|
adapter_id = 7
|
|
manager, layer = _suspend_harness(monkeypatch, adapter_id)
|
|
|
|
# Emulate a packed bind that left slice 1 empty.
|
|
layer.set_lora(0, [torch.ones(2, 2), None], [torch.ones(2, 2), None])
|
|
manager._active_adapter_id = adapter_id
|
|
manager._update_adapter_scale(adapter_id, 0.5)
|
|
assert layer.active_slices == (True, False)
|
|
|
|
manager.set_active_adapter(None)
|
|
assert layer.active_slices == (False, False)
|
|
|
|
manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5)
|
|
assert layer.active_slices == (True, False), "mask must survive unchanged"
|
|
|
|
|
|
def test_removing_suspended_adapter_drops_the_upload(monkeypatch):
|
|
"""A removed adapter can never be resumed, so its weights must not stay
|
|
in the stacked buffers."""
|
|
adapter_id = 7
|
|
manager, layer = _suspend_harness(monkeypatch, adapter_id)
|
|
|
|
manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5)
|
|
manager.set_active_adapter(None)
|
|
assert manager._suspended_adapter_id == adapter_id
|
|
|
|
manager.remove_adapter(adapter_id)
|
|
assert manager._suspended_adapter_id is None
|
|
assert layer.reset_calls == 1, "the upload must be torn down"
|
|
assert layer.suspended_slices is None
|
|
|
|
|
|
@pytest.mark.parametrize("tp_size", [1, 2, 4])
|
|
def test_hunyuan_image3_peft_qkv_forward(tmp_path, monkeypatch, tp_size):
|
|
"""Disk PEFT -> real manager/hooks -> vLLM sharding -> numerical forward.
|
|
|
|
TP4 also covers replicated KV heads. Only distributed rank discovery is
|
|
stubbed: projection construction, adapter loading and all LoRA operations
|
|
execute production code. The oracle indexes checkpoint rows explicitly.
|
|
"""
|
|
import json
|
|
|
|
from safetensors.torch import save_file
|
|
from vllm.config import VllmConfig, set_current_vllm_config
|
|
from vllm.model_executor import parameter
|
|
from vllm.model_executor.layers import linear
|
|
|
|
adapter_dir = tmp_path / "adapter"
|
|
adapter_dir.mkdir()
|
|
rank = 2
|
|
a = torch.arange(16, dtype=torch.float32).reshape(rank, 8) / 16
|
|
b = torch.arange(32, dtype=torch.float32).reshape(16, rank) / 32
|
|
prefix = "base_model.model.model.layers.0.self_attn.qkv_proj"
|
|
save_file(
|
|
{f"{prefix}.lora_A.weight": a, f"{prefix}.lora_B.weight": b},
|
|
str(adapter_dir / "adapter_model.safetensors"),
|
|
)
|
|
(adapter_dir / "adapter_config.json").write_text(
|
|
json.dumps({"r": rank, "lora_alpha": rank, "target_modules": ["qkv_proj"]})
|
|
)
|
|
request = LoRARequest(lora_name="hi3", lora_int_id=6411, lora_path=str(adapter_dir))
|
|
x = torch.arange(24, dtype=torch.float32).reshape(3, 8) / 24
|
|
q_rows = [0, 1, 2, 3, 8, 9, 10, 11]
|
|
k_rows = [4, 5, 12, 13]
|
|
v_rows = [6, 7, 14, 15]
|
|
for tp_rank in range(tp_size):
|
|
monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: tp_size)
|
|
monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: tp_rank)
|
|
monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: tp_rank)
|
|
monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: tp_size)
|
|
with set_current_vllm_config(VllmConfig()):
|
|
pipeline = _make_hunyuan_image3_pipeline()
|
|
base = linear.QKVParallelLinear(8, 2, 4, 2, bias=False, params_dtype=torch.float32)
|
|
base.weight.data.zero_()
|
|
block = torch.nn.Module()
|
|
block.self_attn = torch.nn.Module()
|
|
block.self_attn.qkv_proj = base
|
|
pipeline.model.layers = torch.nn.ModuleList([block])
|
|
manager = DiffusionLoRAManager(pipeline, torch.device("cpu"), torch.float32)
|
|
manager.set_active_adapter(request, lora_scale=0.5)
|
|
wrapped = pipeline.model.layers[0].self_attn.qkv_proj
|
|
assert wrapped is not base
|
|
q_size = 8 // tp_size
|
|
kv_size = 4 // min(tp_size, 2)
|
|
kv_rank = tp_rank // max(tp_size // 2, 1)
|
|
rows = (
|
|
q_rows[tp_rank * q_size : (tp_rank + 1) * q_size]
|
|
+ k_rows[kv_rank * kv_size : (kv_rank + 1) * kv_size]
|
|
+ v_rows[kv_rank * kv_size : (kv_rank + 1) * kv_size]
|
|
)
|
|
expected = (x @ a.T) @ b[rows].T * 0.5
|
|
with torch.inference_mode():
|
|
torch.testing.assert_close(wrapped.apply(x), expected)
|
|
manager.set_active_adapter(None)
|
|
torch.testing.assert_close(wrapped.apply(x), torch.zeros_like(expected))
|
|
manager.set_active_adapter(request, lora_scale=0.5)
|
|
torch.testing.assert_close(wrapped.apply(x), expected)
|
|
|
|
|
|
@pytest.mark.parametrize("projection", ["qkv_proj", "o_proj"])
|
|
def test_hunyuan_image3_lora_namespace_is_model_scoped(projection):
|
|
from vllm.lora.lora_model import LoRAModel
|
|
|
|
name = f"model.layers.0.self_attn.{projection}"
|
|
weights = LoRALayerWeights(name, 2, 2, torch.ones(2, 8), torch.ones(8, 2))
|
|
adapter = LoRAModel(6411, 2, {name: weights})
|
|
wrapper = f"transformer.layers.0.self_attn.{projection}"
|
|
hi3 = DiffusionLoRAManager(_make_hunyuan_image3_pipeline(), torch.device("cpu"), torch.float32)
|
|
other = DiffusionLoRAManager(_DummyPipeline(), torch.device("cpu"), torch.float32)
|
|
assert hi3._get_lora_weights(adapter, wrapper) is weights
|
|
assert other._get_lora_weights(adapter, wrapper) is None
|
|
|
|
|
|
@pytest.mark.parametrize("head_dim_key", ["head_dim", "attention_head_dim", None])
|
|
def test_hunyuan_image3_real_qkv_hook_head_dimension(head_dim_key):
|
|
pipeline = _make_hunyuan_image3_pipeline()
|
|
del pipeline.model.config.attention_head_dim
|
|
if head_dim_key is not None:
|
|
setattr(pipeline.model.config, head_dim_key, 2)
|
|
b = torch.arange(32, dtype=torch.float32).reshape(16, 2)
|
|
expected_rows = [0, 1, 2, 3, 8, 9, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15]
|
|
torch.testing.assert_close(pipeline._deinterleave_fused_qkv_lora_b(b), b[expected_rows])
|
|
assert pipeline._deinterleave_fused_qkv_lora_b(b[:-1]) is None
|