CogVideoX: remove per-step host syncs from the denoising loop, enable regional compilation - #14589
Draft
adrianrfreedman wants to merge 3 commits into
Draft
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.
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? |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Two fixes for CogVideoX, both found by following the
examples/profilingguide. 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 upalphas_cumprod, which lives on the CPU, so each lookup has to copy the value back off the GPU andthe CPU blocks until it arrives.
use_dynamic_cfgadds a fourth copy witht.item().Over 20 steps that is 79 stalls inside the loop:
scheduling_ddim_cogvideox.py:392,alphas_cumprod[prev_timestep]scheduling_ddim_cogvideox.py:391,alphas_cumprod[timestep]pipeline_cogvideox.py:744,t.item()39 and not 40 because the last step reads
final_alpha_cumprodinstead.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: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.equalto main on the same seed.The one
tolist()costs about 11 us, roughly what a single.item()costs, so it pays for itselfon the first step.
Why this is fixed in the pipelines and not the scheduler
The real cause is that
alphas_cumprodsits on the CPU and gets indexed by whatever timestep it isgiven. 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, andunipc) index bystep_indexinset_timesteps, so they are fine. Any pipeline passing a CUDA timestep to one of those twelve paysthe 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_blocksattribute is empty. Set_repeated_blocksfor the classCogVideoXTransformer3DModel". 35 of the 69 transformer modelsset it and
_no_split_modulesalready namesCogVideoXBlock, 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'sreturn hidden_states, encoder_hidden_states(
cogvideox_transformer_3d.py:154). Compiled regions pass tensors straight to each other, so eachreplay 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:
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: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 ontorch.cuda.set_sync_debug_mode("error")fromcallback_on_step_endonce the one-off setup copiesare 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 withRuntimeError: called a synchronizing CUDA operation.CUDA only, so it is behind
require_torch_gpu.Self-review
I ran the
self-reviewskill over the diff. Nothing blocking. It caught one thing, and left a fewI would rather you decided.
The test used to spy on
pipe.scheduler.stepand record which device the timestep arrived on..ai/references/testing.mdrules that out ("don't monkeypatch a component method ... just tocapture 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_cpuis a list, not a tensor. The name says where the values are read rather than whatthe object is. I will rename it to
timesteps_listif you prefer.t_cpu = timesteps_cpu[i]in theuse_dynamic_cfgbranch is only there to keep the line under119 characters. Inlining it makes
ruff formatspread the expression over eight lines. Either isfine by me.
torch.cuda.set_sync_debug_modeis a PyTorch prototype feature and warns as much, and no othertest 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.
CogVideoXPipeline. The other three have the same change and no test fileof their own for this, so one test seemed enough to pin the pattern.
examples/profiling/README.mdbut not thetarget-pipelines table at the top, since that one records the pipelines you picked to start with.
I also used
dtyperather thantorch_dtypefor it, becausetorch_dtypenow warns and theREADME already documents the field as
dtype. Say if you would rather the new entry matched itsneighbours, or that all of them were updated.
One documentation suggestion.
.ai/references/pipelines.mdhas no gotcha for this, and it is easyto get wrong because the broken version looks correct. Something like "pass
scheduler.step()aCPU 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
self-reviewskill on the diff?Who can review?
@dg845 @sayakpaul