Reuse cuBLASMp's all-gathered grad_output in the wgrad GEMM by almogsegal · Pull Request #3629 · NVIDIA/TransformerEngine

almogsegal · 2026-10-05T16:01:35Z

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Reuse cuBLASMp's all-gathered grad_output in the wgrad GEMM

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

greptile-apps · 2026-10-05T16:05:50Z

Confidence Score: 5/5

[High impact] The PR appears safe to merge, with no new blocking issues found.

Summary

Reuses cuBLASMp’s gathered output gradient for weight-gradient GEMMs when the format and library version support it.

Eligible weight-gradient GEMMs reuse cuBLASMp’s gathered gradient output.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Input-gradient GEMM gathers output gradient] --> B{Supported library and format?}
  B -->|Yes, immediate weight gradient| C[View gathered workspace]
  B -->|No, or delayed weight gradient| D[Gather separately]
  C --> E[Weight-gradient GEMM]
  D --> E

_{Reviews (3) · Last reviewed commit: "[PyTorch] Test the fused te.ops comm+GEM..." · Reviewed by Greptile}

With the cuBLASMp comm+GEMM overlap backend, the dgrad AllGather + GEMM gathers grad_output and the wgrad GEMM gathered it a second time, while Userbuffers wgrad reads the tensor that its dgrad all-gather left in the communication buffer. cuBLASMp 0.11.0 leaves the gathered B at the start of its workspace (a documented cublasMpMatmul workspace property), so wgrad now reuses it the same way: - nvte_comm_gemm_workspace() (comm_gemm.h, TE-internal) returns the comm-gemm context's workspace. - With the cuBLASMp backend, CommOverlap/CommOverlapP2P::get_buffer() returns a view of that workspace with the communicator's buffer shape and dtype, like the Userbuffers buffer view. - get_cublasmp_all_gather_output() builds the gathered tensor from it for unquantized and per-tensor FP8 data. Linear, LayerNormLinear, LayerNormMLP (fc2) and the fused te.ops backward use it instead of gather_along_first_dim; MXFP8 and other formats, a delayed wgrad and a skipped dgrad GEMM still gather. This needs cuBLASMp 0.11.0 or newer, since older versions keep the gathered B elsewhere in the workspace; the next commit checks the version. Tested on 4x H100 with cuBLASMp 0.11.0: the layer and multi-layer tests of test_comm_gemm_overlap.py with the cuBLASMp backend pass (39 passed, 13 MXFP8 cases skipped on H100), and test_fusible_ops_with_userbuffers.py --use-cublasmp passes in BF16 and FP8 delayed and current scaling. All-gathers per backward: row-parallel Linear 1 -> 0, LayerNormMLP 2 -> 1 (fc1 still re-gathers its input, as with Userbuffers). Signed-off-by: Almog Segal <[email protected]>

Reusing the gathered grad_output relies on cuBLASMp 0.11.0 leaving the gathered B at the start of its workspace. Earlier versions keep it after their cuBLASLt workspaces, so the view held other data. nvte_comm_gemm_workspace() now checks the loaded library with cublasMpGetVersion() and returns nullptr before 0.11.0; get_buffer() then reports that no gathered input is available, and get_cublasmp_all_gather_output() returns None, so the callers all-gather as before. Tested on 4x H100: with cuBLASMp 0.11.0, the row-parallel Linear backward issues no all-gather and the Linear and LayerNormMLP overlap tests with the cuBLASMp backend pass (9 passed). With cuBLASMp 0.10.0 it all-gathers once per backward and the same tests pass, as they do on TE main. Signed-off-by: Almog Segal <[email protected]>

timmoon10 · 2026-10-12T17:41:37Z

Pipeline 73310656

denera

Overall LGTM except for the cuBLASMp < 0.11 fallback relying on this try/except block.

Since you updated the fusible ops too, can you make sure test_fusible_ops_with_userbuffers.py also tests with use_cublasmp=True? Currently it's only testing the Userbuffers path, never exercises the cuBLASMp path.

I spotted two perf issues in my testing but I believe they are both internal to cuBLASMp, not TE code so I wouldn't consider them blockers for this PR.

In a GEMM+RS overlap with SPLIT_MULTICAST algo, cuBLASMp makes every GPU compute the fully-reduced result for every GPU's slice of the output, then each GPU throws away all but its own slice. It moves about TP-times more data over NVLink than necessary where TP is # of tensor-parallel ranks. Userbuffers avoids this by making each GPU reduce only its own slice. cuBLASMp should probably do the same unless there's a NCCL API limitation at play, in which case we should file a ticket with them to fix it.
With overlap configs set to use SPLIT_P2P reduce-scatter, the new cuBLASMp path is at perf-parity with the UB path in BF16, but there's a 2-3% deficit in FP8 on a DGXH100 node. The difference exclusively traces to the forward pass. The backward passes are at parity between UB and cuBLASMp. I believe the issue here is that TE's FP8 recipe uses fast accumulation in the forward pass and split accumulator in the backward pass (GEMM chunks inside UB respect these recipe settings). However, the cuBLASMp path in FP8 is locked to the split accumulator in both forward and backward because cuBLASMp itself does not expose any API for the user to control this setting. I have confirmed that the FP8 deficit disappears when I change the TE FP8 recipe to use the split accumulator everywhere as well, but that's not a viable fix. We want to be able to use the fast accumulator in the forward pass for FP8.

NOTE: I have not had the chance to compare perf for MXFP8 yet.

denera · 2026-10-16T09:21:06Z

+    try:
+        global_data = comm.get_buffer(shape=global_shape)
+    except RuntimeError as e:
+        if "No gathered input in the cuBLASMp workspace" not in str(e):
+            raise
+        return None


Relying on string-matching an NVTE_ERROR is very fragile and changing the error message would silently turn the fallback into a crash.

Even more importantly, this kind of a fallback costs us an exception PER CALL on cuBLASMp < v0.11, which is every backward pass of every layer with a GEMM+RS overlap in the forward pass. That's a non-trivial CPU overhead.

I would recommend just writing a dedicated cuBLASMp capability/version-check function that caches the result instead.

almogsegal · 2026-10-18T06:33:32Z

Overall LGTM except for the cuBLASMp < 0.11 fallback relying on this try/except block.

Since you updated the fusible ops too, can you make sure test_fusible_ops_with_userbuffers.py also tests with use_cublasmp=True? Currently it's only testing the Userbuffers path, never exercises the cuBLASMp path.

I spotted two perf issues in my testing but I believe they are both internal to cuBLASMp, not TE code so I wouldn't consider them blockers for this PR.

In a GEMM+RS overlap with SPLIT_MULTICAST algo, cuBLASMp makes every GPU compute the fully-reduced result for every GPU's slice of the output, then each GPU throws away all but its own slice. It moves about TP-times more data over NVLink than necessary where TP is # of tensor-parallel ranks. Userbuffers avoids this by making each GPU reduce only its own slice. cuBLASMp should probably do the same unless there's a NCCL API limitation at play, in which case we should file a ticket with them to fix it.

With overlap configs set to use SPLIT_P2P reduce-scatter, the new cuBLASMp path is at perf-parity with the UB path in BF16, but there's a 2-3% deficit in FP8 on a DGXH100 node. The difference exclusively traces to the forward pass. The backward passes are at parity between UB and cuBLASMp. I believe the issue here is that TE's FP8 recipe uses fast accumulation in the forward pass and split accumulator in the backward pass (GEMM chunks inside UB respect these recipe settings). However, the cuBLASMp path in FP8 is locked to the split accumulator in both forward and backward because cuBLASMp itself does not expose any API for the user to control this setting. I have confirmed that the FP8 deficit disappears when I change the TE FP8 recipe to use the split accumulator everywhere as well, but that's not a viable fix. We want to be able to use the fast accumulator in the forward pass for FP8.

NOTE: I have not had the chance to compare perf for MXFP8 yet.

Thanks for the feedback, Alp!

I've noticed that too and it's already fixed in cuBLASMp 0.12.0.
As for the FP8 path on H100, I made some improvements and I believe that perf in 0.12.0 is at least the same if not faster than UB. As for the fast accumulation, I'll expose an API to enable it in 0.12.0. I did some runs a few months ago and cuBLASMp was mostly faster than UB with MXFP8, see Add MXFP8 support with cuBLASMp #3145 (comment).

…er() error get_cublasmp_all_gather_output() fell back to the all-gather for cuBLASMp versions before 0.11.0 by catching the error that get_buffer() raises and matching its message. That breaks silently if the message changes, and with an older cuBLASMp it raises and catches an exception in every backward pass of every layer that overlaps an all-gather. nvte_cublasmp_version() now returns the version of the loaded cuBLASMp library, read once with cublasMpGetVersion() (0 without cuBLASMp), and Python exposes it as tex.nvte_cublasmp_version(). get_cublasmp_all_gather_output() checks a cached version >= 0.11.0 before calling get_buffer(), so get_buffer() errors are no longer caught. nvte_comm_gemm_workspace() uses the same function. Signed-off-by: Almog Segal <[email protected]>

test_fuser_ops_with_userbuffers only ran the Userbuffers backend, although the script supports --use-cublasmp and the fused te.ops backward reuses cuBLASMp's gathered grad_output. It is now parametrized over both backends and skips the cuBLASMp cases when Transformer Engine is built without cuBLASMp. Signed-off-by: Almog Segal <[email protected]>

timmoon10 · 2026-10-19T17:00:31Z

Pipeline 75064094

denera

LGTM after the latest round of changes. We can merge once pipeline comes back clean.

almogsegal requested review from ksivaman, ptrendx and timmoon10 as code owners October 5, 2026 16:01

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Oct 5, 2026

almogsegal changed the title ~~Reuse cuBLASMp gathered B~~ Reuse cuBLASMp's all-gathered grad_output in the wgrad GEMM Oct 5, 2026

greptile-apps Bot reviewed Oct 5, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/csrc/extensions/comm_gemm_overlap.cpp

almogsegal added 2 commits October 6, 2026 23:13

almogsegal force-pushed the cublasmp-reuse-gathered-b branch from cc3944f to 2ff7268 Compare October 8, 2026 07:59

denera self-assigned this Oct 9, 2026

denera self-requested a review October 9, 2026 15:57

timmoon10 requested a review from vcherepanov-nv October 12, 2026 17:25

denera requested changes Oct 16, 2026

View reviewed changes

almogsegal added 2 commits October 18, 2026 09:46

denera approved these changes Oct 19, 2026

View reviewed changes

Navigation Menu

Repository navigation

Reuse cuBLASMp's all-gathered grad_output in the wgrad GEMM - #3629

Reuse cuBLASMp's all-gathered grad_output in the wgrad GEMM#3629
almogsegal wants to merge 4 commits into
NVIDIA:mainfrom
almogsegal:cublasmp-reuse-gathered-b

almogsegal commented Oct 5, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented Oct 5, 2026 •

edited

Loading

Uh oh!

Uh oh!

timmoon10 commented Oct 12, 2026

Uh oh!

denera left a comment

Uh oh!

denera Oct 16, 2026

Uh oh!

almogsegal Oct 18, 2026

Uh oh!

almogsegal commented Oct 18, 2026

Uh oh!

timmoon10 commented Oct 19, 2026 •

edited

Loading

Uh oh!

denera left a comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

almogsegal commented Oct 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Oct 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Diagram

Uh oh!

Uh oh!

timmoon10 commented Oct 12, 2026

Uh oh!

denera left a comment

Choose a reason for hiding this comment

Uh oh!

denera Oct 16, 2026

Choose a reason for hiding this comment

Uh oh!

almogsegal Oct 18, 2026

Choose a reason for hiding this comment

Uh oh!

almogsegal commented Oct 18, 2026

Uh oh!

timmoon10 commented Oct 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

denera left a comment

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

almogsegal commented Oct 5, 2026 •

edited

Loading

greptile-apps Bot commented Oct 5, 2026 •

edited

Loading

timmoon10 commented Oct 19, 2026 •

edited

Loading