Repository navigation
[WS2][CUDA][MiniMax-H3] ws2 tensor / sequence / context parallel and the hierarchical cross-node collective (RFC #420 stack 3/5) - #569
Open
fusheng-ji wants to merge 15 commits into
Conversation
fusheng-ji
requested review from
EthanZero2Hero,
Flink-ddd,
KJLdefeated,
bitborne,
inaniloquentee and
ryankert01
as code owners
October 22, 2026 12:18
|
Important Review skippedToo many files! This PR contains 139 files, which is 39 over the limit of 100. To get a review, reduce the PR to 100 files or fewer by splitting it into smaller PRs or changing its base branch. Upgrade to a paid plan to raise the limit. This review couldn't start because sufficient usage credits or metered capacity aren't available. Add credits or update usage-based reviews in the billing tab, then retry.
 |
fusheng-ji
force-pushed
the
feat/h3-ws2-parallel
branch
from
October 22, 2026 14:37
900f962 to
d086b5c
Compare
fusheng-ji
added a commit
to fusheng-ji/RL-Kernel
that referenced
this pull request
Oct 22, 2026
… dual trajectory (RL-Align#571) on H100 at current branch heads
…fixed-order Triton GEMM Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit 650e2148c6f0e5dcdef136537e0e62f745020273) Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit d099c0b)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit b41f868)
…de collective, 2-node H100 evidence Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit 7f1296a970efb0cfcebd9300ff92551de08cf311) Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit d56c391)
…il-closed test instead of a skip Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit 74e5b83)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit 33cbcd6)
…l keys Review finding (RL-Align#568): jointly permuting x and g changed dW/dbias because the token axis was reduced in call order. All projection ops (audio in/out, text context, Q/K/V) now take an optional logical_keys [T,2|3] (negative first column = padding); rows are reduced in ascending key order with padding dropped, active keys must be unique. Key format follows RL-Align#562. Adds the reviewer's repro plus permutation/padding tests for every row and Q/K/V. Co-authored-by: Yitao Wang <75578057+walker-ai@users.noreply.github.com> Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit d2941f0)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit 5f4e34c)
…ready ordered rows Follows the canonical-row design of RL-Align#562: the Q/K/V node of the block graph supplies (batch row, position) keys so its dW is reduced in logical order, and rows that already arrive in canonical order are reduced in place (same bytes, no gather of x and g). Co-authored-by: Yitao Wang <75578057+walker-ai@users.noreply.github.com> Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit 0da0583)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit d086b5c)
…block 0 node-by-node and 50-block/generation gates Squashes the one-block branch (its row-by-row history: timestep sinusoid, RMSNorm, AdaLN rows, sequence/tensor parallel norm rows, final AdaLN, one block, 50-block qualification) onto the current test-h3 so the pull request carries a reviewable, signed history. RFC RL-Align#420. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
…pull request Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
…hard leaves Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
fusheng-ji
force-pushed
the
feat/h3-ws2-parallel
branch
from
October 22, 2026 21:04
d086b5c to
5ae9f5e
Compare
fusheng-ji
added a commit
to fusheng-ji/RL-Kernel
that referenced
this pull request
Oct 23, 2026
… on one 8-GPU node 128 tests pass (8-GPU NCCL cases included) and all six layouts (TP2, TP4, TP4xCP2, TP4xCP2+SP, TP8 with the node-local tree, TP8 with the 2x4 hierarchical tree inside one node) are bitwise equal to the single-GPU block. Host names replaced by node-a. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Flink-ddd
added
platform: cuda
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
This branch has not been deployed
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
Contract
h3-parallel-v1: eight canonical reduction units, so tensor parallel 1 / 2 / 4 / 8 add the same partials in the same tree.dX: unit j is block j of every fused weight, in q,k,v / hidden,gate order.SeqShard; sequence parallel: all-gather / reduce-scatter;coshardslices metadata with the rows.HierarchicalCollective(hierarchical_ipc_nccl_tree): node-local deterministic tree, NCCL all-gather of the group partials, fixed balanced add. The input dtype is the accumulation dtype, so the result matches the single-node tree bit for bit.h3_parallel_run.pylaunches a layout on N processes and compares every rank with the single-GPU canonical block.Prior art & reuse decision
DeterministicCollectiveDecision: canonical-unit partial sums plus a fixed tree; NCCL only moves bytes between nodes.
Results (H100 / Hopper, 2 nodes x 4 GPUs)
Measured on NVIDIA H100 (93 GiB, sm90), driver 580, CUDA 13.0, torch 2.13.0+cu130, pinned
MiniMax-H3@42ed227weights, at head5ae9f5eb. Block 0 of the pinned checkpoint, 264 packed rows; forward and input gradient compared bitwise with the single-GPU canonical block for every layout, and in addition every weight-gradient shard (to_q/k/v, to_out, ff hidden/gate/down) for the layouts without CP (TP2, TP4, TP8).Numbers:
reports/experiments/h3-ws2-parallel-h100/(one report per layout; host names are replaced bynode-a,node-b).Results (B200 / Blackwell, 1 node x 8 GPUs)
Measured on NVIDIA B200 (179 GiB, sm100), driver 580.126, CUDA 13.0, torch 2.13.0+cu130, pinned
MiniMax-H3@42ed227weights, clean tree at5ae9f5eb. One node with 8 GPUs, one process per GPU; same block, 264 packed rows and comparisons as the H100 table above (forward and input gradient against the single-GPU canonical block for every layout, plus every weight-gradient shard for the layouts without CP).Numbers:
h3-ws2-parallel-b200/(one report per layout,tests.json; host names replaced bynode-a).--local-size 4, which selects the hierarchical IPC + NCCL tree and treats the 8 GPUs as two groups of four inside one physical node. It exercises that code path; it is not a measurement of the cross-node network, which was only measured on H100 (two nodes).test_h3_tp_adaln.py/test_h3_sp_norm_adaln.pywere not run on the B200 either.Status
All five layouts are bitwise equal to the single-GPU reference on real nodes; ranks are packed node-major with 1 / 2 / 4 / 8 nodes supported.
B200 status: all six layouts, including the 8-rank cases on one 8-GPU node, are bitwise equal to the single-GPU block (128 of 128 tests pass); the cross-node collective is not measured on B200.
Tests
python -m pytest tests/models/minimax_h3/test_h3_parallel.py tests/models/minimax_h3/test_h3_projection_rows.py -q # 122 passed (4 x H100)The 8-GPU NCCL cases need eight GPUs in one node or two nodes: we only have 4-GPU Hopper nodes, so they are covered by the two-node runner above and not run on a single 8-GPU Hopper node (open for anyone with that resource, see the note at the top).
B200: the same two files on 8 x B200, 128 passed, 0 failed, 0 skipped (the 8-GPU NCCL cases ran on one node).
Scope / Known limitations
test_h3_tp_adaln.py::andtest_h3_sp_norm_adaln.py::8-GPU cases are skipped here. Help from anyone with a single-node 8 x Hopper machine is welcome.Strict mode is fail-closed: no TF32, no fast-math, no split-K; FP32 accumulation with one final rounding; unsupported shapes, dtypes and devices raise instead of falling back.
ws2_pipeline_parallelis not included. Not part of this PR: video projections (#562), O / gate_up / down GEMMs (#554, #535), full attention (#553), ROCm, SGLang adapter, pipeline parallel.