Generalized Tensor Parallelism (GTP) #3005
Conversation
Greptile SummaryThis PR adds a GTP-agnostic
Confidence Score: 3/5The grouped MLP backward path silently discards the return value of Two issues on the changed backward paths lower confidence. In transformer_engine/pytorch/ops/fused/grouped_mlp.py ( Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Module as TE Module (Linear/LNL/GL)
participant DW as distributed_weight.py
participant Impl as DistributedWeight impl (e.g. GTPShardedParam)
participant GEMM as CUDA GEMM kernel
participant Dist as distributed.py
Note over Module,Dist: Forward pass
Module->>DW: materialize_weight_for_forward(weight)
DW->>Impl: leader.materialize_group_for_forward()
Impl->>Dist: "gather_along_first_dim(shard, group, external_coalescing=True)"
Dist-->>Impl: gathered_weight, handle
Impl-->>DW: gathered_weight(s)
DW-->>Module: [gathered_weight, ...]
Module->>GEMM: GEMM(input, gathered_weight)
GEMM-->>Module: output
Note over Module,Dist: Backward pass
Module->>DW: materialize_weight_for_backward(saved_sharded_weight)
DW->>Impl: leader.materialize_group_for_backward()
Impl->>Dist: gather_along_first_dim(shard, group)
Dist-->>Impl: gathered_weight
Impl-->>DW: gathered_weight(s)
DW-->>Module: [gathered_weight, ...]
Module->>GEMM: wgrad_GEMM(input, grad_output)
GEMM-->>Module: wgrad
Module->>DW: finalize_weight_grads(saved_weight, [wgrad])
DW->>Impl: leader.finalize_group_grads(wgrad)
Impl->>Dist: "reduce_scatter_along_first_dim(wgrad, group, output=main_grad)"
Dist-->>Impl: sharded_wgrad, handle
Impl-->>DW: sharded_wgrad
DW-->>Module: [sharded_wgrad]
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant Module as TE Module (Linear/LNL/GL)
participant DW as distributed_weight.py
participant Impl as DistributedWeight impl (e.g. GTPShardedParam)
participant GEMM as CUDA GEMM kernel
participant Dist as distributed.py
Note over Module,Dist: Forward pass
Module->>DW: materialize_weight_for_forward(weight)
DW->>Impl: leader.materialize_group_for_forward()
Impl->>Dist: "gather_along_first_dim(shard, group, external_coalescing=True)"
Dist-->>Impl: gathered_weight, handle
Impl-->>DW: gathered_weight(s)
DW-->>Module: [gathered_weight, ...]
Module->>GEMM: GEMM(input, gathered_weight)
GEMM-->>Module: output
Note over Module,Dist: Backward pass
Module->>DW: materialize_weight_for_backward(saved_sharded_weight)
DW->>Impl: leader.materialize_group_for_backward()
Impl->>Dist: gather_along_first_dim(shard, group)
Dist-->>Impl: gathered_weight
Impl-->>DW: gathered_weight(s)
DW-->>Module: [gathered_weight, ...]
Module->>GEMM: wgrad_GEMM(input, grad_output)
GEMM-->>Module: wgrad
Module->>DW: finalize_weight_grads(saved_weight, [wgrad])
DW->>Impl: leader.finalize_group_grads(wgrad)
Impl->>Dist: "reduce_scatter_along_first_dim(wgrad, group, output=main_grad)"
Dist-->>Impl: sharded_wgrad, handle
Impl-->>DW: sharded_wgrad
DW-->>Module: [sharded_wgrad]
Reviews (24): Last reviewed commit: "fix comments" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
|
/te-ci L1 pytorch |
3e70bdf to
ed9ce68
Compare
Co-authored-by: Jieming Zhang <jiemingz@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>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
|
Want your agent to iterate on Greptile's feedback? Try greploops. |
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
for more information, see https://pre-commit.ci
…ped MLP" This reverts commit 8bb26f0. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
…red weights Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
- Take a single leader weight in the DistributedWeight dispatchers - Gather the FC2 grouped weight late in the fused grouped MLP Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
…ht before the NVFP4 dgrad dispatch Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
- Rename gather coalescing flag grouped -> external_coalescing; - Clean up DistributedWeight wiring in TE modules - Restructure _all_gather_nvfp4 Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
There was a problem hiding this comment.
We also need to support GTP in ops.GroupedLinear. If the op fusion fails (e.g. if we run with a new activation function or on an unsupport GPU arch), we should gracefully fall back to the unfused implementation.
There was a problem hiding this comment.
Yes you are right, we haven't support GTP in ops.GroupedLinear yet because for now I just focus on the critical path for Nemotron3/Nemotron4 pretraining on GB200/GB300.
I suggest keep the current changes minimal for this turn's review convenience and surely we can that support once any user fall back into the gtp-supported branch. How do you think?
|
/te-ci pytorch L1 |
Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>
| # Distributed weight: reduce-scatter the full per-rank wgrads into each sharded | ||
| # main_grad (also fires the Megatron grad-accum hook). | ||
| if is_dist_weight: | ||
| finalize_weight_grads(weights[0], w_list) | ||
|
|
||
| # Need to return dummy wgrads for Megatron-LM wgrad fusion if grad is already added | ||
| if fc_op._accumulate_into_main_grad: | ||
| # (wgrad fusion, or the distributed-weight reduce-scatter above) so it doesn't double-add. | ||
| if fc_op._accumulate_into_main_grad or is_dist_weight: | ||
| w_list = get_dummy_wgrads_for_params(weights) |
There was a problem hiding this comment.
finalize_weight_grads return value silently discarded — undocumented side-effect contract
finalize_weight_grads(weights[0], w_list) is called for its side effect (reduce-scatter into main_grad) and its return value is immediately replaced by dummy wgrads at line 649. Every other call site (linear.py, layernorm_linear.py, grouped_linear.py) assigns the return value and uses it as the actual wgrad.
This pattern is correct only if finalize_group_grads also writes the reduce-scattered result into main_grad in-place (the "Megatron grad-accum hook" path). Nothing in the DistributedWeight protocol spec or FakeDistributedWeight test stub documents or exercises this in-place side-effect requirement. A DistributedWeight implementation that follows the documented return-value-only contract (as FakeDistributedWeight does — it creates new tensors without touching any main_grad buffer) would have its per-expert gradients silently discarded on the grouped MLP path while the other module paths continue to work correctly. The discrepancy is invisible because _compute_grad_params returns dummy tensors regardless.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
|
/te-ci pytorch L1 |
3 participants
Description
Core-idea: add Generalized Tensor Parallelism (GTP), which is a flexible fine-grained sharding/just-in time materialization of both activations and parameters with efficient computation-communication overlap.
Mission: improve LLM pretraining efficiency through generalized tensor parallelism, enabling high performance, memory efficiency, ease of use, and strong scalability.
GTP introduction
GTP_introduction.md
How Mcore interacts with TE
• TE ships self-contained: Protocol + dispatcher no-op.
• GTP lives 100% in Mcore as one protocol implementer.
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: