From 65bc9684659d28dff1010940f0a3a0d6258fd62e Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Fri, 8 May 2026 10:16:49 -0400 Subject: [PATCH] [Refactor] Rename SupportsModuleOffload to SupportsComponentDiscovery (#3354) Signed-off-by: Nick Cao Co-authored-by: Claude --- .claude/skills/add-diffusion-model/SKILL.md | 8 +++---- .../diffusion/cpu_offload_diffusion.md | 11 +++++---- .../offloader/test_module_collector.py | 24 +++++++++---------- .../diffusion/models/bagel/pipeline_bagel.py | 4 ++-- vllm_omni/diffusion/models/interface.py | 8 +++++-- .../diffusion/models/ltx2/pipeline_ltx2.py | 4 ++-- .../models/ltx2/pipeline_ltx2_image2video.py | 4 ++-- .../models/omnigen2/pipeline_omnigen2.py | 4 ++-- .../diffusion/offloader/module_collector.py | 10 ++++---- 9 files changed, 41 insertions(+), 36 deletions(-) diff --git a/.claude/skills/add-diffusion-model/SKILL.md b/.claude/skills/add-diffusion-model/SKILL.md index 212e8b02c..c080ca412 100644 --- a/.claude/skills/add-diffusion-model/SKILL.md +++ b/.claude/skills/add-diffusion-model/SKILL.md @@ -515,15 +515,15 @@ After adding parallelism support, update: ### Step 11: Add CPU Offload Support -Implement `SupportsModuleOffload` on your pipeline class to enable +Implement `SupportsComponentDiscovery` on your pipeline class to enable `--enable-cpu-offload` and `--enable-layerwise-offload`. The protocol declares which submodules the offloader should manage: ```python from typing import ClassVar -from vllm_omni.diffusion.models.interface import SupportsModuleOffload +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery -class YourPipeline(nn.Module, SupportsModuleOffload): +class YourPipeline(nn.Module, SupportsComponentDiscovery): _dit_modules: ClassVar[list[str]] = ["transformer"] _encoder_modules: ClassVar[list[str]] = ["text_encoder"] _vae_modules: ClassVar[list[str]] = ["vae"] @@ -540,7 +540,7 @@ class YourPipeline(nn.Module, SupportsModuleOffload): All attribute names support dotted paths for nested submodules (e.g. `"pipe.transformer"`, `"bagel.time_embedder"`). -Pipelines without `SupportsModuleOffload` fall back to scanning +Pipelines without `SupportsComponentDiscovery` fall back to scanning well-known attribute names (`transformer`, `text_encoder`, `vae`, etc.), which fails for non-standard names. diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index 5db483276..39dc36648 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -38,14 +38,15 @@ vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload ### To Support a Model -Implement the `SupportsModuleOffload` protocol to declare which -submodules participate in offloading: +Implement the `SupportsComponentDiscovery` protocol to declare which +submodules serve as pipeline components (used by offloading, HSDP +sharding, and other framework features): ```python from typing import ClassVar -from vllm_omni.diffusion.models.interface import SupportsModuleOffload +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery -class MyPipeline(nn.Module, SupportsModuleOffload): +class MyPipeline(nn.Module, SupportsComponentDiscovery): _dit_modules: ClassVar[list[str]] = ["transformer"] _encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"] _vae_modules: ClassVar[list[str]] = ["vae"] @@ -158,7 +159,7 @@ class Flux2Transformer2DModel(nn.Module): The offloader discovers pipeline components in two ways: 1. **Protocol-based** (preferred): If the pipeline implements - `SupportsModuleOffload`, its `_dit_modules`, `_encoder_modules`, + `SupportsComponentDiscovery`, its `_dit_modules`, `_encoder_modules`, `_vae_modules`, and `_resident_modules` class variables are used directly. All attribute names support dotted paths (e.g. `"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules. diff --git a/tests/diffusion/offloader/test_module_collector.py b/tests/diffusion/offloader/test_module_collector.py index ab15ad8df..bc7fd97d1 100644 --- a/tests/diffusion/offloader/test_module_collector.py +++ b/tests/diffusion/offloader/test_module_collector.py @@ -1,14 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for ModuleDiscovery and SupportsModuleOffload.""" +"""Unit tests for ModuleDiscovery and SupportsComponentDiscovery.""" from typing import ClassVar import pytest from torch import nn -from vllm_omni.diffusion.models.interface import SupportsModuleOffload +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] @@ -57,7 +57,7 @@ class DuplicateAttrPipeline(nn.Module): self.vae = nn.Linear(10, 10) -class ProtocolPipeline(nn.Module, SupportsModuleOffload): +class ProtocolPipeline(nn.Module, SupportsComponentDiscovery): """Pipeline with non-standard names, using the protocol.""" _dit_modules: ClassVar[list[str]] = ["gen_transformer"] @@ -74,7 +74,7 @@ class ProtocolPipeline(nn.Module, SupportsModuleOffload): self.transformer = nn.Linear(10, 10) -class MissingAttrPipeline(nn.Module, SupportsModuleOffload): +class MissingAttrPipeline(nn.Module, SupportsComponentDiscovery): """Pipeline that declares a non-existent attribute.""" _dit_modules: ClassVar[list[str]] = ["transformer"] @@ -87,7 +87,7 @@ class MissingAttrPipeline(nn.Module, SupportsModuleOffload): self.vae = nn.Linear(10, 10) -class MissingIntermediatePipeline(nn.Module, SupportsModuleOffload): +class MissingIntermediatePipeline(nn.Module, SupportsComponentDiscovery): """Pipeline with dotted path referencing non-existent intermediate.""" _dit_modules: ClassVar[list[str]] = ["nonexistent.transformer"] @@ -98,7 +98,7 @@ class MissingIntermediatePipeline(nn.Module, SupportsModuleOffload): super().__init__() -class NestedPipeline(nn.Module, SupportsModuleOffload): +class NestedPipeline(nn.Module, SupportsComponentDiscovery): """Pipeline with nested modules accessed via dotted paths.""" _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] @@ -113,7 +113,7 @@ class NestedPipeline(nn.Module, SupportsModuleOffload): self.vae = nn.Linear(10, 10) -class ResidentPipeline(nn.Module, SupportsModuleOffload): +class ResidentPipeline(nn.Module, SupportsComponentDiscovery): """Pipeline with resident modules that must stay on GPU.""" _dit_modules: ClassVar[list[str]] = ["language_model.model"] @@ -134,7 +134,7 @@ class ResidentPipeline(nn.Module, SupportsModuleOffload): self.vae = nn.Linear(10, 10) -class MultiVaePipeline(nn.Module, SupportsModuleOffload): +class MultiVaePipeline(nn.Module, SupportsComponentDiscovery): """Pipeline with multiple VAEs.""" _dit_modules: ClassVar[list[str]] = ["transformer"] @@ -155,13 +155,13 @@ class MultiVaePipeline(nn.Module, SupportsModuleOffload): class TestFallbackDiscovery: - """Test the fallback attribute scan (no SupportsModuleOffload).""" + """Test the fallback attribute scan (no SupportsComponentDiscovery).""" def test_discovers_standard_attrs(self): pipeline = FallbackPipeline() result = ModuleDiscovery.discover(pipeline) - assert not isinstance(pipeline, SupportsModuleOffload) + assert not isinstance(pipeline, SupportsComponentDiscovery) assert result.dit_names == ["transformer"] assert result.dits[0] is pipeline.transformer assert result.encoder_names == ["text_encoder", "text_encoder_2"] @@ -183,13 +183,13 @@ class TestFallbackDiscovery: class TestProtocolDiscovery: - """Test discovery via SupportsModuleOffload protocol.""" + """Test discovery via SupportsComponentDiscovery protocol.""" def test_discovers_declared_attrs_and_ignores_undeclared(self): pipeline = ProtocolPipeline() result = ModuleDiscovery.discover(pipeline) - assert isinstance(pipeline, SupportsModuleOffload) + assert isinstance(pipeline, SupportsComponentDiscovery) assert result.dit_names == ["gen_transformer"] assert result.encoder_names == ["mllm", "vision_model"] assert len(result.vaes) == 1 diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index d8be715d7..2363a690a 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -27,7 +27,7 @@ from vllm.transformers_utils.configs.bagel import BagelConfig from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.models.interface import SupportsModuleOffload +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific @@ -150,7 +150,7 @@ class SiglipNaViTWrapper(nn.Module): return outputs.last_hidden_state.squeeze(0) -class BagelPipeline(nn.Module, SupportsModuleOffload, DiffusionPipelineProfilerMixin): +class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin): """Bagel generation pipeline (MoT) packaged for vllm-omni diffusion engine. This pipeline is self-contained and uses the ported Bagel core files. diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index 00d54420d..f0ded9cdb 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -59,8 +59,12 @@ class SupportsStepExecution(Protocol): @runtime_checkable -class SupportsModuleOffload(Protocol): - """Declares which submodules participate in CPU offload. +class SupportsComponentDiscovery(Protocol): + """Declares which submodules serve as pipeline components. + + Used by the framework to locate DiT, encoder, and VAE modules for + CPU offload, HSDP sharding, and other operations that need to know + the pipeline's internal structure. All attribute names support dotted paths for nested submodules (e.g. ``"pipe.transformer"``). diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index 4f62d72c9..a59e9a524 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -34,7 +34,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin -from vllm_omni.diffusion.models.interface import SupportsModuleOffload +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.lora.request import LoRARequest @@ -1153,7 +1153,7 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): return loader.load_weights(weights) -class LTX2TwoStagesPipeline(nn.Module, SupportsModuleOffload): +class LTX2TwoStagesPipeline(nn.Module, SupportsComponentDiscovery): """LTX2TwoStagesPipeline is for two stages image to video generation""" _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py index 4cc65f749..057c144e9 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py @@ -26,7 +26,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin -from vllm_omni.diffusion.models.interface import SupportsModuleOffload +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.lora.request import LoRARequest @@ -734,7 +734,7 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline): return DiffusionOutput(output=(video, audio)) -class LTX2ImageToVideoTwoStagesPipeline(nn.Module, SupportsModuleOffload): +class LTX2ImageToVideoTwoStagesPipeline(nn.Module, SupportsComponentDiscovery): """LTXImageToVideoTwoStagesPipeline is for two stages image to video generation""" support_image_input = True diff --git a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py index 46d634bfd..9c7db02ff 100644 --- a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py +++ b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py @@ -32,7 +32,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.models.interface import SupportsModuleOffload +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.omnigen2.omnigen2_transformer import ( OmniGen2RotaryPosEmbed, OmniGen2Transformer2DModel, @@ -621,7 +621,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class OmniGen2Pipeline(CFGParallelMixin, nn.Module, SupportsModuleOffload): +class OmniGen2Pipeline(CFGParallelMixin, nn.Module, SupportsComponentDiscovery): """ Pipeline for text-to-image generation using OmniGen2. diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py index dfd81e98b..94cbb1fe3 100644 --- a/vllm_omni/diffusion/offloader/module_collector.py +++ b/vllm_omni/diffusion/offloader/module_collector.py @@ -7,7 +7,7 @@ from operator import attrgetter from torch import nn from vllm.logger import init_logger -from vllm_omni.diffusion.models.interface import SupportsModuleOffload +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery logger = init_logger(__name__) @@ -24,16 +24,16 @@ class PipelineModules: class ModuleDiscovery: - """Discovers pipeline components for offloading. + """Discovers pipeline components. - If the pipeline implements :class:`SupportsModuleOffload`, + If the pipeline implements :class:`SupportsComponentDiscovery`, its ``_dit_modules``, ``_encoder_modules``, and ``_vae_modules`` class variables are used directly. Otherwise, falls back to scanning well-known attribute names. """ # Fallback attribute names for pipelines that do not implement - # SupportsModuleOffload. + # SupportsComponentDiscovery. _FALLBACK_DIT_ATTRS = [ "transformer", "transformer_2", @@ -104,7 +104,7 @@ class ModuleDiscovery: Returns: PipelineModules with lists of discovered modules and names """ - declared = isinstance(pipeline, SupportsModuleOffload) + declared = isinstance(pipeline, SupportsComponentDiscovery) if declared: dit_attrs = pipeline._dit_modules enc_attrs = pipeline._encoder_modules