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