[feat]: run FastH3 across two DGX Sparks with Ray sequence parallel - #1803
Conversation
mergify
Bot
added
type: feat
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI
🔴 PR merge requirementsWaiting for
This rule is failing.
|
850db72 to
91c2a25
Compare
|
#1761 has now been independently rebased onto main and is mergeable with no conflicts. #1803 currently includes the #1761 commit stack because it was built on top of that work. Please merge #1761 first, then rebase #1803 onto the updated main so that #1803 contains only its follow-up changes. This would leave #1803 focused on the geometry fix, the dual-Spark Ray work, and the accompanying recipe documentation, which should be reviewed separately. |
|
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 |
Review: lazy module load + dual-Spark Ray pathRecall-oriented review of the full diff against Correctness1. 2. 3. 4. 5. 6. 7. 8. 9. 10. 11. 12. 13. 14. Docs15. Below the cut (verified, lower impact)
One candidate was refuted: 🤖 Generated with Claude Code |
…#1803 review Keep deferral, compile, and later generate() from fighting each other, stop LoRA bookkeeping from pinning a released DiT, and leave per-node NCCL/Gloo interface names alone.
…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. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…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.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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.
… the DiT before VAE decode Input prep and unpatchify were holding live VAE/DiT proxies just for two integers, which loaded the video VAE before Qwen and kept the DiT resident through decode. Auto-enable --lazy-module-load on unified memory and on single-GPU FastH3 examples.
Ray could not start a two-node FastH3 job because the executor was still abstract and NCCL env was not copied to workers. Document the QSFP bring-up, cookbook recipe, and measured 292 s / 587 s pair runs on top of lazy module load.
…page The pair YAML was already in this PR; the cookbook gallery only lists recipes that declare a family and hardware evidence.
… on worker IPs Sequential load dropped Qwen after the first request, so warmup+repeats crashed. Two 1-GPU Sparks also Gloo'd to 127.0.0.1 when node_gpus looked like a single node.
…d recipe Record the 512x896 1-GPU and dual-Spark medians we measured, and note native 480p is 480x832. Geometry stays a CLI/YAML knob.
On GB10 both flags auto-enable; sequential then strips DiT/VAEs before post_init, so VAE torch.compile logs as enabled but never attaches.
…#1803 review Keep deferral, compile, and later generate() from fighting each other, stop LoRA bookkeeping from pinning a released DiT, and leave per-node NCCL/Gloo interface names alone.
4a5d5ce to
809d74c
Compare
4 participants
Summary
Rebased onto current
mainso this PR does not fight #1793. Sequential MiniMax H3 start (h3_sequential_load, GPU-direct DiT load) stays as merged. This branch adds the rest of the GB10 / dual-Spark work on top of that.From #1761 (authors preserved)
Kyle's
lazy_module_load/LazyModulepath: load each opted-in component on first use and free it after its last stage, with in-process reload for a latergenerate(). MiniMax-H3 opts in. Satyam's follow-ups: compile across lazy reloads, LoRA setup must not materialize lazy DiTs, share compiled VSA graphs across H3 layers. Geometry is read from checkpointconfig.jsonso input prep / unpatchify do not hold a live DiT through VAE decode.GPU-direct DiT load is not re-landed here; it already merged in #1793.
Dual-Spark FastH3 (this PR's original work)
set_log_queue/clear_log_queue(Queue is not picklable across nodes) and copyNCCL_*onto workers.basic_fasth3.py --execution-backend {mp,ray}(auto-ray whenRAY_ADDRESSis set).spark_pair_env.sh, generate YAML, cookbook recipe.Sequence parallel replicates the DiT (~66 GiB/node), so each Spark still needs
h3_sequential_loadandlazy_module_load(both auto on unified memory).Measured on two GB10s over QSFP RoCE (~21 GB/s NCCL busbw), same alpine 768×1344 FastH3 recipe, Triton VSA, no FA4, sequential + lazy load, parallel VAE, cold process:
That 292 s is faster than the ~330 s one-Spark clip from earlier bring-up. It is not a lower bound: first decode still pays VAE
torch.compile; TAEH3 (#1795) is a separate opt-in and was not used.Test plan
pytest fastvideo/tests/worker/test_ray_distributed_executor.pypytest fastvideo/tests/stages/test_minimax_h3_sequential_start.py fastvideo/tests/stages/test_lazy_module_load.pysource examples/inference/optimizations/spark_pair_env.sh,ray starton QSFP IPs withFASTVIDEO_HOST_IPmatching--node-ip-address, thenpython examples/inference/basic/basic_fasth3.py --num-gpus 2 --execution-backend ray --vsa-kernel triton --no-fa4 ...