Skip to content

CogVideoX: remove per-step host syncs from the denoising loop, enable regional compilation - #14589

Draft
adrianrfreedman wants to merge 3 commits into
huggingface:mainfrom
adrianrfreedman:cogvideox-no-host-sync-in-denoising-loop
Draft

CogVideoX: remove per-step host syncs from the denoising loop, enable regional compilation#14589
adrianrfreedman wants to merge 3 commits into
huggingface:mainfrom
adrianrfreedman:cogvideox-no-host-sync-in-denoising-loop

Conversation

@adrianrfreedman

@adrianrfreedman adrianrfreedman commented Aug 24, 2026

Copy link
Copy Markdown

What does this PR do?

Two fixes for CogVideoX, both found by following the examples/profiling guide. Part of #13401,
and the same kind of problem as #11696, #13404, #13406, #13461, and #13564.

Everything below is measured on THUDM/CogVideoX-2b, fp16, 480x720, 49 frames,
use_dynamic_cfg=True, one L40S.

1. The denoising loop stalls the CPU four times a step

The loop hands scheduler.step() the timestep as a CUDA tensor. The scheduler uses it to look up
alphas_cumprod, which lives on the CPU, so each lookup has to copy the value back off the GPU and
the CPU blocks until it arrives. use_dynamic_cfg adds a fourth copy with t.item().

Over 20 steps that is 79 stalls inside the loop:

count where
39 scheduling_ddim_cogvideox.py:392, alphas_cumprod[prev_timestep]
20 scheduling_ddim_cogvideox.py:391, alphas_cumprod[timestep]
20 pipeline_cogvideox.py:744, t.item()

39 and not 40 because the last step reads final_alpha_cumprod instead.

The CPU spends 6623.99 ms of that run blocked, 331.2 ms per step. With this PR it is 1.03 ms and
nothing stalls inside the loop. The GPU does identical work, the CPU just stops waiting on it.

The stall also keeps scheduler.step() out of a CUDA graph, since a sync during capture is fatal:

baseline (CUDA timestep)   CAPTURE FAILED: AcceleratorError: CUDA error: operation failed during capture
fixed (CPU timestep)       CAPTURE OK

The fix

Read the timesteps into a plain list once before the loop, then pass the Python value to the
scheduler and to the dynamic-CFG term. The transformer still gets the CUDA tensor, so the final
latents are torch.equal to main on the same seed.

The one tolist() costs about 11 us, roughly what a single .item() costs, so it pays for itself
on the first step.

Why this is fixed in the pipelines and not the scheduler

The real cause is that alphas_cumprod sits on the CPU and gets indexed by whatever timestep it is
given. Twelve schedulers do that inside step(): consistency_decoder, ddim, ddim_cogvideox,
ddim_inverse, ddim_parallel, ddpm, ddpm_parallel, dpm_cogvideox, lcm, repaint, tcd,
and unclip. The sigma-based ones (dpmsolver, euler, and unipc) index by step_index in
set_timesteps, so they are fine. Any pipeline passing a CUDA timestep to one of those twelve pays
the same cost.

Fixing it in the scheduler would cover all of them at once, but it moves numerics across a lot of
pipelines. I kept this PR to the four CogVideoX pipelines, where the output is provably unchanged.
I will open a follow-up for the scheduler if you want one.

2. Regional compilation does not work at all

compile_repeated_blocks() fails with "_repeated_blocks attribute is empty. Set
_repeated_blocks for the class CogVideoXTransformer3DModel". 35 of the 69 transformer models
set it and _no_split_modules already names CogVideoXBlock, so this looks like an oversight.
Declaring it is worth 8.2% on its own, 6379.8 ms eager down to 5859.1 ms.

That still does not get CogVideoX to cudagraphs. Under mode="reduce-overhead" it fails with
"accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run", from
CogVideoXBlock.forward's return hidden_states, encoder_hidden_states
(cogvideox_transformer_3d.py:154). Compiled regions pass tensors straight to each other, so each
replay overwrites the last region's output buffer. It fails the same way on main, so it is not a
regression, but it is the next thing in the way.

Wall clock

4 steps, 3 timed runs after 1 warmup per cell. I ran each comparison in both orders because
whichever runs second is consistently a bit slower.

Eager, ms:

order main this PR
main first 6348.9, 6241.0, 6326.1 6381.3, 6284.3, 6391.5
PR first 6507.4, 6358.6, 6496.6 6243.5, 6373.2, 6443.8
mean 6379.8 6352.9

No eager win. 0.4%, inside the noise, and the sign flips when I swap the order. That is expected:
the latents are identical so the same kernels run, and at this size CogVideoX-2b is GPU-bound at
230 to 430 ms a step. Freeing the CPU does not make the GPU faster.

Regional compilation (--compile_regional, default mode), ms:

order main this PR
main first 5859.1 5648.8
PR first 5699.7, 5902.4 5556.0, 5873.4
mean 5820.4 5692.7

2.2% here, and unlike the eager numbers it goes the same way in all three pairs and in both orders
(210.3, 143.7, and 29.0 ms). That fits the guide's premise that syncs cost more once compilation
deepens the GPU queue.

Tests

tests/pipelines/cogvideo/: 134 passed, 10 skipped.

Added TestCogVideoXPipelineHostSync::test_denoising_loop_does_not_sync_with_host. It turns on
torch.cuda.set_sync_debug_mode("error") from callback_on_step_end once the one-off setup copies
are done and turns it off before the decode, so any copy back from the GPU inside the loop raises.
On main it fails at the t.item() line with RuntimeError: called a synchronizing CUDA operation.
CUDA only, so it is behind require_torch_gpu.

Self-review

I ran the self-review skill over the diff. Nothing blocking. It caught one thing, and left a few
I would rather you decided.

The test used to spy on pipe.scheduler.step and record which device the timestep arrived on.
.ai/references/testing.md rules that out ("don't monkeypatch a component method ... just to
capture what the code under test passed to it"), and it was the weaker test anyway because it
checked the mechanism rather than the result. The sync-debug test above replaces it and uses only
public API.

Open questions for you:

  • timesteps_cpu is a list, not a tensor. The name says where the values are read rather than what
    the object is. I will rename it to timesteps_list if you prefer.
  • t_cpu = timesteps_cpu[i] in the use_dynamic_cfg branch is only there to keep the line under
    119 characters. Inlining it makes ruff format spread the expression over eight lines. Either is
    fine by me.
  • torch.cuda.set_sync_debug_mode is a PyTorch prototype feature and warns as much, and no other
    test in the repo uses it. It is the only way I know to test the actual behaviour rather than the
    mechanism behind it. Drop the test if you would rather not have that in the suite.
  • The test only covers CogVideoXPipeline. The other three have the same change and no test file
    of their own for this, so one test seemed enough to pin the pattern.
  • I added the profiling entry to the config table in examples/profiling/README.md but not the
    target-pipelines table at the top, since that one records the pipelines you picked to start with.
    I also used dtype rather than torch_dtype for it, because torch_dtype now warns and the
    README already documents the field as dtype. Say if you would rather the new entry matched its
    neighbours, or that all of them were updated.

One documentation suggestion. .ai/references/pipelines.md has no gotcha for this, and it is easy
to get wrong because the broken version looks correct. Something like "pass scheduler.step() a
CPU timestep, because the tables it looks up live on the CPU and a CUDA one costs a copy back on
every step". I will add it here or separately.

Before submitting

  • Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
    • Did you read the Coding with AI agents guide?
    • Did you run the self-review skill on the diff?
    • Did you share the final self-review notes in the PR description or a comment?
  • Did you read the contributor guideline?
  • Did you read our philosophy doc? (important for complex PRs)
  • Was this discussed/approved via a GitHub issue or the forum? Help us profile important pipelines and improve if needed聽#13401 asks for this. CogVideoX was not on the list and nobody has claimed it in the thread. Opening as a draft, since you preferred profiling results on a PR rather than in the issue.
  • Did you make sure to update the documentation with your changes? No user-facing change.
  • Did you write any new necessary tests?
  • Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

@dg845 @sayakpaul

The loop passed the device timestep tensor to `scheduler.step`, which indexes
`alphas_cumprod` -- a CPU tensor -- with it. Each lookup falls back to a
blocking device-to-host copy, so every step stalls the host and the step cannot
be captured into a CUDA graph. `use_dynamic_cfg` added a fourth copy per step
via `t.item()`.

Read the timesteps into a Python list once before the loop, and pass the host
scalar to the scheduler and to the dynamic-CFG term. The transformer still
receives the device tensor, so outputs are bit-identical.
Without `_repeated_blocks`, `compile_repeated_blocks()` refuses with
"`_repeated_blocks` attribute is empty", so CogVideoX cannot use regional
compilation at all. `_no_split_modules` already names `CogVideoXBlock`.
Lets the guide's tooling reproduce the CogVideoX numbers directly.
@sayakpaul

Copy link
Copy Markdown
Member

Can you comment on the end-to-end speedup and if the quality gets affected because of this? Also, CogVideoX seems like an unpopular model for the time being. So maybe we should focus on another model?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants