[BugFix] Enable CPU offloading and Cache-DiT together on Diffusion Model (#1723)

Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -27,7 +27,12 @@ def inference(model_name: str, offload: bool = True):
     torch.cuda.reset_peak_memory_stats(device=device_index)
     monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
-    m = Omni(model=model_name, enable_cpu_offload=offload)
+    m = Omni(
+        model=model_name,
+        # TODO: we might want to add overlapped feature e2e tests
+        # cache_backend="cache_dit",
+        enable_cpu_offload=offload,
+    )
     torch.cuda.reset_peak_memory_stats(device=device_index)
     height = 256
     width = 256

tests/e2e/offline_inference/test_diffusion_layerwise_offload.py

@@ -38,6 +38,8 @@ def run_inference(
     m = Omni(
         model=model_name,
         enable_layerwise_offload=layerwise_offload,
+        # TODO: we might want to add overlapped feature e2e tests
+        # cache_backend="cache_dit",
         boundary_ratio=0.875,
         flow_shift=5.0,
     )

vllm_omni/diffusion/hooks/base.py

@@ -107,7 +107,7 @@ class ModelHook:
             The output of the forward pass.
         """
         args, kwargs = self.pre_forward(module, *args, **kwargs)
-        output = module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
+        output = module._omni_original_forward(*args, **kwargs)  # type: ignore[attr-defined]
         return self.post_forward(module, output)
 
     def reset_state(self, module: nn.Module) -> nn.Module:
@@ -131,7 +131,7 @@ class _WrappedForward:
     def __call__(self, *args: Any, **kwargs: Any):
         registry: HookRegistry | None = getattr(self.module, "_hook_registry", None)
         if registry is None or not registry._hooks:
-            return self.module._original_forward(*args, **kwargs)
+            return self.module._omni_original_forward(*args, **kwargs)
         return registry.dispatch(*args, **kwargs)
 
 
@@ -162,8 +162,10 @@ class HookRegistry:
             setattr(module, "_hook_registry", registry)
 
             # Wrap module.forward once so hooks can intercept calls.
-            if not hasattr(module, "_original_forward"):
-                module._original_forward = module.forward  # type: ignore[attr-defined]
+            # NOTE: Use `_omni_original_forward` to avoid collision with cache-dit's
+            # `_original_forward` attribute on the same module.
+            if not hasattr(module, "_omni_original_forward"):
+                module._omni_original_forward = module.forward  # type: ignore[attr-defined]
                 module.forward = _WrappedForward(module)  # type: ignore[assignment]
 
         return registry
@@ -212,7 +214,7 @@ class HookRegistry:
             The output of the forward pass.
         """
         if not self._hooks:
-            return self.module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
+            return self.module._omni_original_forward(*args, **kwargs)  # type: ignore[attr-defined]
 
         # For single hook case, call directly
         if len(self._hooks) == 1:
@@ -228,7 +230,7 @@ class HookRegistry:
             args, kwargs = hook.pre_forward(self.module, *args, **kwargs)
 
         # Call original forward
-        output = self.module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
+        output = self.module._omni_original_forward(*args, **kwargs)  # type: ignore[attr-defined]
 
         # Apply all post_forward hooks in reverse order
         for _, hook in reversed(sorted_hooks):

vllm_omni/diffusion/offloader/layerwise_backend.py

@@ -50,6 +50,9 @@ class LayerwiseOffloadHook(ModelHook):
         # Per-block synchronization primitive: set after H2D copy completes.
         self._prefetch_done: current_omni_platform.Event | None = None
 
+        # Backward link to the hook that is responsible for prefetching *this* block's weights
+        self._prev_hook: LayerwiseOffloadHook | None = None
+
         self.next_block_parameters: dict[str, nn.Parameter] = {}
         self.next_block_buffers: dict[str, torch.Tensor] = {}
         self.dtype_cpu_flattened_weights: dict[torch.dtype, torch.Tensor] = {}
@@ -128,6 +131,14 @@ class LayerwiseOffloadHook(ModelHook):
 
         return dtype_cpu_flattened_weights, dtype_metadata
 
+    @property
+    def is_materialized(self) -> bool:
+        """Check whether this block's parameters hold real data on device."""
+        for param in self.block_parameters.values():
+            return param.data.dim() > 0
+
+        return True
+
     @torch.compiler.disable
     def prefetch_layer(self, non_blocking: bool = True) -> None:
         """Copy layer weights from CPU -> GPU.
@@ -185,6 +196,12 @@ class LayerwiseOffloadHook(ModelHook):
             buf.data = torch.empty((), device=self.device, dtype=buf.dtype)
 
     def pre_forward(self, module: nn.Module, *args: Any, **kwargs: Any) -> tuple[tuple, dict]:
+        # if the previous hook was skipped and the weights are not on device,
+        # (e.g. by cache-dit block caching), ask the previous hook to
+        # synchronously prefetch *this* block's weights before computation
+        if not self.is_materialized and self._prev_hook is not None:
+            self._prev_hook.prefetch_layer(non_blocking=False)
+
         self.prefetch_layer(non_blocking=True)
 
         return args, kwargs
@@ -293,13 +310,23 @@ class LayerWiseOffloadBackend(OffloadBackend):
             # For subsequent requests, the first layer/block will be pre-fetched
             # during the last layer compute of the previous request.
             last_block, first_block = blocks[-1], blocks[0]
-            hook = apply_block_hook(last_block, first_block, self.device, self.copy_stream, self.config.pin_cpu_memory)
-            hook.prefetch_layer(non_blocking=False)
+            last_hook = apply_block_hook(
+                last_block, first_block, self.device, self.copy_stream, self.config.pin_cpu_memory
+            )
+            last_hook.prefetch_layer(non_blocking=False)
 
+            block_hooks: list[LayerwiseOffloadHook] = [last_hook]
             # Register hook for each of blocks
             for i, block in enumerate(blocks[:-1]):
                 next_block = blocks[(i + 1) % num_blocks]
-                apply_block_hook(block, next_block, self.device, self.copy_stream, self.config.pin_cpu_memory)
+                hook = apply_block_hook(block, next_block, self.device, self.copy_stream, self.config.pin_cpu_memory)
+                block_hooks.append(hook)
+
+            # NOTE(yuanheng-zhao): We make each hook gets a backward reference to the hook
+            # that is responsible for prefetching its block's weights. This is specifically a
+            # workaround for that arbitrary blocks are skipped by caching systems (e.g., cache-dit)
+            for i in range(len(block_hooks)):
+                block_hooks[i]._prev_hook = block_hooks[i - 1]
 
             logger.info(f"Layer-wise offloading enabled on {num_blocks} layers (blocks)")
 

vllm_omni/diffusion/offloader/sequential_backend.py

@@ -35,6 +35,23 @@ class SequentialOffloadHook(ModelHook):
         self.device = device
         self.pin_memory = pin_memory
 
+    @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
+        """
+        for p in module.parameters():
+            if p.data.device != device:
+                p.data = p.data.to(device, non_blocking=True)
+        for b in module.buffers():
+            if b.device != device:
+                b.data = b.data.to(device, non_blocking=True)
+
     def _to_cpu(self, module: nn.Module) -> None:
         """Move module to CPU."""
         try:
@@ -42,13 +59,12 @@ class SequentialOffloadHook(ModelHook):
         except StopIteration:
             return
 
-        previous_device = param.device
         # Skip if already on CPU
-        if previous_device.type == "cpu":
+        if param.device.type == "cpu":
             return
 
-        module.to("cpu")
-        torch.cuda.empty_cache()
+        self._move_params(module, torch.device("cpu"))
+        current_omni_platform.empty_cache()
 
         if self.pin_memory:
             for p in module.parameters():
@@ -64,7 +80,7 @@ class SequentialOffloadHook(ModelHook):
         except StopIteration:
             return
 
-        module.to(self.device)
+        self._move_params(module, self.device)
 
     def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple, dict]:
         # Offload target modules to CPU
@@ -186,19 +202,6 @@ class ModelLevelOffloadBackend(OffloadBackend):
             except Exception as exc:
                 logger.debug("Failed to move VAE to GPU: %s", exc)
 
-        # Initial state: keep DiT modules on CPU (encoders typically run first)
-        # TODO: This part seems to be unnecessary, remove it after testing
-        # for dit_mod in modules.dits:
-        #     dit_mod.to("cpu")
-
-        # torch.cuda.empty_cache()
-
-        # if self.config.pin_cpu_memory:
-        #     for dit_mod in modules.dits:
-        #         for p in dit_mod.parameters():
-        #             if p.data.device.type == "cpu" and not p.data.is_pinned():
-        #                 p.data = p.data.pin_memory()
-
         # Apply sequential offloading hooks
         apply_sequential_offload(
             dit_modules=modules.dits,