Repository navigation
Add mHC pipeline and virtual-pipeline support - #7336
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
/ok to test 83a5e68 |
|
/ok to test bec3054 |
Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
Signed-off-by: jingqiny-99 <jingqiny@nvidia.com>
|
/ok to test d1424bd |
The 64-hidden, four-head BF16 fixture has head_dim=16. FlashInfer 0.6.18.post1 routes it to the RopeQuantize fallback, whose 16-element NeoX loads cross the eight-element half-head boundary. Use head_dim=64 to exercise the standard fused kernel while retaining all existing ASYNC repeatability assertions. No production code changes. Signed-off-by: Deyu Fu <deyuf@nvidia.com>
|
Both attempts failed the same ASYNC repeatability check in the FlashInfer RoPE scenario. The fixture had head_dim=16; the pinned FlashInfer BF16 fallback uses 16-element NeoX loads, which cross the 8-element half-head boundary (kernel source). 63edc5f changes only this fixture to head_dim=64 (hidden_size=256). The fused path and all parity assertions remain enabled. Static address and formatting checks pass; GPU CI on the new head is pending. |
Use the current main RoPE fixture and its matching 64-wide head assertion. Preserve main GDN padding/recompute coverage and the reviewed deterministic normalization guard. Signed-off-by: Deyu Fu <deyuf@nvidia.com>
|
/ok to test 9a6c8de |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38235781252 |
Keep the two GB200 mHC functional recipes compatible with main's explicit GDN backend selection. Deterministic mode requires the torch backend; leaving the default FLA backend aborts before training starts. Signed-off-by: Deyu Fu <deyuf@nvidia.com>
|
Fixed in 6837799: both deterministic GB200 mHC recipes now explicitly select --gdn-kernel-backend torch, as required by current main. This restores the previous deterministic backend choice. Production code, goldens, and deterministic checks are unchanged; CI is running. |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38277837860 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38279931814 |
|
/ok to test 2050273 |
Use full-precision native metrics from GitLab jobs 482990678 and 482990676. Three independent 20-step runs agree bit-exactly on LM loss, MTP loss and num-zeros for each case. The prior baselines were generated with the 26.08 image. Keep the deterministic comparisons unchanged. Signed-off-by: Deyu Fu <deyuf@nvidia.com>
|
Refreshed the two GB200 goldens in eb78363. Both mHC functional tests pass against the updated baselines, with repeatable full-precision LM/MTP loss and zero-gradient counts. Deterministic checks and tolerances are unchanged. |
|
Both mHC functional tests passed on eb78363, including checkpoint saves and exact deterministic LM loss, MTP loss, and zero-gradient-count checks. No skipped or NaN iterations. |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38342805521 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38348993234 |
What does this PR do?
Enable mHC with pipeline parallelism and interleaved virtual pipelines on main. Intermediate stages carry
hidden_size * mhc_num_residual_streamsactivations; both schedules now allocate their P2P buffers with that width, including physical-rank wraparound between virtual chunks.Related to #2919. This ports the pipeline capability from #2943 onto main's mHC interfaces. The separate CUDA graph PR #7337 covers graph execution.
Tracking PR #5795 also lists this capability as P3. The earlier stacked draft #5943 was closed to rebuild P3 on main after P2; its closure comment explains that sequencing. This PR adapts the transport to main's current mHC interfaces.
The generic PP2 batched P2P message-order prerequisite overlaps open draft #7177 (its standalone fix). Current main and dev both retain the old order, so this prerequisite is an independently reproduced pipeline bug fix. Reconcile it with #7177 as that work lands.
The deterministic GDN prerequisite addresses the same missing Q/K-normalization coverage as open dev PR #6613. This implementation also preserves FLA's backward formula using the output rounded to the input dtype, verified against the actual kernel. Main PR #6966 addresses shared autotune-cache configuration more broadly. Main PR #5396 changes the normalization call site behind an optional fusion flag; these changes should be reconciled when combined.
Changes
Validation
Current 26.09 GB200 baselines (
eb7836329c, 2026-10-12): the two deterministic recipes explicitly select the torch GDN backend. Goldens were refreshed from native pipeline 73292070: each case completed three independent 20-step runs with identical full-precision LM/MTP loss and num-zeros trajectories. That pipeline failed comparisons against the older 26.08 goldens; it is not a passing validation run. Fresh verification 73314320 passed both cases with two independent 20-step repeats each on this exact head: checkpoint saves, native pytest checks, and full-precision deterministic LM/MTP loss and num-zeros comparisons all passed, with no skipped or NaN iterations.Historical GB200 functional validation with 26.08 (2026-10-08,
977f638ed8): both registered mHC recipes now run on one 4-GPU GB200 node, with unchanged TP/PP/VPP/MTP settings and global batch size. Native GB200 goldens were generated from pipeline 72332839. Independent verification 72354991 passed both cases with two repeats each: all 20 steps, checkpoint saves, and full-precision deterministic LM loss, MTP loss and zero-gradient-count checks. The earlier H100 results below are historical; registered recipes and checked-in goldens now target GB200.DSv4-width transport validation after merging main
0ac6ffd385(2026-09-18): added six real NCCL schedule tests for hidden sizes 4096 / 7168 with four mHC streams (wire widths 16384 / 28672), covering PP2, PP2/VPP2, and PP2/VPP2/TP2 + sequence parallel. On one OCI node with four GB200 GPUs and the main 26.08 ARM CI image, all four ranks reported 6 passed; native runner exit was 0. An independent serial PyTorch reference checks every microbatch boundary, input gradient, local parameter gradient, and final loss withrtol=0, atol=0; the tests also verify output pseudo-deallocation. Executed atc7cf58bcf; the subsequent83a5e6853commit changes only copyright headers. These are transport-shape tests using small affine stages, not full DSv4 attention/MoE or TP/DDP gradient-finalization validation. Full DSv4 proxy testing on Lyris remains pending.Added real schedule comparisons against identically initialized PP1 models: loss and every finalized parameter gradient, PP2/VPP2, standalone embedding/loss chunks, empty Hybrid chunks, TP/SP + final-stage MTP, and CP with changing sequence lengths.
Added two functional recipes for TP2/PP2/MTP and PP2/VPP2 with an empty stage and MTP.
Main autoformatter's required Black 26.3.0, isort, pylint and Ruff checks pass; copyright and diff checks pass. Its advisory mypy step reports unresolved local GPU dependencies and existing typing errors.
Exact main 26.08 CI image on eight H100 GPUs: the complete 83-case PP compatibility file passed at
eaa79ac5d, including PP2/VPP2, embedding/loss-only and empty Hybrid chunks, TP2/SP/MTP, and all four BF16 CP/variable-length eager/recompute cases. Native runner exit and positive test counts were verified. Related Hybrid mHC and recompute files also passed 15 and 16 cases, respectively.The tests use native FP32 gradient accumulation and CP token-count loss normalization. Independent parity cases reset compiler caches to prevent earlier FP32 cases exhausting Dynamo's default specialization limit; the PP1/PP2 pair shares the same cache and unchanged numerical tolerances.
Both native functional recipes complete 20 training iterations and checkpoint save after the optimizer fixes. Six CUDA checkpoint regressions pass, including mixed-dtype state ownership and actual ShardedTensor global bounds.
GDN normalization: 8 eager/compiled FP32/BF16 kernel parity and exact-repeat cases passed; 12 full GDN/GDN2 deterministic cases passed, including TP2/SP/CP2. Fixed-input kernel characterization and independent optimizer-phase traces localized the original functional drift to this normalization path.
Both native functional recipes passed independent exact golden verification on the original main 26.08 H100 CI image: TP2/PP2/MTP and TP1/PP2/VPP2 with an empty stage and MTP. Each completed all 20 iterations and checkpoint save, with native exit 0, positive verification counts, and unchanged loss, zero-gradient and memory checks. Included goldens are the untouched native output. No tolerance or required metric was relaxed, and no throughput improvement is claimed.
The GB200 verification above used
N_REPEAT=2; both repeats passed for both registered recipes.The implementation is based directly on main and does not depend on the CUDA graph PR. The PR is ready for review; GitHub CI is tracking the current head.