From a473ae78c32ec6590e04831c60dad7c2540401a3 Mon Sep 17 00:00:00 2001 From: Yancy <138764723+Asthenia0412@users.noreply.github.com> Date: Thu, 24 Sep 2026 00:06:54 +0800 Subject: [PATCH] [Bugfix][HunyuanImage3] Fix DiT LoRA binding: namespace alias + GQA QKV de-interleave (#6802) Signed-off-by: Asthenia Signed-off-by: Gao Han Co-authored-by: Asthenia Co-authored-by: Gao Han --- .../common/ci_source_file_dependencies.yml | 2 + .buildkite/cuda/test-nightly.yml | 2 +- .../test_wan_vae_fastpath_install.py | 41 ++- .../test_wan_vae_fastpath_kernels.py | 5 + .../distributed/wan_vae_fastpath_helpers.py | 37 +++ tests/diffusion/lora/test_lora_manager.py | 235 ++++++++++++++++++ .../e2e/accuracy/test_hunyuan_image3_lora.py | 110 ++++++++ vllm_omni/diffusion/lora/manager.py | 64 ++++- .../hunyuan_image3/pipeline_hunyuan_image3.py | 32 +++ 9 files changed, 509 insertions(+), 19 deletions(-) create mode 100644 tests/diffusion/distributed/wan_vae_fastpath_helpers.py create mode 100644 tests/e2e/accuracy/test_hunyuan_image3_lora.py diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index e4a54670a..67231eb8a 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -398,6 +398,8 @@ source_file_dependencies: diffusion_hunyuan_image3_accuracy: - *hunyuan_image3 - tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py + - tests/e2e/accuracy/test_hunyuan_image3_lora.py + - vllm_omni/diffusion/lora/ diffusion_hunyuan_image3_perf: - *hunyuan_image3 diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index fb9977d4d..36df33a41 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -263,7 +263,7 @@ steps: commands: - export HUNYUAN_IMAGE3_MODEL="tencent/HunyuanImage-3.0-Instruct" - export HUNYUAN_IMAGE3_DEVICES="0,1,2,3" - - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m "full_model and H100 and B200 and cards_4" --run-level full_model + - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py tests/e2e/accuracy/test_hunyuan_image3_lora.py -m "full_model and H100 and B200 and cards_4" --run-level full_model - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · Single-GPU" source_file_dependencies: diffusion_qwen_image_perf diff --git a/tests/diffusion/distributed/test_wan_vae_fastpath_install.py b/tests/diffusion/distributed/test_wan_vae_fastpath_install.py index 9f9ae1e63..b5e36b62c 100644 --- a/tests/diffusion/distributed/test_wan_vae_fastpath_install.py +++ b/tests/diffusion/distributed/test_wan_vae_fastpath_install.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """CPU tests for the Wan VAE decoder fast path installer and its exact PyTorch fallbacks.""" from __future__ import annotations import importlib +from collections.abc import Callable from functools import wraps from types import SimpleNamespace @@ -13,6 +14,10 @@ import torch from diffusers.models.autoencoders import AutoencoderKLWan from torch import nn +from tests.diffusion.distributed.wan_vae_fastpath_helpers import ( # noqa: F401 + original_wan_rms_norm, + unpatched_wan_rms_norm, +) from vllm_omni.diffusion import registry as registry_module from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan @@ -25,7 +30,12 @@ from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import ( ) from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import forwards as fastpath_forwards -pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cpu, + pytest.mark.diffusion, + pytest.mark.usefixtures("unpatched_wan_rms_norm"), +] TINY_RESIDUAL = dict( base_dim=8, @@ -201,7 +211,7 @@ def test_failed_install_restores_original_state(failure_stage: str, monkeypatch) with monkeypatch.context() as patch: if failure_stage == "binding": conv_type = type(vae.post_quant_conv) - original_setattr = conv_type.__setattr__ + original_setattr: Callable[..., None] = conv_type.__setattr__ def fail_binding(module, name, value): original_setattr(module, name, value) @@ -534,13 +544,30 @@ def test_upsample_forward_only_fuses_nearest_2x() -> None: assert torch.equal(actual, expected) -def test_rms_norm_vae_substitute_is_not_matched() -> None: +@torch.no_grad() +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_rms_norm_vae_substitute_is_not_matched(dtype: torch.dtype, monkeypatch) -> None: + from diffusers.models.autoencoders import autoencoder_kl_wan + from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE assert not fastpath_forwards.is_diffusers_rms_norm(RMSNormVAE(8, images=False)) - _, vae = _build_pair(TINY_RESIDUAL, torch.float32) - norm = vae.decoder.norm_out - assert fastpath_forwards.is_diffusers_rms_norm(norm) + _, vae = _build_pair(TINY_RESIDUAL, dtype) + assert fastpath_forwards.is_diffusers_rms_norm(vae.decoder.norm_out) + + # Production Wan pipelines use this substitute. The installer must leave + # its different epsilon/upcast behavior intact while optimizing the rest. + monkeypatch.setattr(autoencoder_kl_wan, "WanRMS_norm", RMSNormVAE) + reference, candidate = _build_pair(TINY_RESIDUAL, dtype) + report = install_wan_vae_fastpath(candidate) + assert report.installed + assert "WanRMS_norm" not in report.patched + assert "RMSNormVAE" not in report.patched + assert "forward" not in candidate.decoder.norm_out.__dict__ + latents = torch.randn(1, 4, 3, 6, 8).to(dtype) + expected = reference.decode(latents, return_dict=False)[0] + actual = candidate.decode(latents, return_dict=False)[0] + assert torch.equal(actual, expected) @torch.no_grad() diff --git a/tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py b/tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py index 424ac4a6c..1b76fc87b 100644 --- a/tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py +++ b/tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py @@ -18,6 +18,10 @@ from diffusers.models.autoencoders.autoencoder_kl_wan import ( ) from torch import nn +from tests.diffusion.distributed.wan_vae_fastpath_helpers import ( # noqa: F401 + original_wan_rms_norm, + unpatched_wan_rms_norm, +) from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import forwards as fp from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import ( install_wan_vae_fastpath, @@ -28,6 +32,7 @@ from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton_upsample as up pytestmark = [ + pytest.mark.usefixtures("unpatched_wan_rms_norm"), pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion, diff --git a/tests/diffusion/distributed/wan_vae_fastpath_helpers.py b/tests/diffusion/distributed/wan_vae_fastpath_helpers.py new file mode 100644 index 000000000..ba4462274 --- /dev/null +++ b/tests/diffusion/distributed/wan_vae_fastpath_helpers.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Isolate the upstream numerical reference from Wan's global norm patch.""" + +from __future__ import annotations + +import importlib.util + +import pytest +from diffusers.models.autoencoders import autoencoder_kl_wan +from torch import nn + + +@pytest.fixture(scope="session") +def original_wan_rms_norm() -> type[nn.Module]: + # Wan's package import replaces the public class, including aliases held + # by already collected tests. Load the installed upstream source under a + # private name instead of reloading (and changing) the live module/classes. + spec = importlib.util.spec_from_file_location( + "diffusers.models.autoencoders._wan_fastpath_reference", autoencoder_kl_wan.__file__ + ) + assert spec is not None and spec.loader is not None + reference = importlib.util.module_from_spec(spec) + spec.loader.exec_module(reference) + return reference.WanRMS_norm + + +@pytest.fixture +def unpatched_wan_rms_norm(monkeypatch, request, original_wan_rms_norm): + # Import the patching package before temporarily restoring the reference, + # so importing RMSNormVAE in a test cannot change it halfway through. + from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE + + assert original_wan_rms_norm is not RMSNormVAE + monkeypatch.setattr(autoencoder_kl_wan, "WanRMS_norm", original_wan_rms_norm) + if hasattr(request.module, "WanRMS_norm"): + monkeypatch.setattr(request.module, "WanRMS_norm", original_wan_rms_norm) diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index b2a0505f3..f5d069b6a 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -22,6 +22,8 @@ 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 @@ -550,6 +552,146 @@ def test_lora_manager_warns_when_all_adapters_pinned(monkeypatch): 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="applies to no layer"): + 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. @@ -894,3 +1036,96 @@ def test_removing_suspended_adapter_drops_the_upload(monkeypatch): 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 diff --git a/tests/e2e/accuracy/test_hunyuan_image3_lora.py b/tests/e2e/accuracy/test_hunyuan_image3_lora.py new file mode 100644 index 000000000..f880abffe --- /dev/null +++ b/tests/e2e/accuracy/test_hunyuan_image3_lora.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""HI3 DiT generation with an on-disk PEFT adapter (issue #6411). + +Run with --run-level full_model on four GPUs with at least 80 GiB each. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest +import torch +import yaml +from huggingface_hub import hf_hub_download +from safetensors.torch import save_file + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.lora.request import LoRARequest + +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + + +@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=4) +def test_hunyuan_image3_dit_lora_generation(tmp_path: Path): + model = "tencent/HunyuanImage-3.0-Instruct" + config = json.loads(Path(hf_hub_download(model, "config.json")).read_text()) + hidden = config["hidden_size"] + heads = config["num_attention_heads"] + kv_heads = config["num_key_value_heads"] + head_dim = config.get("head_dim", config["attention_head_dim"]) + rows = (heads + 2 * kv_heads) * head_dim + rank = 8 + generator = torch.Generator().manual_seed(6411) + adapter_dir = tmp_path / "adapter" + adapter_dir.mkdir() + tensors = {} + # Exercise both affected projections in several layers, with different + # updates for every interleaved Q/K/V row, rather than a constant delta + # that would be invariant under a row permutation. + for layer in (0, 1): + for projection, out_dim in (("qkv_proj", rows), ("o_proj", hidden)): + prefix = f"base_model.model.model.layers.{layer}.self_attn.{projection}" + tensors[f"{prefix}.lora_A.weight"] = torch.randn(rank, hidden, generator=generator) * 0.02 + tensors[f"{prefix}.lora_B.weight"] = torch.randn(out_dim, rank, generator=generator) * 0.1 + save_file(tensors, str(adapter_dir / "adapter_model.safetensors")) + (adapter_dir / "adapter_config.json").write_text( + json.dumps({"r": rank, "lora_alpha": rank, "target_modules": ["qkv_proj", "o_proj"]}) + ) + tp = 4 + deploy = tmp_path / "deploy.yaml" + deploy.write_text( + yaml.safe_dump( + { + "pipeline": "hunyuan_image3_dit", + "async_chunk": False, + "trust_remote_code": True, + "stages": [ + { + "stage_id": 0, + "devices": ",".join(map(str, range(tp))), + "max_num_seqs": 1, + "enforce_eager": True, + "trust_remote_code": True, + "parallel_config": {"tensor_parallel_size": tp}, + } + ], + } + ) + ) + request = LoRARequest(lora_name="hi3", lora_int_id=6411, lora_path=str(adapter_dir)) + with OmniRunner( + model, + trust_remote_code=True, + deploy_config=str(deploy), + stage_init_timeout=1800, + init_timeout=2400, + ) as runner: + + def generate(label, lora_request=None, scale=1.0): + outputs = runner.omni.generate( + "A red ceramic teapot on a wooden table.", + OmniDiffusionSamplingParams( + height=512, + width=512, + seed=42, + num_inference_steps=2, + guidance_scale=1.0, + lora_request=lora_request, + lora_scale=scale, + ), + ) + image = outputs[0].images[0] + assert image.size == (512, 512) + image.save(tmp_path / f"{label}.png") + return np.asarray(image).copy() + + baseline = generate("baseline") + adapted = generate("adapted", request) + restored = generate("restored") + repeated = generate("repeated", request) + zero_scale = generate("zero_scale", request, scale=0.0) + assert np.abs(adapted.astype(float) - baseline).mean() > 0.1 + np.testing.assert_array_equal(restored, baseline) + np.testing.assert_array_equal(repeated, adapted) + np.testing.assert_array_equal(zero_scale, baseline) diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 3a8230530..9d6b2af92 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -548,7 +548,21 @@ class DiffusionLoRAManager: return lora_weights module_suffix = full_module_name.split(".")[-1] - return lora_model.get_lora(module_suffix) + lora_weights = lora_model.get_lora(module_suffix) + if lora_weights is not None: + return lora_weights + + # Model-scoped namespace aliases. HunyuanImage-3 registers the DiT + # under ``transformer.layers.*`` (``self.transformer`` aliases + # ``self.model``) while PEFT adapters use ``model.layers.*``. + name_aliases = getattr(self.pipeline, "_get_diffusion_lora_name_aliases", None) + if callable(name_aliases): + for alias in name_aliases(full_module_name) or []: + lora_weights = lora_model.get_lora(alias) + if lora_weights is not None: + return lora_weights + + return None def _is_active_at_scale(self, adapter_id: int, scale: float) -> bool: """True if the adapter_id is active and the current scale matches.""" @@ -646,18 +660,46 @@ class DiffusionLoRAManager: lora_layer.reset_lora(0) continue - 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, + # HunyuanImage-3 fused ``qkv_proj`` LoRA-B rows are + # GQA-interleaved in the checkpoint (per-KV-head + # [Q-group, K, V]); de-interleave them to the block layout + # [all Q, all K, all V] the QKV output slices expect. The + # checkpoint sizes exclude replicated KV heads; set_lora() + # selects the appropriate Q shard and shared KV shard. + deinterleave = getattr(self.pipeline, "_deinterleave_fused_qkv_lora_b", None) + if isinstance(getattr(lora_layer, "base_layer", None), QKVParallelLinear) and callable(deinterleave): + deinterleaved = deinterleave(lora_weights.lora_b) + base = lora_layer.base_layer + output_sizes = ( + base.total_num_heads * base.head_size, + base.total_num_kv_heads * base.head_size, + base.total_num_kv_heads * base.v_head_size, ) - lora_layer.reset_lora(0) - continue + if ( + deinterleaved is None + 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, + ) + 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, + ) + lora_layer.reset_lora(0) + continue + b_splits = list(torch.split(lora_weights.lora_b, list(output_slices), dim=0)) - b_splits = list(torch.split(lora_weights.lora_b, list(output_slices), dim=0)) lora_a_list = [lora_weights.lora_a] * n_slices lora_b_list = [b * scale for b in b_splits] lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index bfe44629a..181bc8dd8 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -440,6 +440,38 @@ class HunyuanImage3Pipeline( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler, ) + def _get_diffusion_lora_name_aliases(self, full_module_name: str) -> list[str]: + """Map LoRA wrapper names to the PEFT adapter namespace. + + ``self.transformer`` aliases ``self.model``, so the generic diffusion + LoRA manager registers the DiT layers under ``transformer.layers.*``. + PEFT adapters store the same weights under the ``model.layers.*`` + namespace (matching the HF checkpoint), so expose the alias to let both + names resolve to the same adapter tensors. + """ + if full_module_name.startswith("transformer."): + return [f"model.{full_module_name[len('transformer.') :]}"] + return [] + + def _deinterleave_fused_qkv_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor | None: + """Re-order a fused ``qkv_proj`` LoRA-B from the checkpoint's + GQA-interleaved layout (per-KV-head ``[Q-group, K, V]``) into the block + layout ``[all Q, all K, all V]`` expected by the QKV output slices. + + Reuses the exact same split as the base-weight loader + (``HunyuanImage3Model._split_qkv_weight``) so the LoRA deltas land on + the same output rows as the checkpoint weights. Returns ``None`` if the + HI3 fused-QKV layout cannot be established, in which case the caller + must fail closed rather than apply the rows to the wrong slices. + """ + split_qkv = getattr(self.transformer, "_split_qkv_weight", None) + if not callable(split_qkv): + return None + try: + return split_qkv(lora_b) + except (RuntimeError, ValueError): + return None + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: skip_prefixes = ["lm_head."] if self.hf_config.tie_word_embeddings else [] # List of unexpected keywords in weight names