Remove per-step host syncs from the CogVideoX denoising loop, enable regional compilation - #14589
adrianrfreedman wants to merge 9 commits into
Conversation
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.
|
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? |
Same pattern as the CogVideoX change: reading a scalar off a device tensor inside the denoising loop forces a device-to-host sync on every step. Z-Image reads a normalised time for the cfg-truncation check. Hoisting the tensor arithmetic and calling .tolist() once keeps the values bit-identical in fp32, fp16 and bf16, which recomputing in Python floats would not. Cosmos reads the current timestep per step, via t.cpu().item() in the 2.5 pipelines and t.item() in Cosmos3 Omni. consisid and Flux2 Klein already avoid this, and scheduler.sigmas is kept on CPU deliberately, so neither is touched. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
ZImagePipeline already hoists the normalised timesteps out of the denoising loop. The other five Z-Image pipelines still read timestep[0].item() every step, which syncs the device to the host. Use the same guard, the same name, and the same expression as the base pipeline rather than a second spelling of it. The precompute only runs when cfg truncation is active, matching ZImagePipeline; the old code paid the sync on every step regardless. Timesteps come from FlowMatchEulerDiscreteScheduler as float32, so .float() is a no-op here and the values are unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
_mask_velocity_predictions guarded each multiply with `if mask.sum() > 0`. That reduction is read on the host, so it syncs the device once per modality per step, and it is inside the denoising loop. The guard is redundant. The mask is 1 - condition_mask over a 0/1 mask, so it is non-negative, and sum() == 0 means every element is zero, in which case pred * mask already equals zeros_like(pred). Both branches return the same tensor, checked in fp32 and bf16 for all-zero, all-one, and mixed masks. The one behavioural difference is a pre-existing NaN or Inf in the predictions: NaN * 0 propagates where zeros_like() used to hide it. Surfacing that seems better than masking it. On Cosmos3-Edge this takes device-to-host copies in the loop from 22 to 10. The loop still syncs inside UniPCMultistepScheduler.step, which is the scheduler-level problem this PR deliberately leaves alone. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The previous commit deleted the `if mask.sum() > 0` guards in _mask_velocity_predictions on the grounds that both branches agree. They do not agree in one case: an all-zero mask with a non-finite prediction, because 0 * inf is NaN and 0 * nan is NaN, where zeros_like() returned zeros. An all-zero mask means the modality is fully conditioned, and zero velocity is the right answer there whatever the model produced, so the guard was a safety net rather than an optimisation. Keep the guards and hoist the read instead. The condition masks are built before the denoising loop and never change, so the pipeline computes the three flags once and passes them in. _mask_velocity_predictions still computes them on demand when they are not supplied, which keeps the modular pipeline callers working unchanged. Same effect on syncs as deleting the guards, since the helper runs twice per step for the conditional and unconditional passes: device-to-host copies in the loop go from 22 to 10 on Cosmos3-Edge. Behaviour is now identical to main in every case. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Ok, sorry for the slow reply. I have answers to both questions, and I have widened the PR. Quality. Output is unchanged. CogVideoX latents are Speedup. No eager wall-clock win on CogVideoX or Z-Image, and I would rather say so than dress it up. Both are GPU-bound at these sizes. What it removes is 6.6 s (CogVideoX) and 1.5 s (Z-Image) of blocked CPU per run, and it lets the loop into a CUDA graph, which it could not enter before because a sync during capture is fatal. Where the queue is deeper the numbers move: 2.2% on CogVideoX under regional compilation, 0.97% on Cosmos3, consistent in both orders. On CogVideoX being unpopular, you were right, so I went looking for how far the pattern spreads. Eight other pipelines read a scalar off a device tensor inside the loop: five Z-Image and three Cosmos. Full numbers are in the description. Three things I would like your call on:
|
|
@sayakpaul I merged main in to clear the conflicts. #14663 added The fix is the one already used elsewhere here, move Would you rather I added it to this PR, or opened a separate one? |
…ync-in-denoising-loop # Conflicts: # src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py
9c9a41d to
14fdaa4
Compare
|
@adrianrfreedman thanks for your hard work! Let's keep the PRs focused and specific to a single model family so that we can hone in, assess the results, take better next steps. Let's keep this PR specific to maybe, CogVideoX? |
Z-Image and Cosmos move to their own PRs so each model family can be assessed on its own. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Ok, done. This PR is now CogVideoX only, and I have pushed the rest as #14861 (Z-Image) and #14862 (Cosmos). The description here has been trimmed to match. One correction to my last comment: I was wrong about #14663. |
|
One at a time please :) These are involved PRs and take efforts to review properly. So, let's please be mindful of that. |
What does this PR do?
Removes the per-step device-to-host syncs from the CogVideoX denoising loop, and enables regional compilation for CogVideoX. Found by following the
examples/profilingguide. Part of #13401, same class of problem as #11696, #13404, #13406, #13461, and #13564.Scoped to CogVideoX only, per @sayakpaul's request. The Z-Image and Cosmos changes that were here before have moved to their own PRs: #14861 and #14862.
Numbers below are
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 passes
scheduler.step()a CUDA timestep. The scheduler indexesalphas_cumprodwith it, which lives on the CPU, so every lookup copies the value back and blocks.use_dynamic_cfgadds a fourth copy viat.item(). Over 20 steps that is 79 stalls:scheduling_ddim_cogvideox.py:392,alphas_cumprod[prev_timestep]scheduling_ddim_cogvideox.py:391,alphas_cumprod[timestep]pipeline_cogvideox.py:744,t.item()39 rather than 40 because the last step reads
final_alpha_cumprod.That is 6623.99 ms of blocked CPU, 331.2 ms a step. With this PR it is 1.03 ms and nothing stalls inside the loop. It also keeps
scheduler.step()out of a CUDA graph, since a sync during capture is fatal:The fix. Read the timesteps into a 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.equalto main on the same seed. The onetolist()costs about 11 us, roughly a single.item().Why in the pipelines and not the scheduler. Twelve schedulers index a CPU
alphas_cumprodinsidestep():consistency_decoder,ddim,ddim_cogvideox,ddim_inverse,ddim_parallel,ddpm,ddpm_parallel,dpm_cogvideox,lcm,repaint,tcd, andunclip. The sigma-based ones index bystep_indexinset_timestepsand are fine. Fixing the scheduler covers all twelve at once, but it moves numerics across a lot of pipelines, so I kept this to pipelines where the output is provably unchanged. Happy to open the scheduler follow-up.2. Regional compilation does not work at all
compile_repeated_blocks()fails with "_repeated_blocksattribute is empty. Set_repeated_blocksfor the classCogVideoXTransformer3DModel". 35 of 69 transformer models set it and_no_split_modulesalready namesCogVideoXBlock, so this looks like an oversight. Declaring it is worth 8.2%, 6379.8 ms eager down to 5859.1 ms.It still will not reach cudagraphs.
mode="reduce-overhead"fails with "accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run", fromCogVideoXBlock.forward'sreturn hidden_states, encoder_hidden_states(cogvideox_transformer_3d.py:154). It fails the same way on main, so not a regression, but it is next in the way.Quality
Unchanged. The final latents are
torch.equalto main on the same seed, so the same kernels run on the same values.Wall clock
4 steps, 3 timed runs after 1 warmup per cell, each comparison in both orders because whichever runs second is consistently a bit slower.
Eager, ms:
No eager win. 0.4%, inside the noise, and the sign flips when I swap the order. CogVideoX-2b is GPU-bound at 230 to 430 ms a step. What the change buys in eager is 6.6 s per run of unblocked CPU, which matters on a busier host, and a loop that can enter a CUDA graph.
Regional compilation (
--compile_regional, default mode), ms:2.2%, and unlike the eager numbers it goes the same way in all three pairs and both orders (210.3, 143.7, and 29.0 ms). That fits the guide's premise that syncs cost more once compilation deepens the queue.
Tests
tests/pipelines/cogvideo/: 146 passed, 45 skipped.Added
TestCogVideoXPipelineHostSync::test_denoising_loop_does_not_sync_with_host. It turns ontorch.cuda.set_sync_debug_mode("error")fromcallback_on_step_endonce the one-off setup copies are done, and off before the decode, so any copy back from the GPU inside the loop raises. On main it fails at thet.item()line withRuntimeError: called a synchronizing CUDA operation. CUDA only, so behindrequire_torch_gpu.Self-review
Ran the
self-reviewskill over the diff. Nothing blocking. It caught one thing: the test used to spy onpipe.scheduler.stepto record which device the timestep arrived on, which.ai/references/testing.mdrules out, and which checked the mechanism rather than the result. The sync-debug test replaces it and uses only public API.Your call on these:
timesteps_cpuis a list, not a tensor. I will rename it totimesteps_listif you prefer.t_cpu = timesteps_cpu[i]in theuse_dynamic_cfgbranch only exists to keep the line under 119 characters. Inlining it makesruff formatspread the expression over eight lines.torch.cuda.set_sync_debug_modeis a PyTorch prototype feature and no other test uses it. Drop the test if you would rather not have it in the suite.CogVideoXPipeline. The other three pipelines take the same change, so one test seemed enough to pin the pattern.examples/profiling/README.mdI added the entry to the config table but not the target-pipelines table, and useddtyperather thantorch_dtypebecausetorch_dtypenow warns. Say if you would rather it matched its neighbours..ai/references/pipelines.mdhas no gotcha for this and the broken version looks correct, so it is easy to get wrong. I will add one here or separately.Before submitting
self-reviewskill on the diff?Who can review?
@dg845 @sayakpaul