From e3b0afbb97a026e8fa348c11a06408bab33d55d1 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Wed, 22 Apr 2026 16:23:00 -0400 Subject: [PATCH] [Refactor] Let diffusion pipelines declare offloadable modules via SupportsModuleOffload (#2427) Signed-off-by: Nick Cao Co-authored-by: Claude --- .claude/skills/add-diffusion-model/SKILL.md | 40 ++- .../diffusion/cpu_offload_diffusion.md | 55 +++- .../offloader/test_module_collector.py | 240 ++++++++++++++++++ .../diffusion/models/bagel/pipeline_bagel.py | 17 +- vllm_omni/diffusion/models/interface.py | 21 ++ .../diffusion/models/ltx2/pipeline_ltx2.py | 9 +- .../models/ltx2/pipeline_ltx2_image2video.py | 9 +- .../models/omnigen2/pipeline_omnigen2.py | 9 +- .../diffusion/offloader/layerwise_backend.py | 13 +- .../diffusion/offloader/module_collector.py | 134 +++++++--- .../diffusion/offloader/sequential_backend.py | 6 +- 11 files changed, 491 insertions(+), 62 deletions(-) create mode 100644 tests/diffusion/offloader/test_module_collector.py diff --git a/.claude/skills/add-diffusion-model/SKILL.md b/.claude/skills/add-diffusion-model/SKILL.md index a7e0bbf9a..0b979e1a9 100644 --- a/.claude/skills/add-diffusion-model/SKILL.md +++ b/.claude/skills/add-diffusion-model/SKILL.md @@ -282,10 +282,11 @@ For Omni or custom models, create: Required updates: 1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table -2. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported -3. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported -4. `examples/offline_inference/xxx/README.md` — offline example docs -5. `examples/online_serve/xxx/README.md` — online serve docs +2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table) +3. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported +4. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported +5. `examples/offline_inference/xxx/README.md` — offline example docs +6. `examples/online_serve/xxx/README.md` — online serve docs ### Step 8: Add E2E Tests (Recommended) @@ -512,6 +513,37 @@ After adding parallelism support, update: 1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table 2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch) +### Step 11: Add CPU Offload Support + +Implement `SupportsModuleOffload` 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 + +class YourPipeline(nn.Module, SupportsModuleOffload): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] # optional +``` + +- `_dit_modules`: denoising submodules (kept on GPU during diffusion loop) +- `_encoder_modules`: encoder/vision submodules (offloaded to CPU during diffusion loop) +- `_vae_modules`: VAE(s) (handled by both sequential and layerwise backends) +- `_resident_modules`: additional modules to pin on GPU during layerwise + offloading (e.g. embedders, connectors). Only used by the layerwise + backend. Optional — defaults to `[]`. + +All attribute names support dotted paths for nested submodules +(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`). + +Pipelines without `SupportsModuleOffload` fall back to scanning +well-known attribute names (`transformer`, `text_encoder`, `vae`, +etc.), which fails for non-standard names. + --- ## Iterative Development Tips diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index f80005ccb..b85c7e408 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -36,6 +36,45 @@ m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True) 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: + +```python +from typing import ClassVar +from vllm_omni.diffusion.models.interface import SupportsModuleOffload + +class MyPipeline(nn.Module, SupportsModuleOffload): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] # optional + + def __init__(self): + super().__init__() + self.transformer = ... # DiT — stays on GPU during denoising + self.text_encoder = ... # Encoder — offloaded to CPU during denoising + self.vision_model = ... # Encoder — offloaded to CPU during denoising + self.vae = ... # VAE — always on GPU +``` + +- `_dit_modules`: attribute names of denoising submodules (kept on GPU + during the diffusion loop). +- `_encoder_modules`: attribute names of encoder/vision submodules + (offloaded to CPU during the diffusion loop). +- `_vae_modules`: attribute names of VAE(s) (always kept on GPU, not + part of the mutual exclusion hooks). +- `_resident_modules`: attribute names of small submodules that must + stay on GPU during layerwise offloading (e.g. embedders, connectors). + Optional — defaults to `[]`. + +All attribute names support dotted paths for nested submodules +(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`). + +Both DiT and encoder lists are needed because the offload hooks use +mutual exclusion: when one group runs, the other moves to CPU. + ### Limitations - Cold start latency increases - Adds overhead from CPU-GPU transfers between encoder and denoising phases @@ -116,11 +155,19 @@ class Flux2Transformer2DModel(nn.Module): **Module Discovery** -The offloader automatically discovers pipeline components: +The offloader discovers pipeline components in two ways: -- **DiT modules**: `transformer`, `transformer_2`, `dit` -- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder` -- **VAE**: `vae` +1. **Protocol-based** (preferred): If the pipeline implements + `SupportsModuleOffload`, 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. + +2. **Fallback attribute scan**: Otherwise, the offloader scans for + well-known attribute names: + - **DiT modules**: `transformer`, `transformer_2`, `dit`, `sr_dit`, `language_model`, `transformer_blocks`, `model` + - **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder` + - **VAE**: `vae`, `audio_vae` **Hook System** diff --git a/tests/diffusion/offloader/test_module_collector.py b/tests/diffusion/offloader/test_module_collector.py new file mode 100644 index 000000000..ab15ad8df --- /dev/null +++ b/tests/diffusion/offloader/test_module_collector.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for ModuleDiscovery and SupportsModuleOffload.""" + +from typing import ClassVar + +import pytest +from torch import nn + +from vllm_omni.diffusion.models.interface import SupportsModuleOffload +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + +# NOTE: tests for skipped/warned attributes verify the *behavioral* +# outcome (attribute excluded from results) but do not assert on log +# output. vllm's logger sets propagate=False, preventing caplog from +# capturing records. See https://github.com/pytest-dev/pytest/issues/3697 + + +# --------------------------------------------------------------------------- +# Test pipelines +# --------------------------------------------------------------------------- + + +class FallbackPipeline(nn.Module): + """Pipeline with standard attribute names (no protocol).""" + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + self.text_encoder = nn.Linear(10, 10) + self.text_encoder_2 = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + +class NonModuleAttrPipeline(nn.Module): + """Pipeline where an attribute is not an nn.Module (fallback path).""" + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + self.text_encoder = "not_a_module" + self.vae = nn.Linear(10, 10) + + +class DuplicateAttrPipeline(nn.Module): + """Pipeline where two encoder attrs point to the same module.""" + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + encoder = nn.Linear(10, 10) + self.text_encoder = encoder + self.text_encoder_2 = encoder + self.vae = nn.Linear(10, 10) + + +class ProtocolPipeline(nn.Module, SupportsModuleOffload): + """Pipeline with non-standard names, using the protocol.""" + + _dit_modules: ClassVar[list[str]] = ["gen_transformer"] + _encoder_modules: ClassVar[list[str]] = ["mllm", "vision_model"] + _vae_modules: ClassVar[list[str]] = ["gen_vae"] + + def __init__(self): + super().__init__() + self.gen_transformer = nn.Linear(10, 10) + self.mllm = nn.Linear(10, 10) + self.vision_model = nn.Linear(10, 10) + self.gen_vae = nn.Linear(10, 10) + # Standard name present but NOT declared — should be ignored + self.transformer = nn.Linear(10, 10) + + +class MissingAttrPipeline(nn.Module, SupportsModuleOffload): + """Pipeline that declares a non-existent attribute.""" + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["nonexistent_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + +class MissingIntermediatePipeline(nn.Module, SupportsModuleOffload): + """Pipeline with dotted path referencing non-existent intermediate.""" + + _dit_modules: ClassVar[list[str]] = ["nonexistent.transformer"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = [] + + def __init__(self): + super().__init__() + + +class NestedPipeline(nn.Module, SupportsModuleOffload): + """Pipeline with nested modules accessed via dotted paths.""" + + _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] + _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__(self): + super().__init__() + self.pipe = nn.Module() + self.pipe.transformer = nn.Linear(10, 10) + self.pipe.text_encoder = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + +class ResidentPipeline(nn.Module, SupportsModuleOffload): + """Pipeline with resident modules that must stay on GPU.""" + + _dit_modules: ClassVar[list[str]] = ["language_model.model"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [ + "bagel.time_embedder", + "bagel.vae2llm", + ] + + def __init__(self): + super().__init__() + self.language_model = nn.Module() + self.language_model.model = nn.Linear(10, 10) + self.bagel = nn.Module() + self.bagel.time_embedder = nn.Linear(10, 10) + self.bagel.vae2llm = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + +class MultiVaePipeline(nn.Module, SupportsModuleOffload): + """Pipeline with multiple VAEs.""" + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"] + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + self.text_encoder = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + self.audio_vae = nn.Linear(10, 10) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestFallbackDiscovery: + """Test the fallback attribute scan (no SupportsModuleOffload).""" + + def test_discovers_standard_attrs(self): + pipeline = FallbackPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert not isinstance(pipeline, SupportsModuleOffload) + assert result.dit_names == ["transformer"] + assert result.dits[0] is pipeline.transformer + assert result.encoder_names == ["text_encoder", "text_encoder_2"] + assert result.vaes[0] is pipeline.vae + assert result.resident_modules == [] + + def test_deduplicates_encoders(self): + pipeline = DuplicateAttrPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert len(result.encoders) == 1 + assert result.encoder_names == ["text_encoder"] + + def test_skips_non_module_attr(self): + pipeline = NonModuleAttrPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert len(result.encoders) == 0 + + +class TestProtocolDiscovery: + """Test discovery via SupportsModuleOffload protocol.""" + + def test_discovers_declared_attrs_and_ignores_undeclared(self): + pipeline = ProtocolPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert isinstance(pipeline, SupportsModuleOffload) + assert result.dit_names == ["gen_transformer"] + assert result.encoder_names == ["mllm", "vision_model"] + assert len(result.vaes) == 1 + # self.transformer exists but is NOT in _dit_modules + assert "transformer" not in result.dit_names + # No _resident_modules declared — defaults to empty + assert result.resident_modules == [] + + def test_skips_missing_attr(self): + pipeline = MissingAttrPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert len(result.encoders) == 0 + + def test_skips_missing_intermediate(self): + result = ModuleDiscovery.discover(MissingIntermediatePipeline()) + + assert len(result.dits) == 0 + + def test_dotted_path_resolves_nested_modules(self): + pipeline = NestedPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert result.dit_names == ["pipe.transformer"] + assert result.dits[0] is pipeline.pipe.transformer + assert result.encoder_names == ["pipe.text_encoder"] + assert result.encoders[0] is pipeline.pipe.text_encoder + assert result.vaes[0] is pipeline.vae + + def test_resident_modules(self): + pipeline = ResidentPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert result.resident_names == [ + "bagel.time_embedder", + "bagel.vae2llm", + ] + assert result.resident_modules[0] is pipeline.bagel.time_embedder + assert result.resident_modules[1] is pipeline.bagel.vae2llm + assert result.dits[0] is pipeline.language_model.model + + def test_multiple_vaes(self): + pipeline = MultiVaePipeline() + result = ModuleDiscovery.discover(pipeline) + + assert len(result.vaes) == 2 + assert result.vaes[0] is pipeline.vae + assert result.vaes[1] is pipeline.audio_vae diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 90baf5f67..0290e9502 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -12,6 +12,7 @@ from collections.abc import Iterable from copy import deepcopy from dataclasses import dataclass from math import isqrt +from typing import ClassVar import numpy as np import torch @@ -26,6 +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.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 @@ -148,12 +150,25 @@ class SiglipNaViTWrapper(nn.Module): return outputs.last_hidden_state.squeeze(0) -class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin): +class BagelPipeline(nn.Module, SupportsModuleOffload, DiffusionPipelineProfilerMixin): """Bagel generation pipeline (MoT) packaged for vllm-omni diffusion engine. This pipeline is self-contained and uses the ported Bagel core files. """ + _dit_modules: ClassVar[list[str]] = ["language_model.model"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [ + "bagel.time_embedder", + "bagel.vae2llm", + "bagel.llm2vae", + "bagel.latent_pos_embed", + "bagel.vit_model", + "bagel.connector", + "bagel.vit_pos_embed", + ] + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__() self.od_config = od_config diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index ef906472b..00d54420d 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -58,6 +58,27 @@ class SupportsStepExecution(Protocol): """Decode output after denoise loop.""" +@runtime_checkable +class SupportsModuleOffload(Protocol): + """Declares which submodules participate in CPU offload. + + All attribute names support dotted paths for nested submodules + (e.g. ``"pipe.transformer"``). + + Attributes: + _dit_modules: Denoising submodules (on GPU during diffusion). + _encoder_modules: Encoder submodules (offloaded during diffusion). + _vae_modules: VAE(s) (always on GPU). + _resident_modules: Extra modules pinned on GPU during layerwise + offloading. Optional, defaults to ``[]``. + """ + + _dit_modules: ClassVar[list[str]] + _encoder_modules: ClassVar[list[str]] + _vae_modules: ClassVar[list[str]] + _resident_modules: ClassVar[list[str]] = [] + + def supports_step_execution(pipeline: object) -> bool: """Return whether `pipeline` implements :class:`SupportsStepExecution`.""" diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index f06ffab16..4f62d72c9 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -9,7 +9,7 @@ import json import os from collections.abc import Iterable from contextlib import nullcontext -from typing import Any +from typing import Any, ClassVar import numpy as np import torch @@ -34,6 +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.progress_bar import ProgressBarMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.lora.request import LoRARequest @@ -1152,9 +1153,13 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): return loader.load_weights(weights) -class LTX2TwoStagesPipeline(nn.Module): +class LTX2TwoStagesPipeline(nn.Module, SupportsModuleOffload): """LTX2TwoStagesPipeline is for two stages image to video generation""" + _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] + _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"] + _vae_modules: ClassVar[list[str]] = ["pipe.vae", "pipe.audio_vae"] + def __init__( self, *, diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py index 50a71a54b..4cc65f749 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py @@ -6,7 +6,7 @@ from __future__ import annotations import copy import os from collections.abc import Iterable -from typing import Any +from typing import Any, ClassVar import numpy as np import PIL.Image @@ -26,6 +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.request import OmniDiffusionRequest from vllm_omni.lora.request import LoRARequest @@ -733,11 +734,15 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline): return DiffusionOutput(output=(video, audio)) -class LTX2ImageToVideoTwoStagesPipeline(nn.Module): +class LTX2ImageToVideoTwoStagesPipeline(nn.Module, SupportsModuleOffload): """LTXImageToVideoTwoStagesPipeline is for two stages image to video generation""" support_image_input = True + _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] + _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"] + _vae_modules: ClassVar[list[str]] = ["pipe.vae", "pipe.audio_vae"] + def __init__( self, *, diff --git a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py index 04720c932..46d634bfd 100644 --- a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py +++ b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py @@ -8,7 +8,7 @@ import os import warnings from collections.abc import Iterable from dataclasses import dataclass -from typing import Any +from typing import Any, ClassVar import numpy as np import PIL.Image @@ -32,6 +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.omnigen2.omnigen2_transformer import ( OmniGen2RotaryPosEmbed, OmniGen2Transformer2DModel, @@ -620,7 +621,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class OmniGen2Pipeline(CFGParallelMixin, nn.Module): +class OmniGen2Pipeline(CFGParallelMixin, nn.Module, SupportsModuleOffload): """ Pipeline for text-to-image generation using OmniGen2. @@ -634,6 +635,10 @@ class OmniGen2Pipeline(CFGParallelMixin, nn.Module): od_config (OmniDiffusionConfig): The OmniDiffusion configuration. """ + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["mllm"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index d1216e2f2..9979d01b1 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -298,13 +298,20 @@ class LayerWiseOffloadBackend(OffloadBackend): for enc in modules.encoders: enc.to(self.device) - # Move VAE to GPU if available - if modules.vae is not None: + # Move VAE(s) to GPU if available + for vae in modules.vaes: try: - modules.vae.to(self.device, non_blocking=True) + vae.to(self.device, non_blocking=True) except Exception as exc: logger.debug("Failed to move VAE to GPU: %s", exc) + # Move resident modules to GPU (small modules needed every forward) + for name, module in zip(modules.resident_names, modules.resident_modules): + try: + module.to(self.device) + except Exception as exc: + logger.debug("Failed to move resident module %s to GPU: %s", name, exc) + logger.info("Applying layer-wise offloading on %s", modules.dit_names) # Apply block-wise offloading hook for each of the blocks in DiT model(s) diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py index a09a33700..dfd81e98b 100644 --- a/vllm_omni/diffusion/offloader/module_collector.py +++ b/vllm_omni/diffusion/offloader/module_collector.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import dataclass +from dataclasses import dataclass, field +from operator import attrgetter from torch import nn from vllm.logger import init_logger +from vllm_omni.diffusion.models.interface import SupportsModuleOffload + logger = init_logger(__name__) @@ -15,15 +18,81 @@ class PipelineModules: dit_names: list[str] encoders: list[nn.Module] encoder_names: list[str] - vae: nn.Module | None = None + vaes: list[nn.Module] + resident_modules: list[nn.Module] = field(default_factory=list) + resident_names: list[str] = field(default_factory=list) class ModuleDiscovery: - """Discovers pipeline components for offloading""" + """Discovers pipeline components for offloading. - DIT_ATTRS = ["transformer", "transformer_2", "dit", "sr_dit", "language_model", "transformer_blocks", "model"] - ENCODER_ATTRS = ["text_encoder", "text_encoder_2", "text_encoder_3", "image_encoder"] - VAE_ATTRS = ["vae", "audio_vae"] + If the pipeline implements :class:`SupportsModuleOffload`, + 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. + _FALLBACK_DIT_ATTRS = [ + "transformer", + "transformer_2", + "dit", + "sr_dit", + "language_model", + "transformer_blocks", + "model", + ] + _FALLBACK_ENCODER_ATTRS = [ + "text_encoder", + "text_encoder_2", + "text_encoder_3", + "image_encoder", + ] + _FALLBACK_VAE_ATTRS = [ + "vae", + "audio_vae", + ] + + @staticmethod + def _collect_modules( + pipeline: nn.Module, + attr_names: list[str], + *, + warn_missing: bool = False, + ) -> tuple[list[nn.Module], list[str]]: + """Resolve attribute names to (module, name) pairs, skipping missing. + + Supports dotted paths via :func:`operator.attrgetter`. + Warns on missing attributes when *warn_missing* is True. + """ + modules: list[nn.Module] = [] + names: list[str] = [] + seen: set[int] = set() + for attr in attr_names: + try: + module = attrgetter(attr)(pipeline) + except AttributeError: + module = None + if module is None: + if warn_missing: + logger.warning( + "Pipeline declares '%s' as offloadable but the attribute does not exist or is None", + attr, + ) + continue + if not isinstance(module, nn.Module): + logger.warning( + "Expected '%s' to be nn.Module, got %r", + attr, + type(module), + ) + continue + if id(module) not in seen: + seen.add(id(module)) + modules.append(module) + names.append(attr) + return modules, names @staticmethod def discover(pipeline: nn.Module) -> PipelineModules: @@ -35,46 +104,29 @@ class ModuleDiscovery: Returns: PipelineModules with lists of discovered modules and names """ - # Collect DiT/transformer modules - dit_modules: list[nn.Module] = [] - dit_names: list[str] = [] - for attr in ModuleDiscovery.DIT_ATTRS: - if not hasattr(pipeline, attr): - continue - module_obj = getattr(pipeline, attr) - if module_obj is None: - continue + declared = isinstance(pipeline, SupportsModuleOffload) + if declared: + dit_attrs = pipeline._dit_modules + enc_attrs = pipeline._encoder_modules + vae_attrs = pipeline._vae_modules + res_attrs = pipeline._resident_modules + else: + dit_attrs = ModuleDiscovery._FALLBACK_DIT_ATTRS + enc_attrs = ModuleDiscovery._FALLBACK_ENCODER_ATTRS + vae_attrs = ModuleDiscovery._FALLBACK_VAE_ATTRS + res_attrs = [] - if not isinstance(module_obj, nn.Module): - logger.warning(f"Expected {attr} to be nn.Module, got {type(module_obj)!r}") - continue - - if module_obj in dit_modules: - continue - - dit_modules.append(module_obj) - dit_names.append(attr) - - # Collect all encoders - encoders: list[nn.Module] = [] - encoder_names: list[str] = [] - for attr in ModuleDiscovery.ENCODER_ATTRS: - if hasattr(pipeline, attr) and getattr(pipeline, attr) is not None: - encoders.append(getattr(pipeline, attr)) - encoder_names.append(attr) - - # Collect VAE - vae = None - for attr in ModuleDiscovery.VAE_ATTRS: - module = getattr(pipeline, attr, None) - if module is not None: - vae = module - break + dit_modules, dit_names = ModuleDiscovery._collect_modules(pipeline, dit_attrs, warn_missing=declared) + encoders, encoder_names = ModuleDiscovery._collect_modules(pipeline, enc_attrs, warn_missing=declared) + vaes, _ = ModuleDiscovery._collect_modules(pipeline, vae_attrs, warn_missing=declared) + residents, resident_names = ModuleDiscovery._collect_modules(pipeline, res_attrs, warn_missing=declared) return PipelineModules( dits=dit_modules, dit_names=dit_names, encoders=encoders, encoder_names=encoder_names, - vae=vae, + vaes=vaes, + resident_modules=residents, + resident_names=resident_names, ) diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 46f48e99c..06454ad5c 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -210,10 +210,10 @@ class ModelLevelOffloadBackend(OffloadBackend): for enc in modules.encoders: enc.to(self.device) - # Move VAE to GPU if available - if modules.vae is not None: + # Move VAE(s) to GPU if available + for vae in modules.vaes: try: - modules.vae.to(self.device, non_blocking=True) + vae.to(self.device, non_blocking=True) except Exception as exc: logger.debug("Failed to move VAE to GPU: %s", exc)