[perf]: MiniMax H3 on GB10 - load pipeline components on demand instead of all at once - #1761
[perf]: MiniMax H3 on GB10 - load pipeline components on demand instead of all at once#1761KyleNeverGivesUp wants to merge 10 commits into
Conversation
mergify
Bot
added
type: perf
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI
🔴 PR merge requirementsWaiting for
This rule is failing.
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: aa30a9815b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
9fab27e to
3690d74
Compare
29109a7 to
0875b78
Compare
mergify
Bot
added
scope: attention
a3d9725 to
8b2bde2
Compare
8b2bde2 to
cc9bfd0
Compare
aryan5v
left a comment
aryan5v
left a comment
There was a problem hiding this comment.
Reviewed the lazy-load path, H3 opt-in, GPU-direct DiT iterator, and the Spark/offload docs against GB10 FastH3 behavior.
The LazyModule contract is in good shape: is_lazy_module avoids __class__ materialization, self-returning .to() keeps the proxy, compile is a materialize transform, LoRA config reads are deferred, nested-stage walk and failure-path release are covered, and the CPU tests hit the production load_modules / H3 stage constructors. The to_cpu=cpu_offload change in fsdp_load.py is the right GB10 load path.
Please fix the docs before merge. They currently describe always-on sequential H3 start (encode, drop the encoder, then load DiT/VAE, arch-config geometry, one prompt per worker). This PR is an opt-in lazy_module_load flag that defaults to off, reloads from disk on the next generate(), and materializes the video VAE on the first input-prep attribute read. Spark operators following those docs will still OOM.
Two residency issues on the motivating model, neither blocked by tests:
MiniMaxH3InputPreparationStagereadsvae.spatial_compression_ratioandvae.latent_channelsin forward, so the 9.7 GiB video VAE loads before Qwen3-VL. That is the 57.7 GiB encode peak in the PR table. Those values are already onMiniMaxH3VideoVAEArchConfig/MiniMaxH3AudioVAEArchConfigand do not need the weights.MiniMaxH3VideoDecodingStageholdstransformeronly forpatch_size, so the 65 GiB DiT stays resident through VAE decode (the 75.8 GiB FastH3 row). Stashpatch_sizeon the batch during latent prep / denoise and drop the DiT before decode.
Nit: VSA prepare_for_compile graph sharing is a separate change. Fine if it is required for compile-across-reload, otherwise it is easier to review on its own.
| encoder is still resident, the process is a typical `earlyoom` kill (Python is | ||
| preferred). The CUDA pipeline now encodes first, releases the encoder, then | ||
| loads DiT and VAEs onto the accelerator (`to_cpu` follows `cpu_offload`, which | ||
| is off here). See [Offloading](../../inference/offloading.md). |
There was a problem hiding this comment.
This reads as if CUDA H3 now sequential-loads by default. In this PR it does not: lazy_module_load defaults to False, and basic_fasth3.py only turns it on with --lazy-module-load.
On one GB10, FastH3 still loads encoder + DiT + VAEs together unless that flag is set, which is the earlyoom case this bullet describes.
Please state the flag, and either auto-enable it on unified memory or show --lazy-module-load in the Spark command. The GPU-direct to_cpu=cpu_offload sentence can stay.
| (spatial ratio, latent channels, audio sample rate) comes from the VAE arch | ||
| configs until those weights load. A later `generate()` on the same worker | ||
| currently re-enters conditioning after the encoder has been released; start a | ||
| new generator for a new prompt until prompt-cache reload exists. |
There was a problem hiding this comment.
This paragraph does not match the code in this PR.
- Sequential encode-then-release is not default CUDA H3 behavior; it is
--lazy-module-load. - Input-prep geometry is not taken from VAE arch configs here. The first input-prep access of
spatial_compression_ratio/latent_channelsmaterializes the full video VAE. - A later
generate()on a lazy worker re-runs the loader afterrelease(); it does not require a new process. The broken second-prompt behavior and the MLX phase-order sentence belong to a different design.
Please describe lazy_module_load only (the section below already does that) and drop this always-on sequential-start writeup.
| help="load each heavy component on first use and free it after the last stage that " | ||
| "needs it, so peak memory is the largest overlapping set instead of the sum of every " | ||
| "component. Enable when the model does not fit at load time; costs a reload per " | ||
| "generation, so leave it off when it does fit") |
There was a problem hiding this comment.
Default all on one GB10 cannot keep FastH3's four components resident (PR body: 124 GiB vs 121 GiB). Leaving this off means the Spark example still dies at load unless the user already knows the flag.
Either auto-enable on unified memory (same place disable_offload_on_unified_memory already runs) or default this on when num_gpus == 1 and document the reload cost. Off-by-default is reasonable on 4×GB200.
| # Deferral is safe here: no stage reads a component's attributes while it | ||
| # is being constructed, and `initialize_pipeline` only inspects the | ||
| # schedulers, which are never deferred. | ||
| _lazy_module_names = ("text_encoder", "transformer", "vae", "audio_vae") |
There was a problem hiding this comment.
Opt-in list is right, but last-holder scheduling still overlaps more than H3 needs:
- Input prep holds
vaeand reads two ints in forward, so the video VAE loads before the text encoder (your r16 table: 9.7 then 57.7 GiB peak). - Video decode holds
transformeronly forpatch_size, so the DiT stays through VAE decode (FastH3 75.8 GiB resident).
If those scalars come from arch config / the batch instead of the live modules, encode peak drops the VAE and decode can free the DiT. Construction-time access is already tested; first-forward attribute reads are not.
|
This PR has merge conflicts with the base branch. Please rebase: git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease |
|
Addressed the CHANGES_REQUESTED review on a branch you can merge into this PR: KyleNeverGivesUp#2 ( What changed:
Spark2 GB10 check, alpine 768×1344×124, VSA-DataFree, seed 2026,
E2E 399.3 s. Denoise matches the sequential-start full-VAE run (~138 vs 139 s). The extra wall time vs 336.1 s sequential is cold Qwen load inside CUDA TAEH3 is not in this follow-up; that is a separate PR against |
|
#1803 supersedes this PR for the remaining FastH3-on-Spark work. Already on mainThe GB10 “do not load encoder + DiT + VAE all at once” path landed in #1793 ( Preserved from this PRWhat survived is the behavior, not this branch’s
Still only on this branch (not on main, not in the Spark-pair commit): generic What #1803 adds on top#1803 is the follow-on, stacked on main + #1793, not on this lazy-load implementation:
Please review / land Spark-pair on #1803. This PR should not merge as-is on top of #1793. |
…last stage A pipeline materializes every component before the first stage runs, so peak memory is the sum of all components even though no two are needed at the same moment. The CPU offload flags cannot help with this: they act after loading, and on a unified-memory device moving weights to the host frees nothing because it is the same pool. Add lazy_module_load, off by default. Heavy components become a LazyModule proxy that loads on first use, and the pipeline installs a release hook on the last stage that holds each one. Peak becomes the largest overlapping set instead of the sum. Measured on MiniMax-H3 r16, 121 GiB GB10, 1 GPU, 192x320, 4 steps. Peak CUDA allocated is 57.7 GiB, reached during conditioning where the text encoder and video VAE overlap. The four deferred components account for 96.3 GiB together, which is what stays resident without the flag: text encoder 48.0, DiT 37.8, video VAE 9.7, audio VAE 0.8. Generation completed and wrote a video. Details worth flagging: The release hook lives on PipelineStage.__call__, not in the pipeline stage loop, so pipelines that override forward still free. The proxy forwards __class__, so isinstance stays honest. A proxy answering False to isinstance(module, FSDPModule) would take the wrong branch silently. Self-returning methods hand back the proxy rather than the component. Stages write self.vae = self.vae.to(device) in a dozen places; returning the component there would replace the proxy with a reference the pipeline cannot release, and the run would look normal while freeing nothing. Releasing is a latency cost, never a correctness one: a released component reloads on next access. Training keeps everything resident and warns if the flag is set. If no stage holds a deferred component the pipeline warns rather than silently doing nothing.
The example builds its OffloadConfig directly rather than going through FastVideoArgs.add_cli_args, so the new flag needs its own switch to be reachable from the motivating command line.
Same reason as the MiniMax H3 T2V example: this script builds its OffloadConfig directly rather than going through FastVideoArgs.add_cli_args, so the flag needs its own switch here to be reachable. FastH3 is the case the flag exists for, since its components sum past what a 121 GiB unified-memory device can hold.
The schedule maps each deferred component to the last stage that holds it, derived once after create_pipeline_stages. Every pipeline in the tree builds its stages there, but nothing enforced it. A stage appended afterwards could hold a component an earlier stage had already been told to free, and would then be handed a released component mid-run with no error. add_stage now rebuilds the schedule and says so, turning an invariant nothing checked into a visible self-correction.
…lure
Review of the first version found that the flag was unsafe as a global
default. Releasing a component and loading it again is only correct when
nothing outside the loader has changed it, and two habits in this tree break
that without raising:
* mutating a component after load. LongCatPipeline.initialize_pipeline turns
on block-sparse attention and writes parameters into every transformer
block. That runs once, so a re-materialized component silently comes back
with the feature off and the second generation quietly degrades.
* reading a component's attributes while stages are built. The shared
DenoisingStage.__init__ derives the attention backend from
transformer.hidden_size, which materializes the DiT during post_init and
defeats the deferral it was meant to gain.
_lazy_module_names is now empty in the base class, so an unchecked pipeline
gets no deferral and says so. MiniMax-H3 opts in to the four components this
PR measured, and nothing else changes behaviour.
Also from review:
Releasing on the failure path. A stage's hook frees only what that stage is
the last user of, and it ran only after a successful forward. Both the stage
and the whole run now release on the way out, so the retry a memory
constrained caller attempts does not start from a worse position than the
request that just failed. A failing release cannot replace the exception
being propagated.
Walking into nested stages. Cosmos25AutoDenoisingStage keeps the transformer
inside child stages, so a one-level scan called it unreferenced and never
freed it. The scan now recurses through stages and containers with cycle
protection.
Keeping the proxy when torch.compile is skipped. The FSDP check ran after the
proxy had already been replaced by the real module, so an FSDP-wrapped
component lost both the compile and its release hook.
Not attaching the activation trace to a deferred component, since the hook
manager pins every module it wraps.
test_parser.py asserted on a whole serialized config dict, so adding a field
to OffloadConfig broke it. Grepping the field name could not find that; only
running the suite could.
The existing tests build stages and pipelines by hand, so they would stay green through both defects review found: a load_modules that never reaches the deferral, and a stage constructor that reads a component's attributes and materializes it during post_init. Three tests now run the production paths. Two call the real ComposedPipelineBase.load_modules with the component loader stubbed and a counter on it, asserting that only opted-in names become proxies and that the loader is never asked for them. The third builds the real MiniMax-H3 stages over tracked proxies and asserts nothing materialized. The third one was mutation-checked: adding a single transformer.patch_size read to MiniMaxH3DenoisingStage.__init__ makes it fail, which is the habit that defeats deferral in the shared DenoisingStage today.
Register pipeline-level compilation as a lazy materialization transform so every freshly loaded component receives the same compile setup. Keep whole-module compilation behind the proxy, preserve release scheduling, and add regression coverage for conditional and whole-module compile paths.
Defer transformer config access until LoRA conversion is actually requested so base inference keeps lazy DiTs unloaded through conditioning. Initialize transformer bookkeeping per pipeline and cover deferred no-LoRA construction, on-demand exclusion setup, and instance isolation.
…re on main defers the same four MiniMax-H3 components through h3_sequential_load, which auto-arms on unified memory, and this branch defers them through lazy_module_load. With both active they interleave over one set of modules: the H3 override strips the denoise modules from the load list, calls the base loader which wraps the text encoder in a proxy, then pops that encoder while the release schedule still holds it. lazy_module_load is off by default and has to be asked for, so when it is on it takes ownership and the H3 sequential path stands down. Passing neither leaves main's behaviour untouched.
770d2df to
df4cf71
Compare
|
This PR has merge conflicts with the base branch. Please rebase: git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease |
|
This PR has merge conflicts with the base branch. Please rebase: git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease |
3 participants
Problem
A pipeline materializes every component before the first stage runs, so peak memory is the sum of all components even though no two of them are needed at the same moment. For MiniMax-H3 r16 that sum is 96.3 GiB: text encoder 48.0, DiT 37.8, video VAE 9.7, audio VAE 0.6. For FastH3 it is 124.0 GiB.
The existing CPU offload flags cannot help, for two reasons. They act after loading has already finished, so they do nothing about the load-time peak. And on a unified-memory device such as GB10 there is no second pool to offload into, so moving weights to the host frees nothing.
A GB10 has 121 GiB. FastH3's components do not fit on it at all.
What this changes
Adds
lazy_module_load, off by default. Heavy components become aLazyModuleproxy that loads on first use. The pipeline derives, from what its stages actually hold, which stage is the last user of each component, and installs a release hook there. Peak becomes the largest overlapping set rather than the sum.Measured
One GB10, 1 GPU, single prompt, 124 frames. Every number below is
torch.cuda.memory_allocated()sampled where a component is loaded orreleased, so it measures component residency rather than the run's true peak,
which also includes activations.
MiniMax-H3 r16, 192x320, 4 steps
Peak 57.7 GiB against 96.3 GiB resident without the flag. The same command without the flag was killed by the machine's out-of-memory daemon during the DiT load an hour earlier, at
fsdp_load.py:584.The decisive moment is two adjacent log lines, 4 ms apart:
The DiT starts loading from 9.7 GiB rather than from 58.3 GiB.
FastH3 Preview v0.2, VSA-H3 attention on the Triton kernel
Component residency does not move with resolution, because it is set by the
weights. What moves is activation memory, which this change does not touch.
The last two rows load every component and then die in
minimax_h3_denoising.py:190, so the ceiling on one GB10 sits between 222,208and 347,200 packed tokens. That ceiling also reflects whatever else the shared
machine was running, since the out-of-memory daemon watches the host rather
than this process, so treat it as a lower bound.
DiT 35.05B parameters, 65.5 GiB. The four components sum to 124.0 GiB against the device's 121 GiB, so without the flag this model cannot load at all and there is no eager baseline to compare against.
Still applies to FastH3 Preview v1
FastVideo has since released FastH3 Preview v1. The tables above were measured on v0.2 and are left as recorded, but nothing this change depends on moved: the two checkpoints ship a byte-identical
transformer/config.json, the same 35.05B parameters, and the same 688 student tensors, so the component sum that overflows a 121 GiB device is unchanged.Measured with this flag on v1, using KyleNeverGivesUp/FastH3-4-step-Preview-v1-r16, which is v1 with the rank-reduced AdaLN converter from #1699 applied:
Both peaks match the same conversion applied to v0.2 to the byte, which is what you would expect if peak is set by the architecture and the packed token count rather than by which training run produced the weights.
The 345-frame row is the longest MiniMax-H3 generates, 14.38 s of synchronized video and audio in one pass on a single GB10. It needs FP8 to fit, and FP8 only reaches the feed-forward stack after #1780.
Cost
A released component is read from disk again on the next generation. Within one generation nothing reloads, since the release point is the last stage that holds the component, not the last stage that uses it. Across generations the whole set reloads, which for this model is several minutes. That is why the flag is off by default and documented as something to enable only when the model does not otherwise fit.
Design notes
Four things here would fail silently if done the obvious way.
The release hook lives on
PipelineStage.__call__, not in the pipeline's stage loop. Several pipelines overrideforward, and a loop-based hook would leave those pipelines deferring but never freeing, with no error.The proxy forwards
__class__. Callers branch onisinstance(module, FSDPModule)and similar. A proxy answeringFalsewould take the wrong branch and produce a wrong result rather than a crash. The cost is that an isinstance check materializes, which is the correct trade.Self-returning methods hand back the proxy, not the component. Stages write
self.vae = self.vae.to(device)in a dozen places. Returning the component there would replace the proxy with a strong reference the pipeline cannot release, and the run would look completely normal while freeing nothing.add_stagerebuilds the schedule if it runs after the schedule was built. Every pipeline in the tree builds its stages insidecreate_pipeline_stages, but nothing enforced it. A stage appended later could hold a component an earlier stage had already been told to free.Building the schedule also has to avoid
isinstanceon proxy values, since that would materialize every deferred component before the run starts. There is a test for exactly that.Training keeps every component resident and warns if the flag is set. If no stage holds a deferred component the pipeline warns rather than silently doing nothing.
Scope
Deferral is opt-in per pipeline.
_lazy_module_namesis empty in the base class and MiniMax-H3 lists the four components measured above. Every other pipeline is unchanged, and setting the flag there logs a warning and does nothing.The reason is that releasing a component and loading it again is only correct when nothing outside the loader has changed it, and two habits in this tree break that without raising:
LongCatPipeline.initialize_pipelineturns on block-sparse attention and writes parameters into every transformer block. It runs once, so a re-materialized component silently comes back with the feature off.DenoisingStage.__init__derives the attention backend fromtransformer.hidden_size, which materializes the DiT duringpost_initand defeats the deferral it was meant to gain.A pipeline opts in once someone has checked it against both.
LTX2Pipeline,MagiHumanPipeline, andLingbotVideoPipelineadditionally overrideload_moduleswithout delegating to the base implementation, so they would need that too.The flag lowers weight residency only. It does not change activation memory, so a resolution whose activations do not fit still does not fit. On one GB10, FastH3 runs at 512x896 and is killed at 640x1120, where q, k, v, and the VSA gate alone account for 18.5 GiB against roughly 30 GiB left after the weights. The repository's reference configuration is 768x1344 on four GPUs with
sp_size=4, which splits those activations across ranks.Relationship to the existing MPS special cases
denoising.py:653,decoding.py:343, andsr_denoising.py:279already do a hard-coded version of this behindif torch.backends.mps.is_available(), with matching reload paths keyed onfastvideo_args.model_loaded. Apple Silicon is unified memory too, so that is the same problem found on a different device.Those three differ from this change in two ways. They still load everything eagerly first, so they lower the steady-state footprint but not the load-time peak, which is where H3 on GB10 dies. And they never release the text encoder, which for H3 is the largest component at 48 GiB.
This change could subsume all three later. It does not touch them here.
Tests
31 cases in
fastvideo/tests/stages/test_lazy_module_load.py, CPU only, covering the proxy contract, the release schedule, the identity rule for self-returning methods, the late-stage rebuild, and the enablement rules. Placed understages/becausefastvideo/tests/loader/is not collected by any CI lane on this branch's base.lazy_module_loadis registered indocs/design/inference_schema_parity_inventory.yamland mapped inapi/compat.py, sotest_fastvideo_args_fields_are_classifiedstays green.