From 04a6e5a992d5dcb0f2bf7e6f66aff6806912f80e Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Thu, 26 Mar 2026 10:08:13 +0800 Subject: [PATCH] [Bugfix] Maintain model-level CPU offload in a blocking way (#1978) Signed-off-by: Yuanheng Zhao --- vllm_omni/diffusion/offloader/sequential_backend.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 638d97335..46f48e99c 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -43,9 +43,17 @@ class SequentialOffloadHook(ModelHook): module: nn.Module, target_device: torch.device, *, - non_blocking: bool = True, + non_blocking: bool = False, pin_memory: bool = False, ) -> 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 != target_device: data = p.data.to(target_device, non_blocking=non_blocking) @@ -83,7 +91,7 @@ class SequentialOffloadHook(ModelHook): except StopIteration: return - self._move_params(module, self.device) + self._move_params(module, self.device, non_blocking=False) def pre_forward(self, module: nn.Module, *args, **kwargs) -> tuple[tuple, dict]: # Offload target modules to CPU