439 lines
16 KiB
Python
439 lines
16 KiB
Python
from collections import defaultdict
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
from diffusers.loaders.lora_conversion_utils import (
|
|
_convert_non_diffusers_qwen_lora_to_diffusers,
|
|
)
|
|
from pytest_mock import MockerFixture
|
|
from torch.testing import assert_close
|
|
|
|
from vllm_omni.diffusion.lora.loader import (
|
|
LoraLoaderMixin,
|
|
QwenImageLoraLoaderMixin,
|
|
WanLoraLoaderMixin,
|
|
_prepare_lora_delta,
|
|
_remap_state_dict_keys,
|
|
)
|
|
|
|
HEAD_DIM = 24
|
|
RANK = 2
|
|
NUM_LAYERS = 10
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class DummyTransformerBlock(nn.Module):
|
|
def __init__(
|
|
self,
|
|
d_model: int,
|
|
bias: bool = False,
|
|
stacked_params: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.attn = nn.Module()
|
|
if stacked_params:
|
|
self.attn.to_qkv = nn.Linear(d_model * 3, d_model, bias=bias)
|
|
else:
|
|
self.attn.to_q = nn.Linear(d_model, d_model, bias=bias)
|
|
self.attn.to_k = nn.Linear(d_model, d_model, bias=bias)
|
|
self.attn.to_v = nn.Linear(d_model, d_model, bias=bias)
|
|
self.attn.to_out = nn.Linear(d_model, d_model, bias=bias)
|
|
|
|
|
|
class DummyTransformer(nn.Module):
|
|
def __init__(self, num_layers: int, d_model: int, bias: bool = False, stacked_params: bool = False):
|
|
super().__init__()
|
|
|
|
self.blocks = nn.ModuleList(
|
|
[DummyTransformerBlock(d_model, bias=bias, stacked_params=stacked_params) for _ in range(num_layers)]
|
|
)
|
|
if stacked_params:
|
|
self.stacked_params_mapping = [
|
|
(".attn.to_qkv", ".attn.to_q", "q"),
|
|
(".attn.to_qkv", ".attn.to_k", "k"),
|
|
(".attn.to_qkv", ".attn.to_v", "v"),
|
|
]
|
|
|
|
|
|
class DummyPipeline(nn.Module, LoraLoaderMixin):
|
|
def __init__(self, num_layers: int, d_model: int, bias: bool = False, stacked_params: bool = False):
|
|
super().__init__()
|
|
self.transformer = DummyTransformer(num_layers, d_model, bias=bias, stacked_params=stacked_params)
|
|
|
|
|
|
# ======================================================
|
|
# helper functions for test
|
|
# ======================================================
|
|
def _filter_parameters(named_params, mode: str = "all"):
|
|
if mode == "all":
|
|
return named_params
|
|
|
|
filtered = []
|
|
for name, param in named_params:
|
|
if name.endswith(mode):
|
|
filtered.append((name, param))
|
|
|
|
return filtered
|
|
|
|
|
|
def assert_params_equal(module, original_params, mode: str = "all"):
|
|
for name, param in _filter_parameters(module.named_parameters(), mode):
|
|
assert_close(param, original_params[name])
|
|
|
|
|
|
def assert_params_not_equal(module, original_params, mode: str = "all"):
|
|
for name, param in _filter_parameters(module.named_parameters(), mode):
|
|
assert not torch.allclose(param, original_params[name])
|
|
|
|
|
|
def make_dummy_lora_state_dict(
|
|
base_key,
|
|
weight_a: torch.Tensor | None = None,
|
|
weight_b: torch.Tensor | None = None,
|
|
weight_bias: torch.Tensor | None = None,
|
|
bias: bool = False,
|
|
lora_a_suffix: str = "lora_A.weight",
|
|
lora_b_suffix: str = "lora_B.weight",
|
|
lora_bias_suffix: str = "bias",
|
|
):
|
|
"""
|
|
Create a dummy LoRA state dict.
|
|
|
|
Returns:
|
|
tuple: (state_dict, expected_delta)
|
|
"""
|
|
key_a = f"{base_key}.{lora_a_suffix}"
|
|
key_b = f"{base_key}.{lora_b_suffix}"
|
|
key_bias = f"{base_key}.{lora_bias_suffix}" if bias else None
|
|
|
|
if weight_a is None:
|
|
weight_a = torch.randn(RANK, HEAD_DIM)
|
|
if weight_b is None:
|
|
weight_b = torch.randn(HEAD_DIM, RANK)
|
|
|
|
state_dict = {}
|
|
expected_delta = None
|
|
|
|
if bias:
|
|
state_dict[key_a] = weight_a
|
|
state_dict[key_b] = weight_b
|
|
weight_bias = torch.randn(HEAD_DIM)
|
|
state_dict[key_bias] = weight_bias
|
|
expected_delta = weight_bias
|
|
else:
|
|
state_dict[key_a] = weight_a
|
|
state_dict[key_b] = weight_b
|
|
expected_delta = weight_b @ weight_a
|
|
|
|
return state_dict, expected_delta
|
|
|
|
|
|
def make_lora_state_dict_for_module(
|
|
module,
|
|
lora_rank=RANK,
|
|
bias=False,
|
|
lora_a_suffix: str = "lora_A.weight",
|
|
lora_b_suffix: str = "lora_B.weight",
|
|
lora_bias_suffix: str = "bias",
|
|
prefix: str = "transformer",
|
|
remap_proj_out: bool = True,
|
|
):
|
|
param_to_weight_names = defaultdict(list)
|
|
if hasattr(module, "stacked_params_mapping"):
|
|
for param_name, weight_name, _ in module.stacked_params_mapping:
|
|
param_to_weight_names[param_name].append(weight_name)
|
|
|
|
state_dict = {}
|
|
for name, param in module.named_parameters(prefix=prefix):
|
|
if name.endswith(".weight"):
|
|
base_key = name[: -len(".weight")]
|
|
if remap_proj_out and base_key.endswith(".to_out"):
|
|
base_key = base_key.replace(".to_out", ".to_out.0")
|
|
|
|
d_in, d_out = param.shape
|
|
is_stacked_param = False
|
|
for param_name, weight_names in param_to_weight_names.items():
|
|
if param_name not in name:
|
|
continue
|
|
is_stacked_param = True
|
|
target_in = d_in // len(weight_names)
|
|
for weight_name in weight_names:
|
|
lora_key = base_key.replace(param_name, weight_name)
|
|
weight_a = torch.randn(lora_rank, d_out)
|
|
weight_b = torch.randn(target_in, lora_rank)
|
|
state_dict[f"{lora_key}.{lora_a_suffix}"] = weight_a
|
|
state_dict[f"{lora_key}.{lora_b_suffix}"] = weight_b
|
|
|
|
if is_stacked_param:
|
|
continue
|
|
|
|
weight_a = torch.randn(lora_rank, d_out)
|
|
weight_b = torch.randn(d_in, lora_rank)
|
|
key_a = f"{base_key}.{lora_a_suffix}"
|
|
key_b = f"{base_key}.{lora_b_suffix}"
|
|
state_dict[key_a] = weight_a
|
|
state_dict[key_b] = weight_b
|
|
elif name.endswith(".bias") and bias:
|
|
base_key = name[: -len(".bias")]
|
|
if remap_proj_out and base_key.endswith(".to_out"):
|
|
base_key = base_key.replace(".to_out", ".to_out.0")
|
|
|
|
d_bias = param.shape[0]
|
|
is_stacked_param = False
|
|
for param_name, weight_names in param_to_weight_names.items():
|
|
if param_name not in name:
|
|
continue
|
|
is_stacked_param = True
|
|
target_bias = d_bias // len(weight_names)
|
|
for weight_name in weight_names:
|
|
lora_key = base_key.replace(param_name, weight_name)
|
|
weight_bias = torch.randn(target_bias)
|
|
state_dict[f"{lora_key}.{lora_bias_suffix}"] = weight_bias
|
|
|
|
if is_stacked_param:
|
|
continue
|
|
|
|
key_bias = f"{base_key}.{lora_bias_suffix}"
|
|
state_dict[key_bias] = torch.randn_like(param)
|
|
|
|
return state_dict
|
|
|
|
|
|
# ======================================================
|
|
# Test loader helper functions
|
|
# ======================================================
|
|
def test_prepare_lora_delta_with_stacked_params():
|
|
base_key = "transformer.attn.to_qkv"
|
|
|
|
sd0, delta0 = make_dummy_lora_state_dict("transformer.attn.to_q")
|
|
sd1, delta1 = make_dummy_lora_state_dict("transformer.attn.to_k")
|
|
sd2, delta2 = make_dummy_lora_state_dict("transformer.attn.to_v")
|
|
|
|
lora_state_dict = {}
|
|
lora_state_dict.update(sd0)
|
|
lora_state_dict.update(sd1)
|
|
lora_state_dict.update(sd2)
|
|
expected_delta = torch.concat([delta0, delta1, delta2])
|
|
|
|
stacked_mapping = {".attn.to_qkv": [".attn.to_q", ".attn.to_k", ".attn.to_v"]}
|
|
|
|
actual_delta, used_keys = _prepare_lora_delta(
|
|
lora_state_dict,
|
|
base_key,
|
|
stacked_mapping,
|
|
)
|
|
|
|
# validate the delta
|
|
assert actual_delta is not None
|
|
assert_close(actual_delta, expected_delta)
|
|
|
|
# validate the used keys
|
|
key_q_a = "transformer.attn.to_q.lora_A.weight"
|
|
key_q_b = "transformer.attn.to_q.lora_B.weight"
|
|
key_k_a = "transformer.attn.to_k.lora_A.weight"
|
|
key_k_b = "transformer.attn.to_k.lora_B.weight"
|
|
key_v_a = "transformer.attn.to_v.lora_A.weight"
|
|
key_v_b = "transformer.attn.to_v.lora_B.weight"
|
|
|
|
for key in [key_q_a, key_q_b, key_k_a, key_k_b, key_v_a, key_v_b]:
|
|
assert key in used_keys
|
|
|
|
|
|
def test_remap_state_dict_keys_matched_two():
|
|
lora_state_dict, _ = make_dummy_lora_state_dict(
|
|
"transformer.attn.to_out.0", bias=True, lora_bias_suffix="lora_B.bias"
|
|
)
|
|
|
|
lora_state_dict = _remap_state_dict_keys(
|
|
lora_state_dict,
|
|
[(".to_out.0.", ".to_out."), (".lora_B.bias", ".bias")],
|
|
)
|
|
|
|
assert len(lora_state_dict) == 3
|
|
assert "transformer.attn.to_out.bias" in lora_state_dict
|
|
|
|
|
|
# ======================================================
|
|
# Test LoraLoaderMixin
|
|
# ======================================================
|
|
class TestLoraLoaderMixin:
|
|
@pytest.mark.parametrize(
|
|
"bias, stacked_params",
|
|
[
|
|
(False, False),
|
|
(True, False),
|
|
(False, True),
|
|
(True, True),
|
|
],
|
|
ids=[
|
|
"basic",
|
|
"bias",
|
|
"stacked_params",
|
|
"bias_stacked_params",
|
|
],
|
|
)
|
|
def test_load_lora_into_module(self, bias: bool, stacked_params: bool):
|
|
pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM, bias=bias, stacked_params=stacked_params)
|
|
original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
|
|
|
|
lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer, bias=bias)
|
|
assert len(lora_state_dict) > 0
|
|
|
|
# remap the keys to vllm-omni format
|
|
lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])
|
|
|
|
used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
|
|
# validate the parameters are updated
|
|
assert_params_not_equal(pipeline.transformer, original_params)
|
|
|
|
# validate all keys are used
|
|
missing_keys = set(lora_state_dict.keys()) - used_keys
|
|
assert len(missing_keys) == 0
|
|
|
|
def test_unload_module_lora(self):
|
|
pipeline = DummyPipeline(NUM_LAYERS, HEAD_DIM)
|
|
original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
|
|
|
|
lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
|
|
# validate the state dict is valid
|
|
assert len(lora_state_dict) > 0
|
|
|
|
# remap the keys to vllm-omni format
|
|
lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])
|
|
|
|
used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
|
|
|
|
# validate the weights are updated
|
|
assert_params_not_equal(pipeline.transformer, original_params)
|
|
|
|
# validate all keys are used
|
|
missing_keys = set(lora_state_dict.keys()) - used_keys
|
|
assert len(missing_keys) == 0
|
|
|
|
unload_used_keys = pipeline.unload_module_lora(lora_state_dict, pipeline.transformer)
|
|
missing_keys = set(lora_state_dict.keys()) - unload_used_keys
|
|
assert len(missing_keys) == 0
|
|
assert len(unload_used_keys - used_keys) == 0
|
|
|
|
# validate the weights are restored
|
|
assert_params_equal(pipeline.transformer, original_params)
|
|
|
|
|
|
# ======================================================
|
|
# Test QwenImageLoraLoaderMixin
|
|
# ======================================================
|
|
class DummyQwenImagePipeline(nn.Module, QwenImageLoraLoaderMixin):
|
|
def __init__(self, num_layers: int, d_model: int, bias: bool = False):
|
|
super().__init__()
|
|
self.transformer = DummyTransformer(num_layers, d_model, bias=bias, stacked_params=True)
|
|
|
|
|
|
class TestQwenImageLoraLoaderMixin:
|
|
def test_load_lora_weights(self, mocker: MockerFixture):
|
|
pipeline = DummyQwenImagePipeline(NUM_LAYERS, HEAD_DIM)
|
|
original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
|
|
|
|
lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
|
|
assert len(lora_state_dict) > 0
|
|
|
|
mocker.patch(
|
|
"vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
|
|
return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
|
|
)
|
|
pipeline.load_lora_weights(lora_state_dict, "adapter0")
|
|
|
|
# validate the weights are updated after lora loaded
|
|
assert_params_not_equal(pipeline.transformer, original_params)
|
|
|
|
# validate lora_loaded map is updated after lora loaded
|
|
assert "adapter0" in pipeline.lora_loaded
|
|
|
|
def test_unload_lora_weights(self, mocker: MockerFixture):
|
|
pipeline = DummyQwenImagePipeline(NUM_LAYERS, HEAD_DIM)
|
|
original_params = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
|
|
|
|
lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
|
|
assert len(lora_state_dict) > 0
|
|
|
|
mocker.patch(
|
|
"vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
|
|
return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
|
|
)
|
|
pipeline.load_lora_weights(lora_state_dict, "adapter0")
|
|
# validate the weights are updated after lora loaded
|
|
assert_params_not_equal(pipeline.transformer, original_params)
|
|
|
|
# validate lora_loaded map is updated after lora loaded
|
|
assert "adapter0" in pipeline.lora_loaded
|
|
|
|
pipeline.unload_lora_weights("adapter0")
|
|
# validate the weights are restored after lora unloaded
|
|
assert_params_equal(pipeline.transformer, original_params)
|
|
|
|
# validate lora_loaded map is updated
|
|
assert "adapter0" not in pipeline.lora_loaded
|
|
|
|
def test_load_diffusers_lora_applies_non_unit_alpha(self, mocker: MockerFixture):
|
|
pipeline = DummyQwenImagePipeline(1, HEAD_DIM)
|
|
to_out = pipeline.transformer.blocks[0].attn.to_out
|
|
original_weight = to_out.weight.detach().clone()
|
|
base_key = "blocks.0.attn.to_out.0"
|
|
lora_a = torch.randn(RANK, HEAD_DIM)
|
|
lora_b = torch.randn(HEAD_DIM, RANK)
|
|
alpha = torch.tensor(6.0)
|
|
state_dict = {
|
|
f"{base_key}.lora_A.weight": lora_a,
|
|
f"{base_key}.lora_B.weight": lora_b,
|
|
f"{base_key}.alpha": alpha,
|
|
}
|
|
mocker.patch(
|
|
"vllm_omni.diffusion.lora.loader.get_converter_by_pipeline",
|
|
return_value=_convert_non_diffusers_qwen_lora_to_diffusers,
|
|
)
|
|
|
|
pipeline.load_lora_weights(state_dict, "adapter0")
|
|
|
|
expected_weight = original_weight + (lora_b @ lora_a) * (alpha.item() / RANK)
|
|
assert_close(to_out.weight, expected_weight)
|
|
|
|
pipeline.unload_lora_weights("adapter0")
|
|
assert_close(to_out.weight, original_weight)
|
|
|
|
|
|
# ======================================================
|
|
# Test WanLoraLoaderMixin
|
|
# ======================================================
|
|
class DummyWanPipeline(nn.Module, WanLoraLoaderMixin):
|
|
def __init__(self, num_layers: int, d_model: int):
|
|
super().__init__()
|
|
self.has_transformer_2 = True
|
|
self.transformer = DummyTransformer(num_layers, d_model, stacked_params=True)
|
|
self.transformer_2 = DummyTransformer(num_layers, d_model, stacked_params=True)
|
|
|
|
|
|
def test_wan_load_lora_weights_targets_both_transformers(mocker: MockerFixture):
|
|
pipeline = DummyWanPipeline(NUM_LAYERS, HEAD_DIM)
|
|
original_high = {name: param.clone() for name, param in pipeline.transformer.named_parameters()}
|
|
original_low = {name: param.clone() for name, param in pipeline.transformer_2.named_parameters()}
|
|
high_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
|
|
low_state_dict = make_lora_state_dict_for_module(pipeline.transformer_2)
|
|
|
|
mocker.patch(
|
|
"vllm_omni.diffusion.lora.loader._load_lora_state_dict",
|
|
side_effect=[high_state_dict, low_state_dict],
|
|
)
|
|
|
|
pipeline.load_lora_weights(["high.safetensors", "low.safetensors"], "adapter0")
|
|
|
|
assert_params_not_equal(pipeline.transformer, original_high)
|
|
assert_params_not_equal(pipeline.transformer_2, original_low)
|
|
assert set(pipeline.lora_loaded["adapter0"]) == {"transformer", "transformer_2"}
|
|
|
|
pipeline.unload_lora_weights("adapter0")
|
|
|
|
assert_params_equal(pipeline.transformer, original_high)
|
|
assert_params_equal(pipeline.transformer_2, original_low)
|
|
assert "adapter0" not in pipeline.lora_loaded
|