From 28291df1795129ae03ce86a7b9b0a230166b5ac5 Mon Sep 17 00:00:00 2001 From: Alicia <115451386+congw729@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:30:03 +0800 Subject: [PATCH] [Refactor][Diffusion] Remove model-specific names from LoRA and ModelOpt loader defaults (#5907) Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- tests/diffusion/lora/test_lora_manager.py | 4 +++- tests/diffusion/models/bagel/test_bagel_lora.py | 9 +++++++-- vllm_omni/diffusion/lora/manager.py | 17 ++++++++++------- .../checkpoint_adapters/modelopt.py | 9 ++++++++- .../diffusion/models/bagel/pipeline_bagel.py | 7 ++++++- .../models/z_image/z_image_transformer.py | 9 +++++++++ 6 files changed, 43 insertions(+), 12 deletions(-) diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index 10b023bf0..b2a0505f3 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -729,8 +729,10 @@ def test_lora_manager_discovers_bagel_component(monkeypatch): lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls), ) - # Pipeline with a 'bagel' component (no 'transformer') + # 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() diff --git a/tests/diffusion/models/bagel/test_bagel_lora.py b/tests/diffusion/models/bagel/test_bagel_lora.py index c285758fe..1e73318ce 100644 --- a/tests/diffusion/models/bagel/test_bagel_lora.py +++ b/tests/diffusion/models/bagel/test_bagel_lora.py @@ -122,8 +122,10 @@ class TestStage1DiTLoRA: lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls), ) - # Build pipeline with bagel component + # Build pipeline with bagel component declared via _lora_components, + # matching how BagelPipeline opts in. pipeline = torch.nn.Module() + pipeline._lora_components = ["bagel"] pipeline.bagel = torch.nn.Module() lm = torch.nn.Module() lm.stacked_params_mapping = [ @@ -194,8 +196,11 @@ class TestBagelLoRARoundTrip: monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", DummyBaseLayerWithLoRA) - # Build pipeline with bagel.language_model.foo (simple non-packed layer) + # Build pipeline with bagel.language_model.foo (simple non-packed layer); + # the bagel component is declared via _lora_components, matching how + # BagelPipeline opts in. pipeline = torch.nn.Module() + pipeline._lora_components = ["bagel"] pipeline.bagel = torch.nn.Module() lm = torch.nn.Module() lm.foo = _FakeLinearBase() diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 11d193531..3a8230530 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -394,17 +394,20 @@ class DiffusionLoRAManager: fully_sharded_loras=False, ) - # Default denoising components, declared DiT components, and any a - # pipeline opts into via ``_lora_components``. + # Components scanned for LoRA-capable layers: framework defaults, + # declared DiT components, and any a pipeline opts into via + # ``_lora_components``. The defaults only cover the generic diffusers + # naming convention: ``transformer`` (plus ``transformer_2`` for + # dual-DiT pipelines such as Wan2.2) and ``unet`` for SDXL-style + # pipelines. Model-specific attribute names must be declared by the + # pipeline itself via ``_dit_modules`` or ``_lora_components``. # - # NOTE: SDXL-style pipelines expose the denoiser as ``unet``. - # Without scanning this component, adapters can load/activate while - # effectively applying to zero layers, producing base-identical output. + # NOTE: if the denoiser component is not scanned here, adapters can + # load/activate while effectively applying to zero layers, producing + # base-identical output. default_components = ( "transformer", "transformer_2", - "dit", - "bagel", "unet", ) declared_components = tuple(getattr(self.pipeline, "_dit_modules", ()) or ()) diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py index 68ec98d5a..f3c703b09 100644 --- a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py +++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py @@ -17,10 +17,17 @@ MODEL_OPT_SCALE_SUFFIXES = ( ".weight_scale_2", ".weight_scale_inv", ) +# Fused projections following the generic diffusers Attention naming +# convention. ``to_q``/``to_k``/``to_v`` and the joint-attention +# ``add_q_proj``/``add_k_proj``/``add_v_proj`` shard names are fused into +# ``to_qkv``/``add_kv_proj`` by most vllm-omni DiT implementations, many of +# which do not declare their own ``packed_modules_mapping``; keeping these +# two entries here avoids repeating them in every such model. Model-specific +# fusions (e.g. Z-Image's ``w13``) must be declared on the model class via +# ``packed_modules_mapping``, which is merged on top of these defaults. DEFAULT_PACKED_MODULES_MAPPING = { "to_qkv": ("to_q", "to_k", "to_v"), "add_kv_proj": ("add_q_proj", "add_k_proj", "add_v_proj"), - "w13": ("w1", "w3"), } FP8_DTYPES = tuple( dtype diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index f27601148..d805fd3b8 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ BagelPipeline implementation for vLLM-Omni. """ @@ -316,6 +316,11 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf "bagel.connector", "bagel.vit_pos_embed", ] + # LoRA scanning target. ``_dit_modules`` holds the dotted path + # "language_model.model", which DiffusionLoRAManager cannot resolve as a + # pipeline attribute, so the top-level ``bagel`` module (LLM plus ViT, + # connector, and VAE projections) is declared explicitly here. + _lora_components: ClassVar[list[str]] = ["bagel"] supports_step_execution: ClassVar[bool] = True def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index 1399eb0f0..f25c01fca 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -612,6 +612,15 @@ class ZImageTransformer2DModel(CachedTransformer): Note: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers. """ + # Fused projections and their checkpoint shard names (mirrors + # stacked_params_mapping in load_weights). ``w13`` is the Z-Image FFN + # fusion of ``w1``/``w3``; declared here so quantized-checkpoint adapters + # can resolve shard keys without a framework-level default entry. + packed_modules_mapping = { + "to_qkv": ["to_q", "to_k", "to_v"], + "w13": ["w1", "w3"], + } + _repeated_blocks = ["ZImageTransformerBlock"] _layerwise_offload_blocks_attrs = ["layers"]