603 lines
23 KiB
Python
603 lines
23 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
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,
|
|
_resolve_lora_compute_device,
|
|
)
|
|
|
|
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 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
|
|
|
|
|
|
def test_prepare_lora_delta_disambiguates_prefix_overlapping_stacked_names():
|
|
"""A fused name can be a prefix of another fused name.
|
|
|
|
SenseNova-U1 carries both `.qkv_proj` (understanding) and `.qkv_proj_mot_gen`
|
|
(generation). Selecting stacked mappings by substring matched both rules for
|
|
the generation weight and concatenated six shards into a delta twice as tall
|
|
as the parameter.
|
|
"""
|
|
param_to_weight_names = {
|
|
".qkv_proj_mot_gen": [".q_proj_mot_gen", ".k_proj_mot_gen", ".v_proj_mot_gen"],
|
|
".qkv_proj": [".q_proj", ".k_proj", ".v_proj"],
|
|
}
|
|
state_dict = {}
|
|
for part in ("q", "k", "v"):
|
|
for tower in ("_mot_gen", ""):
|
|
base = f"layers.0.self_attn.{part}_proj{tower}"
|
|
state_dict[f"{base}.lora_A.weight"] = torch.ones(RANK, HEAD_DIM)
|
|
state_dict[f"{base}.lora_B.weight"] = torch.ones(HEAD_DIM, RANK)
|
|
|
|
for base_key in ("layers.0.self_attn.qkv_proj_mot_gen", "layers.0.self_attn.qkv_proj"):
|
|
delta, _ = _prepare_lora_delta(state_dict, base_key, param_to_weight_names)
|
|
assert delta.shape == (3 * HEAD_DIM, HEAD_DIM), base_key
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"fused, shards",
|
|
[
|
|
# qwen_image_transformer.py stacked_params_mapping
|
|
(".to_qkv", (".to_q", ".to_k", ".to_v")),
|
|
(".add_kv_proj", (".add_q_proj", ".add_k_proj", ".add_v_proj")),
|
|
# wan2_2_transformer.py stacked_params_mapping
|
|
(".attn1.to_qkv", (".attn1.to_q", ".attn1.to_k", ".attn1.to_v")),
|
|
# wan2_2_s2v_transformer.py stacked_params_mapping
|
|
(".self_attn.to_qkv", (".self_attn.q", ".self_attn.k", ".self_attn.v")),
|
|
],
|
|
)
|
|
def test_prepare_lora_delta_unchanged_for_existing_fused_layouts(fused, shards):
|
|
"""None of the fused names shipped today overlap by prefix, so matching on the
|
|
tail must select exactly the shards substring matching selected."""
|
|
param_to_weight_names = {fused: list(shards)}
|
|
base_key = f"blocks.0{fused}"
|
|
state_dict = {}
|
|
for shard in shards:
|
|
state_dict[f"blocks.0{shard}.lora_A.weight"] = torch.ones(RANK, HEAD_DIM)
|
|
state_dict[f"blocks.0{shard}.lora_B.weight"] = torch.ones(HEAD_DIM, RANK)
|
|
|
|
delta, used_keys = _prepare_lora_delta(state_dict, base_key, param_to_weight_names)
|
|
assert delta.shape == (len(shards) * HEAD_DIM, HEAD_DIM)
|
|
assert len(used_keys) == 2 * len(shards)
|
|
|
|
|
|
def test_lora_loader_mixin_lora_is_fused_property(mocker: MockerFixture):
|
|
pipeline = DummyQwenImagePipeline(1, HEAD_DIM)
|
|
assert pipeline.lora_is_fused is False
|
|
|
|
lora_state_dict = make_lora_state_dict_for_module(pipeline.transformer)
|
|
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")
|
|
assert pipeline.lora_is_fused is True
|
|
|
|
# Test explicit setter
|
|
pipeline2 = DummyQwenImagePipeline(1, HEAD_DIM)
|
|
assert pipeline2.lora_is_fused is False
|
|
pipeline2.lora_is_fused = True
|
|
assert pipeline2.lora_is_fused is True
|
|
|
|
|
|
def test_prepare_lora_delta_with_compute_device():
|
|
base_key = "blocks.0.attn.to_q"
|
|
state_dict = {
|
|
f"{base_key}.lora_A.weight": torch.ones(RANK, HEAD_DIM),
|
|
f"{base_key}.lora_B.weight": torch.ones(HEAD_DIM, RANK),
|
|
}
|
|
delta, used_keys = _prepare_lora_delta(state_dict, base_key, compute_device=torch.device("cpu"))
|
|
assert delta is not None
|
|
assert delta.shape == (HEAD_DIM, HEAD_DIM)
|
|
assert len(used_keys) == 2
|
|
|
|
|
|
def test_resolve_lora_compute_device_probes_unusable_accelerator(mocker: MockerFixture):
|
|
"""cuda:N/npu:N may be reported without a physical device; the probe must reject it."""
|
|
mocker.patch(
|
|
"vllm_omni.diffusion.lora.loader.get_local_device",
|
|
return_value=torch.device("cuda:999"),
|
|
)
|
|
assert _resolve_lora_compute_device() is None
|
|
|
|
|
|
def test_load_lora_into_module_cpu_weights_match_cpu_reference(mocker: MockerFixture):
|
|
"""Regression test: params on CPU used to force compute_device=None and skip the
|
|
#7005 accelerator path. With an unusable accelerator, fusion must fall back to
|
|
CPU and still produce base + lora_B @ lora_A."""
|
|
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)
|
|
lora_state_dict = _remap_state_dict_keys(lora_state_dict, [(".to_out.0.", ".to_out.")])
|
|
|
|
mocker.patch(
|
|
"vllm_omni.diffusion.lora.loader.get_local_device",
|
|
return_value=torch.device("cuda:999"),
|
|
)
|
|
|
|
used_keys = pipeline.load_lora_into_module(lora_state_dict, pipeline.transformer)
|
|
assert len(used_keys) > 0
|
|
assert not (set(lora_state_dict.keys()) - used_keys)
|
|
|
|
for name, param in pipeline.transformer.named_parameters():
|
|
if not name.endswith(".weight"):
|
|
continue
|
|
base_key = name[: -len(".weight")]
|
|
key_a = f"{base_key}.lora_A.weight"
|
|
key_b = f"{base_key}.lora_B.weight"
|
|
if key_a in lora_state_dict and key_b in lora_state_dict:
|
|
expected = original_params[name] + (
|
|
lora_state_dict[key_b].to(param.dtype) @ lora_state_dict[key_a].to(param.dtype)
|
|
)
|
|
assert_close(param, expected)
|
|
|
|
|
|
def test_load_lora_into_module_forwards_usable_accel_device(mocker: MockerFixture):
|
|
"""A usable accelerator must be forwarded as compute_device for CPU params.
|
|
A spy forces the actual math back onto CPU so the test runs without a device."""
|
|
fake_dev = torch.device("cuda:0")
|
|
mocker.patch(
|
|
"vllm_omni.diffusion.lora.loader._resolve_lora_compute_device",
|
|
return_value=fake_dev,
|
|
)
|
|
|
|
seen_devices = []
|
|
original = _prepare_lora_delta
|
|
|
|
def spy(state_dict, base_key, *args, **kwargs):
|
|
seen_devices.append(kwargs.get("compute_device"))
|
|
kwargs["compute_device"] = None
|
|
return original(state_dict, base_key, *args, **kwargs)
|
|
|
|
mocker.patch(
|
|
"vllm_omni.diffusion.lora.loader._prepare_lora_delta",
|
|
side_effect=spy,
|
|
)
|
|
|
|
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)
|
|
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)
|
|
|
|
# Every CPU param with a LoRA delta must have received the accelerator device.
|
|
assert seen_devices
|
|
assert all(d == fake_dev for d in seen_devices)
|
|
assert not (set(lora_state_dict.keys()) - used_keys)
|
|
assert_params_not_equal(pipeline.transformer, original_params)
|