[feat] Generalized Tensor Parallelism (GTP)#4967
Conversation
Co-authored-by: Jieming Zhang <jiemingz@nvidia.com> Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
fanshiqing
changed the title
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
- DDP: route the backward post-hook for GTP params through register_grad_accum_hook and skip the autograd AccumulateGrad hook, so grad-ready fires only after the GTP wgrad add. - GTP finalize path (_wait_reduce_scatter, finalize_grad=True): fire _handle_megatron_grad_accum after the add so terminal/async-only weights are not orphaned once the autograd path is suppressed. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
|
/ok to test c33667a |
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
|
/claude review |
There was a problem hiding this comment.
Large, well-documented feature PR with comprehensive test coverage (9 new test files). The architecture is sound — GTP/EGTP are correctly integrated as orthogonal parallelism axes with proper fallbacks when inactive.
One bug flagged inline: the new __getattr__ on ProcessGroupCollection breaks the hasattr-based fallback checks for the new GTP-related groups when a custom ProcessGroupCollection is constructed without those fields (backward-compat path). The fix is straightforward — use name in vars(pg_collection) instead of hasattr for the optional-group fallbacks, consistent with how the required-group checks were already updated.
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
…an-readable dtype in GTP weight cache Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Make megatron.core self-contained: it must not import from megatron.experimental, which is not shipped with the core wheel. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
…unner.stream fence Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
…ction fix2: make GTP module import gracefully without TransformerEngine Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
|
@ericharper @jaredcasper Hey Eric and Jared, can you help talk a look for this MR? Let me know if you have any concern for this MR so that I can change accordingly~ Any comments are welcome! |
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
|
/ok to test 00c9d20 |
| @@ -1191,6 +1282,7 @@ def __init__( | |||
| tp_comm_buffer_name: str | None = None, # Not used | |||
| tp_group: Optional[torch.distributed.ProcessGroup] = None, | |||
There was a problem hiding this comment.
make tp_group a list of [tp_group, gtp_remat_group].
6 participants
What does this PR do ?
Generalized Tensor Parallelism (GTP) is a light-weight, high-performance and memory-efficient distributed training strategy implemented in Megatron-LM and TransformerEngine. It shards weight tensors across an GTP process group and reconstructs them on-demand via async all-gather, enabling training of larger models without sacrificing throughput by overlapping communication with computation.
Design doc: GTP_design.docx
GTP_README.md
Features
1.1 Fine-grained, per-weight materialization & gradient reduction
1.2 CUDA graph compatibility
1.3 Low-precision quantize-then-gather
1.4 Composability with TP / SP / EP / DDP
1.5 Opt-in, minimally invasive integration
1.6 Optimizer-agnostic (Adam + Muon)
1.7 Scaling
1.8 Native distributed checkpointing (DCP)
Usage
2.1 Required flags
2.2 High-priority streams (Blackwell and later)
2.3 Minimal end-to-end example
2.4 Tuning knobs
Implementation details
3.1 GTP architecture (Mcore ↔ TE integration)
3.2 DDP buckets with (E)GTP
3.3 Distributed checkpointing (DCP)
Testing
TE related MR: TE-3005
GPT's Architecture (Mcore + TE)
① Mcore registers callbacks into TE at import time.
② TE calls back into Mcore runtime during te.Linear(gtp_group=…) init AND during fwd/bwd (
weight.all_gather_and_prefetch/wgrad_reduce_scatter).③ Mcore extensions forward gtp_group= at module init.
④ TE provides MXFP8 / NVFP4 tensor types AND the quantize-then-AG / RS collectives (
gather_along_first_dim,reduce_scatter_along_first_dim) — imported by Mcore runtime; GTP wraps them with its own schedule, buffer cache, and stream choreography.Changes Summary
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.