Skip to content

Fix Diffusers CPU offload lifecycle - #14

Open
iamwavecut wants to merge 1 commit into
Robbyant:masterfrom
iamwavecut:fix/cpu-offload-lifecycle
Open

Fix Diffusers CPU offload lifecycle#14
iamwavecut wants to merge 1 commit into
Robbyant:masterfrom
iamwavecut:fix/cpu-offload-lifecycle

Conversation

@iamwavecut

Copy link
Copy Markdown

Replaces #6 on the current upstream base. The old PR became conflicted after 316da877309523b9a2d335550aedbd0269ef0b6c (feat(inference): improve distributed inference and memory efficiency).

Summary

  • 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;
  • cover the loader, hook lifecycle, compatibility guards, refiner preload, and pipeline release with 11 focused tests.

Upstream conflict resolution

The new upstream commit adds FSDP placement, base/refiner auxiliary-component sharing, VAE tiling, base-pipeline release, and an optional manual VAE move during refiner denoising. It does not add Diffusers/Accelerate model or sequential CPU offload, hook-aware VAE execution, or the Qwen sequential-offload preload.

This replay preserves the new upstream behavior and adapts CPU offload to it:

  • both DiT FSDP and VLM FSDP are rejected early with CPU offload, as are CFG/CP multi-process modes and the native engine;
  • base/refiner text_encoder and VAE sharing remains enabled normally, but is disabled for CPU offload because one module cannot safely belong to two independent Accelerate hook chains;
  • --release_base_before_refiner removes Accelerate hooks before dropping pipeline components, so hook-owned references do not retain the released models;
  • VAE tiling and the default non-offloaded placement path remain unchanged.

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 preload_module_classes path for that visual parent instead of excluding the full text encoder from offload.

Post-rebase validation

The replacement commit is based directly on upstream 316da877309523b9a2d335550aedbd0269ef0b6c. Local checks used Python 3.12.9, Diffusers 0.37.1, Accelerate 1.14.0, Torch 2.13.0, and Transformers 5.14.1:

python -m unittest discover -s tests -v       11 tests, OK
python -m compileall -q lingbot_video scripts tests
python scripts/inference.py --help            exposes all three CPU-offload modes
ruff check --select E9,F63,F7,F82 ...         passed
git diff --check                              passed

No GPU measurements were rerun solely for this conflict resolution. The following GPU evidence belongs to the original implementation commit a999b14d135c23d031ea7fe783225c26be1446c8; the lifecycle logic is retained here, while the new upstream compatibility paths are covered by the tests above.

Existing GPU validation

Validated 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 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.

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