diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index 1d3f1811a..5db483276 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -17,7 +17,7 @@ Both strategies use pinned memory for faster CPU-GPU transfers. The strategies a Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks: - **When encoders run**: DiT transformer is offloaded to CPU -- **When DiT runs**: Encoders are offloaded to CPU +- **When DiT runs**: Encoders are offloaded to CPU, if more than one dit models, only one loaded on GPU, others get offloaded to CPU. - **VAE**: Stays resident on GPU Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed. diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 06454ad5c..3f55372db 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -141,11 +141,12 @@ def apply_sequential_offload( ... ) >>> # Modules of pipeline now automatically swap between CPU and GPU """ - # Register hooks on DiT modules (offload encoders when DiT runs) - for dit_mod in dit_modules: + # Register hooks on DiT modules (offload encoders AND other DiTs when a DiT runs) + for i, dit_mod in enumerate(dit_modules): + other_dits = [d for j, d in enumerate(dit_modules) if j != i] registry = HookRegistry.get_or_create(dit_mod) hook = SequentialOffloadHook( - offload_targets=encoder_modules, + offload_targets=encoder_modules + other_dits, device=device, pin_memory=pin_memory, use_hsdp=use_hsdp,