[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}
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>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
…hrough module APIs Stop passing a dedicated argument through the upper-level module constructors. The tensor-parallel TE linear classes now resolve the active group themselves, so attention / Mamba / MLP / embedding / MTP keep their original APIs. Updated in generalized_tensor_parallelism.md accordingly. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
|
/ok to test 5065e9c |
| if not parallel_state.is_initialized(): | ||
| return None | ||
| if is_expert: | ||
| return parallel_state.get_expert_gtp_weight_remat_group(check_initialized=False) |
There was a problem hiding this comment.
I think we try to refrain from using parallel_state directly :/ See the above method with the deprecation notice.
| def release(self, ticket: int): | ||
| """Return the buffer to the pool (ticket stays valid). | ||
|
|
||
| slot.buf is intentionally NOT cleared: get() must stay idempotent so CUDA-graph-captured |
There was a problem hiding this comment.
[CRITICAL Correctness] release() publishes slot.buf back to the pool while deliberately retaining the same object in the persistent ticket. A later same-key ticket can therefore pop that buffer while the old ticket still returns it, so multiple weights permanently alias one full-weight storage.
This becomes a real read/write race for adjacent same-key weights: all_gather_and_prefetch() issues the next weight's async AG before TransformerEngine launches the current weight's GEMM. The next AG can overwrite the shared output buffer while the current GEMM is reading it. The existing two-same-shape second-pass test reaches this path but only checks isfinite, which cannot detect finite outputs computed with the wrong weight.
Please add exclusive buffer liveness (or a capture-stable double/ring buffer with a consumer-done event) and a numerical parity test that gives the adjacent weights deliberately different values.
| # so param_is_not_tensor_parallel_duplicate still classifies it without GTP-specific code. | ||
| from megatron.core.tensor_parallel import copy_tensor_model_parallel_attributes | ||
|
|
||
| copy_tensor_model_parallel_attributes(gtp_shard, param) |
There was a problem hiding this comment.
[CRITICAL Distributed Correctness] Native ColumnParallelLinear / RowParallelLinear set weight.allreduce=False for expert-parallel weights before calling wrap_module_params_gtp(), but this replacement only copies _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS; that list does not include allreduce. The new GTPShardedParam consequently defaults to allreduce=True.
For native SequentialMLP with EGTP, DDP and _get_param_groups() then classify the expert shard as dense and select the dense DDP/optimizer/broadcast groups instead of the expert groups. Please preserve allreduce (and any other expert ownership metadata) when replacing the parameter, and add a native SequentialMLP + EGTP test; the current MoE coverage uses TE grouped experts and does not exercise this path.
| self.fp4_enabled = self.base_module.config.fp4 is not None | ||
| self.fp8_runtime_enabled = None | ||
| self.fp4_runtime_enabled = None | ||
| self.gtp_remat = self.base_module.config.gtp_weight_remat_size > 1 |
There was a problem hiding this comment.
[CRITICAL CUDA Graph] This gate only checks dense gtp_weight_remat_size. In the supported dense GTP=1, expert EGTP>1 configuration, a captured MoE module still issues EGTP AG/RS, but runner.gtp_remat remains false, so the runner skips EGTP side-stream registration/drains and never builds the post-replay GTP finalize-hook plan.
The missing hook path is deterministic with overlap-grad-reduce: GTP parameters use a manually driven DDP hook, that hook returns during graph capture, and the only post-replay invocation is also guarded by runner.gtp_remat. EGTP parameters are therefore never registered grad-ready. Please include expert_gtp_weight_remat_size > 1 in this decision (or detect GTP params owned by the runner) and add an EGTP-only MoE CUDA-graph test.
| intra_expt_dp_no_egtp_remat_group = process_groups_dict['intra_expt_dp_no_egtp_remat_group'] | ||
| mp_group = process_groups_dict['mp_group'] | ||
| expt_tp_pp_group = process_groups_dict['expt_tp_pp_group'] | ||
| expt_tp_pp_with_egtp_remat_group = process_groups_dict['expt_tp_pp_with_egtp_remat_group'] |
There was a problem hiding this comment.
[CRITICAL Numerical Correctness] The new EGTP-merged model-parallel group is used by the standard Adam/SGD expert path below, but _get_megatron_emerging_optimizer() returns before reaching this code and still assigns pg_collection.tp_ep_pp to non-layer-wise expert Muon/AdaptiveMuon groups. That group excludes the EGTP axis.
As a result, each EGTP rank reduces gradient statistics for only its local expert-weight shard. When clipping (or a grad-norm skip threshold) is enabled, peers can compute different partial norms and apply different clipping coefficients instead of one norm over all EGTP shards. Please use tp_ep_pp_with_egtp_remat for expert emerging-optimizer grad statistics as well, and add a non-layer-wise Muon + EGTP clipping test.
Victarry
left a comment
Victarry
left a comment
There was a problem hiding this comment.
LGTM for the changes in MoE.
GTP_remat peers hold distinct weight shards AND distinct micro-batches (like DP), but were entangled with DP through special-cased dp_cp_no_gtp_remat groups. This promotes gtp_remat to a first-class axis, orthogonal to DP/CP by default. Highlights: * get_data_parallel_group() defaults to the REPLICATE axis (DDP + optimizer); with_gtp_remat=True gives the full DP x CP x gtp_remat axis. The entire no_gtp_remat plumbing is gone. * Batch split, num-microbatches, and gradient scaling now address the full DP x gtp_remat axis explicitly. * Correct gradient scaling in both regimes: reduce-scatter mean + finalize AVG by default; SUM over gtp_remat under calculate_per_token_loss (where total_global_tokens already counts gtp_remat peers' tokens). * Folding gtp_remat into the tp_dp_cp group: only FP8 amax + the router expert-bias, both of which want the full distinct-data set. * GTP-awareness is concentrated in parallel_state + process_groups_config + the GTP module -- DDP, optimizers, and broadcast_params are now GTP-agnostic. * DCP elects writers correctly: sharded weights over the replicate group, egtp-replicated _extra_state over the full group (no duplicate-writer collisions). * GTP_remat metric/loss logging: average over the full data axis. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
- get_data_parallel_group / _world_size / _rank and the three expert siblings now default with_gtp_remat=True (gtp-inclusive). No-op when gtp_remat_size=1. - Explicit with_gtp_remat=False at the weight-replica/grad sites only: process_groups_config DDP/optimizer/FSDP field builders (dp, dp_cp, intra_dp_cp, expt_dp, intra_expt_dp); grad-norm sharded reductions (common_utils); weight-hash replica check (utils); GTP-chunk replica rank (generalized_tensor_parallelism); grouped-expert weight replica (moe_utils). - Drop the now-redundant with_gtp_remat=True from ~65 data-distribution callers. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
…#6) Parameters without is_gtp_weight_remat (e.g., norms, embeddings) were unnecessarily all-gathered across the GTP group before Newton-Schulz orthogonalization, wasting communication bandwidth and running NS on an incorrectly shaped matrix. Signed-off-by: Deepak Narayanan <deepakn94@gmail.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
|
/ok to test 0a2581a |
…roup Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
…e references Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
…ic path; gtp+fp8 param supported Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
…it; rename GTPConfig to GTPRematConfig Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
- Run GTP weight-remat init on the new config-container build path - Fix GTP checkpoint save and load - API cleanup: narrow the GTP public API to production symbols Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
|
/ok to test 934e7d9 |
7 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
Outline:
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_remat architecture (Mcore ↔ TE integration)
Class hierarchy: which linears shard
3.2 DDP buckets with (E)GTP_remat
3.3 Distributed checkpointing (DCP)
Testing
TE related MR: TE-3005
GPT's Architecture (Mcore + TE)
• TE carries no GTP code — yet its fwd/bwd still drive GTP_AG / GTP_RS via the weight subclass.
• Mcore passes the pre-sharded out_features into stock TE, then attaches the GTP surface POST-construction.
• The GTP weight IS a native MXFP8 tensor — it reuses TE's FP8; no stale BF16 copy.
• Save dequantizes FP8→BF16; load re-quantizes under
gtp_native_fp8_load_context.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.