[Refactor] Rename SupportsModuleOffload to SupportsComponentDiscovery (#3354)

Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>

.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.
 

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.

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

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.

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"``).

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"]

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

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.
 

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