[Bugfix] fix offload and hsdp incompatibility (#1888)

Signed-off-by: Lancer <maruixiang6688@gmail.com>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: gcanlin <canlinguosdu@gmail.com>

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"

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):

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,
         )
 
 

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