[Feat] support for multi-block layerwise offloading, fix top-level parameters/buffers staying on CPU (#1486)

Signed-off-by: Lancer <maruixiang6688@gmail.com>
Signed-off-by: Lancer <402430575@qq.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>

docs/contributing/model/adding_diffusion_model.md

@@ -802,7 +802,7 @@ omni = Omni(model="your-model", enable_layerwise_offload=True)
 
 ```python
 class WanTransformer3DModel(nn.Module):
-    _layerwise_offload_blocks_attr = "blocks"  # Attribute name containing transformer blocks
+    _layerwise_offload_blocks_attrs = ["blocks"]  # Attribute name containing transformer blocks
 
     def __init__(self):
         self.blocks = nn.ModuleList([...])  # Transformer blocks

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -91,12 +91,19 @@ Models must define the blocks attribute name for layerwise offloading:
 
 ```python
 class WanTransformer3DModel(nn.Module):
-    _layerwise_offload_blocks_attr = "blocks"  # Attribute name containing transformer blocks
+    _layerwise_offload_blocks_attrs = ["blocks"]  # Attribute names containing transformer blocks
 
     def __init__(self):
         self.blocks = nn.ModuleList([...])  # Transformer blocks
 ```
 
+For models with multiple block types:
+
+```python
+class Flux2Transformer2DModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
+```
+
 ### Limitations
 - Cold start latency increases because of
     1) components are loaded to CPU first at the very first during initialization,
@@ -140,4 +147,4 @@ Factory function `get_offload_backend()` selects the appropriate backend based o
 
 **Notes:**
 - Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
-- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attr` pointing to transformer blocks
+- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks

tests/diffusion/offloader/test_layerwise_backend.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Unit tests for LayerwiseOffloadHook."""
+"""Unit tests for LayerwiseOffloadHook and LayerWiseOffloadBackend utilities."""
 
 import gc
 import os
@@ -15,7 +15,7 @@ from torch import nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
 
 import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module
-from vllm_omni.diffusion.offloader.layerwise_backend import LayerwiseOffloadHook
+from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend, LayerwiseOffloadHook
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -127,3 +127,116 @@ class TestLayerwiseOffloadHook:
         assert current_block.weight.to_local().is_meta
         assert current_block.weight.to_local().shape == torch.Size([4])
         assert not hook.is_materialized
+
+
+class _DummyBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(10, 10))
+
+
+class _SingleBlockModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["blocks"]
+
+    def __init__(self, num_blocks: int = 3):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class _MultiBlockModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
+
+    def __init__(self, num_transformer: int = 2, num_single: int = 2):
+        super().__init__()
+        self.transformer_blocks = nn.ModuleList([_DummyBlock() for _ in range(num_transformer)])
+        self.single_transformer_blocks = nn.ModuleList([_DummyBlock() for _ in range(num_single)])
+
+
+class _EmptyBlocksModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["blocks"]
+
+    def __init__(self):
+        super().__init__()
+        self.blocks = nn.ModuleList([])
+
+
+class _InvalidAttrModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["nonexistent_blocks", "blocks"]
+
+    def __init__(self, num_blocks: int = 2):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class _DeprecatedSingleAttrModel(nn.Module):
+    _layerwise_offload_blocks_attr = "blocks"
+
+    def __init__(self, num_blocks: int = 2):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class _NoAttrsModel(nn.Module):
+    def __init__(self, num_blocks: int = 2):
+        super().__init__()
+        self.blocks = nn.ModuleList([_DummyBlock() for _ in range(num_blocks)])
+
+
+class TestGetBlocksFromDit:
+    def test_get_blocks_from_dit_single_block_attr(self):
+        model = _SingleBlockModel(num_blocks=3)
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert attr_names == ["blocks"]
+        assert len(blocks) == 3
+        assert all(isinstance(b, _DummyBlock) for b in blocks)
+
+    def test_get_blocks_from_dit_multi_block_attrs(self):
+        model = _MultiBlockModel(num_transformer=2, num_single=3)
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert set(attr_names) == {"transformer_blocks", "single_transformer_blocks"}
+        assert len(blocks) == 5
+        assert all(isinstance(b, _DummyBlock) for b in blocks)
+
+    def test_get_blocks_from_dit_empty_blocks(self):
+        model = _EmptyBlocksModel()
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert attr_names == []
+        assert blocks == []
+
+    def test_get_blocks_from_dit_invalid_attr_name(self):
+        model = _InvalidAttrModel(num_blocks=2)
+        with pytest.raises(
+            AttributeError,
+            match="Attribute 'nonexistent_blocks' declared in _layerwise_offload_blocks_attrs does not exist",
+        ):
+            LayerWiseOffloadBackend.get_blocks_from_dit(model)
+
+    def test_get_blocks_from_dit_no_attrs_defined(self):
+        model = _NoAttrsModel(num_blocks=3)
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert attr_names == []
+        assert blocks == []
+
+    def test_get_blocks_from_dit_deprecated_single_attr(self):
+        model = _DeprecatedSingleAttrModel(num_blocks=2)
+        attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model)
+        assert attr_names == ["blocks"]
+        assert len(blocks) == 2
+
+
+class TestGetBlocksAttrNames:
+    def test_get_blocks_attr_names_new_format(self):
+        model = _MultiBlockModel()
+        attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model)
+        assert attrs == ["transformer_blocks", "single_transformer_blocks"]
+
+    def test_get_blocks_attr_names_no_attrs(self):
+        model = _NoAttrsModel()
+        attrs = LayerWiseOffloadBackend.get_blocks_attr_names(model)
+        assert attrs == []
+
+    def test_set_blocks_attr_names(self):
+        model = _NoAttrsModel()
+        LayerWiseOffloadBackend.set_blocks_attr_names(model, ["new_blocks"])
+        assert hasattr(model.__class__, "_layerwise_offload_blocks_attrs")
+        assert model.__class__._layerwise_offload_blocks_attrs == ["new_blocks"]

tests/e2e/online_serving/test_flux2_expansion.py

@@ -1,6 +1,12 @@
 """
 Tests for Flux2 Klein; currently Dev is implemented separately,
 but ideally these models will fold together in the future.
+
+Coverage:
+- FP8 + CacheDiT + Ulysses=2 + TP=2
+- Layerwise CPU offload + Ulysses=2 + Ring=2
+- Layerwise CPU offload + TP=2
+- Layerwise CPU offload + HSDP
 """
 
 import pytest
@@ -42,6 +48,45 @@ def _get_diffusion_feature_cases(model: str):
             ),
             marks=FOUR_CARD_FEATURE_MARKS,
         ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-layerwise-offload",
+                    "--ulysses-degree",
+                    "2",
+                    "--ring",
+                    "2",
+                ],
+            ),
+            id="layerwise_ulysses2_ring2",
+            marks=FOUR_CARD_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-layerwise-offload",
+                    "--tensor-parallel-size",
+                    "2",
+                ],
+            ),
+            id="layerwise_tp2",
+            marks=FOUR_CARD_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--enable-layerwise-offload",
+                    "--use-hsdp",
+                    "--hsdp-shard-size",
+                    "2",
+                ],
+            ),
+            id="layerwise_hsdp",
+            marks=FOUR_CARD_FEATURE_MARKS,
+        ),
     ]
 
 

tests/e2e/online_serving/test_zimage_expansion.py

@@ -3,9 +3,12 @@ Tests of common diffusion feature combinations in online serving mode
 for Z-Image.
 
 Coverage is intentionally limited to the minimal 4xL4 cases that
-exercise Z-Image's supported parallel feature combinations:
+exercise Z-Image's supported feature combinations:
 - CacheDiT + FP8 + Ring=2 + TP=2
 - TeaCache + FP8 + Ulysses=2 + Ring=2
+- Layerwise CPU offload + Ulysses=2 + Ring=2
+- Layerwise CPU offload + TP=2
+- Layerwise CPU offload + HSDP
 """
 
 import pytest
@@ -64,12 +67,39 @@ def _get_diffusion_feature_cases():
             OmniServerParams(
                 model=MODEL,
                 server_args=[
+                    "--enable-layerwise-offload",
+                    "--ulysses-degree",
+                    "2",
+                    "--ring",
+                    "2",
+                ],
+            ),
+            id="layerwise_ulysses2_ring2",
+            marks=FOUR_CARD_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    "--enable-layerwise-offload",
+                    "--tensor-parallel-size",
+                    "2",
+                ],
+            ),
+            id="layerwise_tp2",
+            marks=FOUR_CARD_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(
+                model=MODEL,
+                server_args=[
+                    "--enable-layerwise-offload",
                     "--use-hsdp",
                     "--hsdp-shard-size",
                     "2",
                 ],
             ),
-            id="parallel_hsdp",
+            id="layerwise_hsdp",
             marks=[*FOUR_CARD_MARKS, pytest.mark.skip(reason="issue #2435")],
         ),
     ]

vllm_omni/diffusion/models/flux/flux_transformer.py

@@ -510,6 +510,7 @@ class FluxTransformer2DModel(nn.Module):
     # -- typically a transformer layer
     # used for torch compile optimizations
     _repeated_blocks = ["FluxTransformerBlock"]
+    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
     @staticmethod
     def _is_transformer_block(name: str, module) -> bool:

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

@@ -742,6 +742,7 @@ class Flux2Transformer2DModel(nn.Module):
     """
 
     _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
+    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
 
     @staticmethod
     def _is_transformer_block(name: str, module) -> bool:

vllm_omni/diffusion/models/helios/helios_transformer.py

@@ -576,7 +576,7 @@ class HeliosTransformer3DModel(nn.Module):
     """
 
     _repeated_blocks = ["HeliosTransformerBlock"]
-    _layerwise_offload_blocks_attr = "blocks"
+    _layerwise_offload_blocks_attrs = ["blocks"]
     packed_modules_mapping = {
         "to_qkv": ["to_q", "to_k", "to_v"],
     }

vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py

@@ -539,7 +539,7 @@ class HunyuanVideo15Transformer3DModel(nn.Module):
     """
 
     _repeated_blocks = ["HunyuanVideo15TransformerBlock"]
-    _layerwise_offload_blocks_attr = "transformer_blocks"
+    _layerwise_offload_blocks_attrs = ["transformer_blocks"]
     packed_modules_mapping = {
         "to_qkv": ["to_q", "to_k", "to_v"],
         "add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"],

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -882,7 +882,7 @@ class QwenImageTransformer2DModel(CachedTransformer):
     # -- typically a transformer layer
     # used for torch compile optimizations
     _repeated_blocks = ["QwenImageTransformerBlock"]
-    _layerwise_offload_blocks_attr = "transformer_blocks"
+    _layerwise_offload_blocks_attrs = ["transformer_blocks"]
     packed_modules_mapping = {
         "to_qkv": ["to_q", "to_k", "to_v"],
         "add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"],

vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py

@@ -725,7 +725,7 @@ class WanTransformer3DModel(nn.Module):
     """
 
     _repeated_blocks = ["WanTransformerBlock"]
-    _layerwise_offload_blocks_attr = "blocks"
+    _layerwise_offload_blocks_attrs = ["blocks"]
     packed_modules_mapping = {
         "to_qkv": ["to_q", "to_k", "to_v"],
     }

vllm_omni/diffusion/models/z_image/z_image_transformer.py

@@ -579,6 +579,7 @@ class ZImageTransformer2DModel(CachedTransformer):
     """
 
     _repeated_blocks = ["ZImageTransformerBlock"]
+    _layerwise_offload_blocks_attrs = ["layers"]
 
     @staticmethod
     def _is_transformer_block(name: str, module) -> bool:

vllm_omni/diffusion/offloader/layerwise_backend.py

@@ -312,10 +312,9 @@ class LayerWiseOffloadBackend(OffloadBackend):
             dit_name = modules.dit_names[i]
             logger.info(f"Applying hooks on {dit_name} ({dit_module.__class__.__name__})")
 
-            blocks_attr_name = LayerWiseOffloadBackend.get_blocks_attr_name(dit_module)
-            blocks = LayerWiseOffloadBackend.get_blocks_from_dit(dit_module)
+            blocks_attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(dit_module)
 
-            if not blocks_attr_name or not blocks:
+            if not blocks:
                 logger.warning(
                     "Target layers (blocks) not found. Skipping offloading on %s (%s)",
                     dit_name,
@@ -336,11 +335,20 @@ class LayerWiseOffloadBackend(OffloadBackend):
 
             # Move non-block modules to GPU (they stay resident)
             for name, m in dit_module.named_children():
-                if name == blocks_attr_name:
+                if name not in blocks_attr_names:
+                    m.to(self.device)
+                    logger.debug(f"Moved {name} to device {self.device}")
+                else:
                     logger.debug(f"Skipped blocks module {name}")
-                    continue
-                m.to(self.device)
-                logger.debug(f"Moved {name} to device {self.device}")
+
+            # Move top-level params/buffers to GPU (dit_module's own, not sub-modules)
+            for param in dit_module._parameters.values():
+                if param is not None:
+                    param.data = param.data.to(self.device, non_blocking=True)
+
+            for buffer in dit_module._buffers.values():
+                if buffer is not None:
+                    buffer.data = buffer.data.to(self.device, non_blocking=True)
 
             # Pre-fetch the first layer by manually calling the hook function on the last layer;
             # For subsequent requests, the first layer/block will be pre-fetched
@@ -395,40 +403,84 @@ class LayerWiseOffloadBackend(OffloadBackend):
         logger.info("Layer-wise offloading disabled")
 
     @staticmethod
-    def get_blocks_attr_name(model: nn.Module) -> str | None:
-        """Retrieve blocks attribute name from provided DiT model"""
-        return getattr(model.__class__, "_layerwise_offload_blocks_attr", None)
+    def get_blocks_attr_names(model: nn.Module) -> list[str]:
+        """Get block attribute names from model class."""
+        attrs: list[str] = getattr(model.__class__, "_layerwise_offload_blocks_attrs", [])
+
+        if not attrs:
+            old_attr = getattr(model.__class__, "_layerwise_offload_blocks_attr", None)
+            if old_attr is not None:
+                logger.warning(
+                    "'_layerwise_offload_blocks_attr' is deprecated, "
+                    "please use '_layerwise_offload_blocks_attrs' instead. "
+                    "Example: _layerwise_offload_blocks_attrs = ['blocks']"
+                )
+                attrs = [old_attr] if isinstance(old_attr, str) else list(old_attr)
+
+        return attrs
 
     @staticmethod
-    def set_blocks_attr_name(model: nn.Module, name: str) -> None:
-        if not hasattr(model.__class__, "_layerwise_offload_blocks_attr"):
-            setattr(model.__class__, "_layerwise_offload_blocks_attr", name)
+    def set_blocks_attr_names(model: nn.Module, names: list[str]) -> None:
+        if not hasattr(model.__class__, "_layerwise_offload_blocks_attrs"):
+            setattr(model.__class__, "_layerwise_offload_blocks_attrs", names)
 
     @staticmethod
-    def get_blocks_from_dit(model: nn.Module) -> list[nn.Module]:
+    def get_blocks_from_dit(model: nn.Module) -> tuple[list[str], list[nn.Module]]:
         """
-        Retrieve a list of blocks from provided DiT model. Blocks attribute name
-        are found by `_layerwise_offload_blocks_attr` set to DiT models. For example,
+        Retrieve blocks and attribute names from provided DiT model. Blocks attribute names
+        are found by `_layerwise_offload_blocks_attrs` set to DiT models. For example,
 
         ```
         class WanTransformer3DModel(nn.Module):
-            _layerwise_offload_blocks_attr = "blocks"
+            _layerwise_offload_blocks_attrs = ["blocks"]
         ```
+
+        Returns:
+            Tuple of (blocks_attr_names, blocks)
         """
-        blocks_attr_name = LayerWiseOffloadBackend.get_blocks_attr_name(model)
-        if blocks_attr_name is None:
+        blocks_attr_names = LayerWiseOffloadBackend.get_blocks_attr_names(model)
+        if not blocks_attr_names:
             logger.warning(
-                f"No _layerwise_offload_blocks_attr defined for {model.__class__.__name__}, "
+                f"No _layerwise_offload_blocks_attrs defined for {model.__class__.__name__}, "
                 "skipping layerwise offloading"
             )
-            return []
+            return [], []
 
-        _blocks = getattr(model, blocks_attr_name, None)
-        if _blocks is None:
+        blocks = []
+        for name in blocks_attr_names:
+            attr = getattr(model, name, None)
+            if attr is None:
+                raise AttributeError(
+                    f"Attribute '{name}' declared in _layerwise_offload_blocks_attrs "
+                    f"does not exist on model {model.__class__.__name__}"
+                )
+            try:
+                attr_iter = iter(attr)
+            except TypeError:
+                if isinstance(attr, nn.Module):
+                    logger.warning(
+                        "Attribute '%s' on %s is not iterable; treating it as one block.",
+                        name,
+                        model.__class__.__name__,
+                    )
+                    blocks.append(attr)
+                    continue
+
+                logger.warning(
+                    "Attribute '%s' on %s is not iterable (got %s); skipping it.",
+                    name,
+                    model.__class__.__name__,
+                    type(attr).__name__,
+                )
+            else:
+                blocks.extend(attr_iter)
+
+        if not blocks:
             logger.warning(
-                f"Blocks (layers) '{blocks_attr_name}' not found on {model.__class__.__name__}, "
-                "skipping layerwise offloading"
+                "No blocks found in %s for %s, skipping layerwise offloading",
+                blocks_attr_names,
+                model.__class__.__name__,
             )
-            return []
+            return [], []
 
-        return list(_blocks)
+        return blocks_attr_names, blocks