[diffusion] feat: cross-node sequence parallelism (Ulysses x Ring) - #33327
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
67eeaaa to
a917ca5
Compare
|
/tag-and-rerun-ci |
|
Follow-up plan for which parallelism axes to take cross-node once this lands, ordered by expected ease (the per-axis communication analysis is now a section in #33704):
|
5a4a523 to
12a3341
Compare
num_gpus stays the total world size across every node (srt's tp_size-per-node convention); nnodes splits it evenly into local worker counts, node_rank offsets into the global rank space. Node 0 keeps the TokenizerManager/HTTP surface; other nodes run worker-only (mirrors srt's node_rank>=1 path) and block until their local workers exit. Rendezvous stays on loopback for nnodes==1 (unchanged single-node behavior); --dist-init-addr provides a node-reachable address otherwise. CPU intra-op thread budget now divides by the co-located (local) worker count, not the global total.
Row split becomes 2D: ring first (an outer, contiguous ring_chunk_len slice of the packed sequence), Ulysses second (an inner slice within that chunk) -- only Ulysses shards heads, ring never does. Inside attention, ring gathers every rank's local KV chunk once, attends each chunk's real (non-padding) prefix locally against the full local query set, and combines the per-chunk partial outputs via the standard online-softmax (logsumexp) merge -- exact up to floating-point reassociation, done in fp32. The real/padding boundary (used) is read from the already-host-synced cu_seqlens_q_host tuple, never via a new device sync, so the ring path doesn't regress the existing hot-loop no-sync invariant. ring_degree==1 is byte-identical to before this change (same row-split arithmetic, same non-ring FA call) -- zero regression risk for the production hash gate. ring_degree>1 changes accumulation order and is therefore NOT expected to be bit-identical to a ring_degree==1 run of the same request; validation is a same-config-determinism + SSIM/PSNR quality gate instead, pending on-box confirmation. Only a first-cut data-movement strategy (one all_gather of every rank's KV chunk per attention call); a pipelined P2P rotation that overlaps transfer with compute is a follow-up once correctness is established. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Replace the correctness-first all_gather (blocking full-ring transfer before any compute) with a rotating P2P double-buffer: each hop sends this step's KV to the next ring rank and receives the following hop from the previous rank before this step's attention compute is issued, so the transfer overlaps compute instead of gating it. ring_degree=1 is unaffected (this path is only reachable when ring_active).
Stack K and V into a single buffer before each ring hop so the P2P rotation is one send/recv pair instead of two, and drop the now- redundant .contiguous() calls (torch.stack already materializes a contiguous buffer). Same data moved, same overlap, no behavior change.
_build_local_embedding_layout's call site computed row_start/row_stop from Ulysses world_size/rank alone (both here and in the preceding block_token_tags slice) -- correct when ring_degree=1 (ring_rank is always 0, so the combined formula degenerates to the Ulysses-only one, which is why single-node testing didn't catch this), wrong once ring actually shards rows across nodes. Combine world_size = ulysses*ring and rank = ring_rank*ulysses_world_size + ulysses_rank (matching minimax_h3.py's row_start formula and the process-group's ring-outer/ ulysses-inner rank layout) before both call sites. Found via cross-node re-validation: a v2v request crashed with 'size of tensor a (5100) must match tensor b (6131)' inside _embed's text-row copy, isolated to the 2-node (ring_degree=2) config -- an identical request on 8 GPUs (ring_degree=1) with the same code was unaffected, pointing straight at this ulysses-only formula.
Same class of bug as the previous commit's local-embedding-layout fix: build_rope_cache computed row_start from Ulysses world_size/rank alone (named sp_ws/sp_rank but only ever holding the Ulysses component), degenerating to the right answer only when ring_degree=1. Use the same combined ring-outer/Ulysses-inner formula as forward()'s own row split. Found by the same cross-node v2v re-test, one layer further in after the previous fix: 'query, key and positions must have the same number of tokens' inside _apply_rope_qk, i.e. the RoPE position cache built for a larger (Ulysses-only) local_seq_len than the Q/K rows actually ring-sharded down to.
Adds the verified 2-node H200 (Ulysses8 x Ring2) topology to the deployment picker and cookbook: a new H200 "Cross-node" profile cell, a feature-table row replacing the stale "Ring is not compatible" claim, and a benchmark subsection with the measured per-step and pipelined denoise-stage numbers plus the determinism/precision caveats.
Extends the general Sequence Parallelism doc with a cross-node section (--nnodes/--node-rank/--dist-init-addr, the required --encoder-parallel replicate workaround, and the cross-topology numerics caveat) plus a reference benchmark using MiniMax-H3's measured cross-node numbers, alongside the existing single-node Wan2.2 example.
The contract test asserted that ring_size=2 raises NotImplementedError, which is exactly the limitation this branch removes, so it failed the moment it ran against the implementation. It now checks what ring actually constrains: nothing head-wise, since ring splits rows rather than heads, but the packed-sequence alignment must divide the combined ulysses x ring degree, because ring adds an outer row split on top of Ulysses's inner one. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Ring numeric validation — done (2×B200)Promised validation for the new ring paths, plus one pre-existing silent-corruption bug this surfaced and fixed. The bug (exists on main, not introduced here)With Ring execution added
Non-ring suffix keeps its bitwise tail-order contract; unsupported combos (batch>1, non-FA masked) raise loudly. Validation (u1r2 vs single GPU / u2r1, qwen-image 1024², 20 steps, FA, seed 42)2-rank kernel-level (vs fp32 SDPA reference, bf16): tail-pad + prefix + suffix + kv-prefix all MAE ~2e-4, max ~2.4e-3 ✓ e2e images (probe counters confirm each path engaged all 2400 block-calls):
Ring parity is at the same FP-reorder level as the long-validated Ulysses path. Unit tests (7) cover dispatch, merge math, and the tail-order contract. /tag-and-rerun-ci |
|
/tag-and-rerun-ci |
The ring>1 fallback gathered the full unified sequence and skipped SP attention, but its output is corrupt (silently: single vs u1r2 MAE 67/255, SSIM 0.25, deterministic on main). With the replicated-suffix path now ring-capable, zimage takes the standard suffix route under ring — validated at MAE 0.7/255 vs single GPU on qwen's twin path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Per-model ring validation, round 2: zimageFound another silent corruption on main: zimage's ring fallback ( Fix ( After (Z-Image-Turbo 1024², 9 steps, seed 42, FA, 2×B200): probe confirms the suffix-ring path engaged all 540 layer-calls; single vs u1r2 MAE 1.343/255, SSIM 0.9986 — normal FP-reorder for a 9-step turbo model (few-step models amplify float reordering; the Ulysses-migration acceptance for zimage sat at the same level). Per-model status so far: qwen ✅ (tail-pad ring 0.58/255 + replicated-prefix ring 0.74/255), zimage ✅ (0.79→suffix ring 1.34/255, main corruption fixed), wan (plain ring) running. |
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
/tag-and-rerun-ci |
Per-model ring validation — final matrix (2×B200, u1r2 vs single GPU, fixed seed)
kv-prefix (mova) shares the same helper + merge mechanism as prefix/suffix and is covered by the 2-rank kernel test (MAE ~2e-4 vs fp32 SDPA); no e2e run — the weights are large and the code path is identical up to the rep/shard split point. hunyuan-style generic varlen and explicit-mask/batch>1 combos raise loudly by design. Every MAE above is FP-reorder territory (few-step turbo amplifies it, consistent with the zimage Ulysses-migration acceptance). Unit suite: 30 passed. CI note: the earlier |
|
/tag-and-rerun-ci extra |
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
/tag-and-rerun-ci extra |
… era The qwen_image_t2i_2_gpus GT was generated while u=1,r>1 silently double-counted replicated text, so it froze the corrupt trajectory — the correctness fix now reads as a consistency failure against it. New GT published from this branch's CI output (ci-data-diffusion 211b2602f9, verified byte-identical), revision pin bumped, and the override calibrated against the old GT dropped back to defaults. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
/tag-and-rerun-ci extra |
CI red root-caused: the GT froze the bugThe Handled per the standard GT-refresh flow, scoped to this PR so main stays green until merge:
The earlier |
The consistency lookup prefers <platform>/<file> over the root file, so the refreshed root-level GT was shadowed by the stale h100 copy. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
/tag-and-rerun-ci extra |
CI: green on NV ✅All NVIDIA jobs pass, including Remaining red is the AMD mi300 / NPU pool, which is failing identically on unrelated concurrent PRs (RL, HiCache) — pre-existing, not from this change. Summary of what this PR now carries beyond the original H3 cross-node work:
Ready for review. |
…helpers usp.py's _ring_attention_varlen local-imported flash_attn_varlen_func, parallel_state, and attention.backends every call — none of the three create a cycle with usp.py, so hoisted to module level. denoise_loop.py's _ulysses_ctx/_ring_ctx were a local-import duplicate of parallel_state's get_ulysses_ctx/get_ring_ctx (added earlier this session for exactly this purpose); deleted the duplicates and call the shared versions directly. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
/tag-and-rerun-ci extra |
attention.flash_attention sorts before diffusion.triton.ulysses_qkv alphabetically; my earlier lint pass used bare 'ruff check' instead of this repo's actual pre-commit scope (isort==7.0.0 is the import-sort hook here, not ruff -- ruff is scoped to F401/F821/UP037 only), which missed it. Verified against the exact CI tool versions this time. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
/tag-and-rerun-ci extra |
…-node merge Three corrections, all traced to PRs merged since this page was drafted: - Ring now supports USPAttention's masked/tail-pad text path and its replicated-prefix/-suffix/-kv-prefix paths (sgl-project#33327) -- covers the joint text+image attention most models use. The page previously said these paths didn't support ring; that's no longer true. - Cross-node launch (--nnodes/--node-rank/--dist-init-addr) is merged (sgl-project#33327), not pending. Kept the honest caveat: only H3 is validated end-to-end across nodes; other models' new ring support is same-node- validated so far. - The legacy stacked-QKV UlyssesAttention set shrank to one user (Wan's VSA variant) after sgl-project#33923 migrated zimage/hunyuanvideo onto USPAttention; reworded from a migration-debt framing to the current, narrower and more precise one (a kernel without softmax LSE, not an unmigrated path). - Fixed a since-reversed default: the page said Q/K/V packing was opt-in (enable_packed_qkv_input_a2a); sgl-project#33667 made single-collective packing the default, and that flag now selects the older stream-pipelined alternative instead. - Ring's backend gate is now a declared capability (supports_ring_rotation(), sgl-project#33928) rather than a hardcoded enum check; reworded to match. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…gl-project#33327) Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
…gl-project#33327) Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Summary
--nnodes/--node-rank/--dist-init-addrlaunch infra (worker-only mode on non-head nodes, head node keeps the HTTP/TokenizerManager surface).batch_isend_irecv) instead of a blockingall_gather.denoise_loop.py's local embedding layout,minimax_h3.py's RoPE cache) found only under real cross-node re-validation after rebasing onto the current tree.Validation
--num-gpus 16 --nnodes 2 --sp-degree 16 --ulysses-degree 8 --ring-degree 2 --encoder-parallel replicate.Known limitation
--encoder-parallel auto's fold decision is not yet node-boundary aware and will try to fold the text encoder across nodes, crashing the Ref2VA reference encoder. Cross-node deployments must pass--encoder-parallel replicateexplicitly (documented in the cookbook and enforced nowhere in code yet — tracked as follow-up).USP ring coverage (follow-up commits on this PR)
The H3 ring components were generalized (
usp.py/parallel_state.py) and the remainingUSPAttentionbranches were made ring-capable, closing silent-corruption holes found during validation (details in the validation comment):ring_attn)real_seq_len = pad_start(batch-1, FA)NotImplementedError(was silently corrupt)NotImplementedErrorAlso fixed (pre-existing on main): with
u=1, r>1the replicated-token dispatch keyed on the Ulysses world size, silently double-counting replicated text on the ring (qwen e2e MAE 15/255 → 0.58–0.74/255 after the fix).Test plan
py_compileon all changed Python filesCI States
Latest PR Test (Base): ✅ Run #31231999230
Latest PR Test (Extra): ✅ Run #31231998922