diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index dfa550173..8d85eb4f6 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/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 diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index 8786ae964..be72efffa 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/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 diff --git a/tests/diffusion/offloader/test_layerwise_backend.py b/tests/diffusion/offloader/test_layerwise_backend.py index 7df3c1bb1..5fd80e75c 100644 --- a/tests/diffusion/offloader/test_layerwise_backend.py +++ b/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"] diff --git a/tests/e2e/online_serving/test_flux2_expansion.py b/tests/e2e/online_serving/test_flux2_expansion.py index 0e9e8c89a..336bd83a1 100644 --- a/tests/e2e/online_serving/test_flux2_expansion.py +++ b/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, + ), ] diff --git a/tests/e2e/online_serving/test_zimage_expansion.py b/tests/e2e/online_serving/test_zimage_expansion.py index bef12e55d..9f90ec855 100644 --- a/tests/e2e/online_serving/test_zimage_expansion.py +++ b/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")], ), ] diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py index 362fb4446..680b8bfbb 100644 --- a/vllm_omni/diffusion/models/flux/flux_transformer.py +++ b/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: diff --git a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py index 1d375ca8d..9cf2fb756 100644 --- a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py +++ b/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: diff --git a/vllm_omni/diffusion/models/helios/helios_transformer.py b/vllm_omni/diffusion/models/helios/helios_transformer.py index 812da7db1..b3d2621ad 100644 --- a/vllm_omni/diffusion/models/helios/helios_transformer.py +++ b/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"], } diff --git a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py index 263e39e01..6600b17d5 100644 --- a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py +++ b/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"], diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index c21156706..b34f19e95 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/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"], diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index c4e3b40cd..efaab5a8f 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/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"], } 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 fd8b0e490..3ffad221b 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/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: diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index 20af5b5d8..7876b0094 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/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