Skip to content

Fix Diffusers CPU offload lifecycle - #6

Closed
iamwavecut wants to merge 1 commit into
Robbyant:masterfrom
iamwavecut:agent/fix-vae-offload-root
Closed

iamwavecut wants to merge 1 commit into
Robbyant:masterfrom
iamwavecut:agent/fix-vae-offload-root

Conversation

@iamwavecut

@iamwavecut iamwavecut commented Jul 9, 2026

Copy link
Copy Markdown

Summary

This keeps the fix limited to the three runtime files involved in Diffusers CPU offload:

  • add --cpu_offload {none,model,sequential} for single-process Diffusers inference;
  • resolve VAE inputs from the VAE hook's execution device, with the real parameter device as the no-hook fallback;
  • reset model-offload hooks around out-of-order image/video VAE encodes while preserving text_encoder->transformer->vae;
  • configure Accelerate sequential offload to preload the Qwen visual subtree that reads a child embedding weight directly.

No documentation, dependency, CI, benchmark, or test files are included.

Root causes

  1. Under Accelerate hooks, VAE parameters can be on CPU or meta; their current parameter device is not the VAE execution device.
  2. Refiner and TI2V perform VAE encode outside the linear model-offload order. Model offload must be reset before and after that phase, including the exception path.
  3. Qwen3VLVisionModel reads pos_embed.weight.device before invoking the child Embedding. Sequential offload therefore needs Accelerate's intended preload_module_classes path for that visual parent instead of excluding the full text encoder from offload.

GPU validation

Validated at commit a999b14 on one RTX PRO 4000 Blackwell 24 GB pod, PyTorch 2.8.0/CUDA 12.8, Diffusers 0.39.0, Accelerate 1.14.0, and robbyant/lingbot-video-dense-1.3b. All runs used the same prompts, seeds, dimensions, steps, dtypes, and Torch SDPA backend; only offload mode changed.

TI2V, 320x192, 49 frames, 40 steps:

Mode Exit Wall time Peak VRAM
none 0 56.50 s 14,330 MiB
model 0 68.00 s 9,890 MiB
sequential 0 96.00 s 1,780 MiB

Decoded RGB output from both offload modes is identical to none: 49 frames, 9,031,680 values compared, maximum absolute difference 0.

Synthetic dense-refiner smoke, 320x192, 49 frames, 4+4 steps:

Mode Exit Wall time Peak VRAM
model 0 80.25 s 9,484 MiB
sequential 0 77.75 s 2,062 MiB

Base and refined decoded RGB outputs are identical between model and sequential offload. They are also identical to the pre-fix de29175 model-offload baseline, showing that the lifecycle fix does not change generated pixels.

The synthetic refiner none control does not fit this 24 GB GPU because it preloads two complete dense pipelines; it reaches 23.41/23.42 GiB before generation. This is a fixture capacity limit, not an offload failure.

Additional checks: focused hook/device tests passed externally (10/10), both supported Diffusers/Accelerate API pairs expose the required preload API, py_compile passes, and git diff --check is clean.

@iamwavecut
iamwavecut force-pushed the agent/fix-vae-offload-root branch from de29175 to a999b14 Compare July 9, 2026 21:20
@iamwavecut iamwavecut changed the title Fix VAE CPU offload device flow Fix Diffusers CPU offload lifecycle Jul 9, 2026
@iamwavecut
iamwavecut marked this pull request as ready for review July 9, 2026 21:22
@iamwavecut

Copy link
Copy Markdown
Author

Superseded by #14, replayed on the current upstream base with the new FSDP/refiner lifecycle integrated and focused regression tests.

@iamwavecut iamwavecut closed this Jul 31, 2026
@iamwavecut
iamwavecut deleted the agent/fix-vae-offload-root branch July 31, 2026 10:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant