diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py index 0627520c9..6e5f449c1 100644 --- a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py +++ b/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 diff --git a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py b/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py index 998e6232e..beac14099 100644 --- a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py +++ b/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, ) diff --git a/vllm_omni/diffusion/hooks/base.py b/vllm_omni/diffusion/hooks/base.py index 8b330a19f..3c975408c 100644 --- a/vllm_omni/diffusion/hooks/base.py +++ b/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): diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index c1841dadf..5b66ae5ee 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/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)") diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 945cfc413..ebf85624f 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/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,