From 4cfa9bf8435d81cbfca26438708ce05d340e149c Mon Sep 17 00:00:00 2001 From: Lancer <402430575@qq.com> Date: Tue, 24 Mar 2026 14:50:57 +0800 Subject: [PATCH] [Bugfix] fix offload and hsdp incompatibility (#1888) Signed-off-by: Lancer Signed-off-by: gcanlin Co-authored-by: gcanlin --- .../offloader/test_sequential_backend.py | 128 ++++++++++++++++++ .../models/wan2_2/wan2_2_transformer.py | 2 +- vllm_omni/diffusion/offloader/base.py | 5 + .../diffusion/offloader/sequential_backend.py | 54 ++++---- 4 files changed, 165 insertions(+), 24 deletions(-) create mode 100644 tests/diffusion/offloader/test_sequential_backend.py diff --git a/tests/diffusion/offloader/test_sequential_backend.py b/tests/diffusion/offloader/test_sequential_backend.py new file mode 100644 index 000000000..d18637a78 --- /dev/null +++ b/tests/diffusion/offloader/test_sequential_backend.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for SequentialOffloadBackend.""" + +from unittest.mock import patch + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + + +@pytest.fixture +def accelerator_device() -> torch.device: + """Fixture that provides accelerator device or skips test if unavailable.""" + if current_omni_platform.get_device_count() == 0: + pytest.skip("Accelerator required for this test") + return current_omni_platform.get_torch_device(0) + + +def _create_simple_module() -> nn.Module: + class SimpleModule(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(10, 20) + + return SimpleModule() + + +def _track_pin_memory_calls(): + tracker = {"called": False} + original = torch.Tensor.pin_memory + + def mock(self): + tracker["called"] = True + return original(self) + + return tracker, mock + + +class TestMoveParamsPinMemory: + def test_dtensor_skips_pin_memory(self, accelerator_device): + """DTensor should skip pin_memory to avoid RuntimeError.""" + module = _create_simple_module().to(accelerator_device) + tracker, mock_pin = _track_pin_memory_calls() + + original_isinstance = isinstance + + def fake_isinstance(obj, cls): + if cls.__name__ == "DTensor": + return True + return original_isinstance(obj, cls) + + with patch.object(torch.Tensor, "pin_memory", mock_pin): + with patch("builtins.isinstance", fake_isinstance): + hook = SequentialOffloadHook( + offload_targets=[], + device=accelerator_device, + pin_memory=True, + use_hsdp=False, + ) + hook._move_params( + module, + torch.device("cpu"), + non_blocking=False, + pin_memory=True, + ) + assert not tracker["called"], "pin_memory should not be called for DTensor" + + def test_regular_tensor_calls_pin_memory(self, accelerator_device): + """Regular tensor should call pin_memory when moving to CPU.""" + module = _create_simple_module().to(accelerator_device) + tracker, mock_pin = _track_pin_memory_calls() + + with patch.object(torch.Tensor, "pin_memory", mock_pin): + hook = SequentialOffloadHook( + offload_targets=[], + device=accelerator_device, + pin_memory=True, + use_hsdp=False, + ) + hook._move_params( + module, + torch.device("cpu"), + non_blocking=False, + pin_memory=True, + ) + assert tracker["called"], "pin_memory should be called for regular tensors" + + def test_pin_memory_skipped_when_disabled(self, accelerator_device): + """pin_memory should not be called when pin_memory=False.""" + module = _create_simple_module().to(accelerator_device) + tracker, mock_pin = _track_pin_memory_calls() + + with patch.object(torch.Tensor, "pin_memory", mock_pin): + hook = SequentialOffloadHook( + offload_targets=[], + device=accelerator_device, + pin_memory=False, + use_hsdp=False, + ) + hook._move_params( + module, + torch.device("cpu"), + non_blocking=False, + pin_memory=False, + ) + assert not tracker["called"], "pin_memory should not be called when disabled" + + def test_pin_memory_skipped_for_non_cpu_target(self, accelerator_device): + """pin_memory should not be called for non-CPU targets.""" + module = _create_simple_module().to("cpu") + tracker, mock_pin = _track_pin_memory_calls() + + with patch.object(torch.Tensor, "pin_memory", mock_pin): + hook = SequentialOffloadHook( + offload_targets=[], + device=torch.device("cpu"), + pin_memory=True, + use_hsdp=False, + ) + hook._move_params(module, accelerator_device, non_blocking=False, pin_memory=True) + assert not tracker["called"], "pin_memory should not be called for non-CPU target" 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 2ff604159..a4ae3118a 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -221,7 +221,7 @@ class WanRotaryPosEmbed(nn.Module): freqs_cos = torch.cat([freqs_cos_f, freqs_cos_h, freqs_cos_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1) freqs_sin = torch.cat([freqs_sin_f, freqs_sin_h, freqs_sin_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1) - return freqs_cos, freqs_sin + return freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device) class WanImageEmbedding(nn.Module): diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index 6f838d4dd..e10438765 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/vllm_omni/diffusion/offloader/base.py @@ -24,6 +24,7 @@ class OffloadStrategy(Enum): class OffloadConfig: strategy: OffloadStrategy pin_cpu_memory: bool = True + use_hsdp: bool = False @classmethod def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig": @@ -42,6 +43,9 @@ class OffloadConfig: enable_layerwise_offload = getattr(od_config, "enable_layerwise_offload", False) pin_cpu_memory = getattr(od_config, "pin_cpu_memory", True) + parallel_config = getattr(od_config, "parallel_config", None) + use_hsdp = getattr(parallel_config, "use_hsdp", False) if parallel_config else False + # Determine strategy (mutual exclusion, layer-wise takes priority) if enable_layerwise_offload: strategy = OffloadStrategy.LAYER_WISE @@ -58,6 +62,7 @@ class OffloadConfig: return cls( strategy=strategy, pin_cpu_memory=pin_cpu_memory, + use_hsdp=use_hsdp, ) diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index ebf85624f..638d97335 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -3,6 +3,7 @@ import torch from torch import nn +from torch.distributed._tensor import DTensor # type: ignore[attr-defined] from vllm.logger import init_logger from vllm_omni.diffusion.hooks import HookRegistry, ModelHook @@ -29,52 +30,54 @@ class SequentialOffloadHook(ModelHook): offload_targets: list[nn.Module], device: torch.device, pin_memory: bool = True, + use_hsdp: bool = False, ): # Modules to offload to CPU before this module runs self.offload_targets = offload_targets self.device = device self.pin_memory = pin_memory + self.use_hsdp = use_hsdp @staticmethod - def _move_params(module: nn.Module, device: torch.device) -> None: - """Move module parameters and buffers to device. - - This cls method specifically prevents recursion device movement, - E.g., Cache-DiT CachedBlocks has attr `transformer` as a ref to original - transformer blocks, thus `module.to(device)` will fail for recursion calling, - refer to - https://github.com/vipshop/cache-dit/blob/v1.2.3/src/cache_dit/caching/cache_blocks/__init__.py#L83 - """ + def _move_params( + module: nn.Module, + target_device: torch.device, + *, + non_blocking: bool = True, + pin_memory: bool = False, + ) -> None: for p in module.parameters(): - if p.data.device != device: - p.data = p.data.to(device, non_blocking=True) + if p.data.device != target_device: + data = p.data.to(target_device, non_blocking=non_blocking) + if pin_memory and target_device.type == "cpu" and not isinstance(data, DTensor): + data = data.pin_memory() + p.data = data for b in module.buffers(): - if b.device != device: - b.data = b.data.to(device, non_blocking=True) + if b.device != target_device: + data = b.data.to(target_device, non_blocking=non_blocking) + if pin_memory and target_device.type == "cpu" and not isinstance(data, DTensor): + data = data.pin_memory() + b.data = data def _to_cpu(self, module: nn.Module) -> None: - """Move module to CPU.""" try: param = next(module.parameters()) except StopIteration: return - # Skip if already on CPU if param.device.type == "cpu": return - self._move_params(module, torch.device("cpu")) + self._move_params( + module, + torch.device("cpu"), + non_blocking=not self.use_hsdp, + pin_memory=self.pin_memory, + ) current_omni_platform.empty_cache() - if self.pin_memory: - for p in module.parameters(): - if p.data.device.type == "cpu" and not p.data.is_pinned(): - p.data = p.data.pin_memory() - def _to_gpu(self, module: nn.Module) -> None: - """Move module to GPU.""" try: - # Skip if already on target device if next(module.parameters()).device == self.device: return except StopIteration: @@ -107,6 +110,7 @@ def apply_sequential_offload( encoder_modules: list[nn.Module], device: torch.device, pin_memory: bool = True, + use_hsdp: bool = False, ) -> None: """Apply sequential offloading hooks to DiT and encoder modules. @@ -119,6 +123,7 @@ def apply_sequential_offload( encoder_modules: Encoder modules to register hooks on device: Target GPU device for loading pin_memory: Whether to pin CPU memory for faster transfers + use_hsdp: Whether HSDP is enabled (affects non_blocking behavior) Example: >>> apply_sequential_offload( @@ -135,6 +140,7 @@ def apply_sequential_offload( offload_targets=encoder_modules, device=device, pin_memory=pin_memory, + use_hsdp=use_hsdp, ) registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook) logger.debug("Registered offload hook for %s", dit_mod.__class__.__name__) @@ -146,6 +152,7 @@ def apply_sequential_offload( offload_targets=dit_modules, device=device, pin_memory=pin_memory, + use_hsdp=use_hsdp, ) registry.register_hook(SequentialOffloadHook._HOOK_NAME, hook) logger.debug("Registered offload hook for %s", enc.__class__.__name__) @@ -208,6 +215,7 @@ class ModelLevelOffloadBackend(OffloadBackend): encoder_modules=modules.encoders, device=self.device, pin_memory=self.config.pin_cpu_memory, + use_hsdp=self.config.use_hsdp, ) # Track modules for cleanup