[PyTorch] Support distributed weights in GroupedLinear's grouped-tensor path by fanshiqing · Pull Request #3517 · NVIDIA/TransformerEngine

fanshiqing · 2026-09-15T09:58:32Z

Description

#3005 added the DistributedWeight protocol and materializes weights upstream
of module.GroupedLinear's split-quantize / grouped-tensor fork, but wired up
only the split-quantize side. Two bugs followed:

Grouped-tensor path crashes. It closed its wgrad callables over the
transient gathered copies, so backward read .main_grad off a plain
tensor: AttributeError: 'Tensor' object has no attribute 'main_grad'
Re-gathering drops fp8 quantization. Forward discards the quantized
weights (the gathered buffer is transient) and backward re-gathers at the
shards' own precision, so bf16 shards under an fp8 recipe fail the dgrad
GEMM's scaling-mode check. Natively-fp8 shards are unaffected — which is why
this went unnoticed.

Why #3005 missed it:

Date	PR
05-30	#3038	Adds the grouped-tensor path, gated behind `NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM`, default off
07-23	#3005	Adds `DistributedWeight`; wires the split-quantize path only — reasonable, the other was opt-in
08-12	#3224	Adds the `use_grouped_tensor` constructor arg, making the unwired path reachable from the normal API

Known gap: bug 2 also affects ops/basic/grouped_linear.py, which feeds
re-gathered weights into general_grouped_gemm the same way
(a_is_low_precision == b_is_low_precision). This PR fixes the module paths;
the ops path is left for a follow-up.

Type of change

Bug fix (non-breaking change which fixes an issue)

Changes

Grouped-tensor: save the shards, take main_grad from grad_buffer(),
re-gather in backward, finalize via finalize_weight_grads.
Re-quantize a re-gathered weight unless it is already quantized.
Read weight_requires_grad from the parameters, not the materialized copy —
the protocol does not require materialize to propagate it, and a False
silently skipped the whole wgrad path.
Never cache a gathered copy in _fp8_workspaces (both paths): it would pin an
unsharded weight for the module's lifetime.
Reject delay_wgrad_compute + distributed weight, matching the existing guard
in ops/fused/grouped_mlp.py.
New tests/pytorch/test_module_grouped_linear_distributed_weight.py (8/8),
registered in qa/L0_pytorch_unittest. Existing grouped/dist-weight suites:
2023 passed, 624 skipped.

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-09-15T10:01:31Z

Confidence Score: 5/5

[Medium risk] Adds distributed weight support to grouped linear forward path.

The production changes appear safe to merge; the remaining issue is a non-blocking loss of split-quantize test coverage on pre-Hopper systems.

Findings

Fallback Tests Are Skipped ▶

Summary

This PR adds distributed-weight support to both GroupedLinear execution paths and registers focused regression tests.

Preserves original distributed parameters for gradient-buffer access and finalization.
Re-materializes and, when necessary, re-quantizes weights during backward.
Avoids caching transient gathered weights.
Adds tests for detached materializations, fused and unfused gradient accumulation, and single grouped weights.
The latest test capability guard is overly broad and suppresses supported split-quantize coverage on pre-Hopper GPUs.

_{Reviews (6) · Last reviewed commit: "Fix CI"}

…or path The DistributedWeight protocol (NVIDIA#3005) materializes weights upstream of the split-quantize / grouped-tensor fork, but only updated the former: * Grouped-tensor closed its wgrad callables over the transient gathered copies, so backward read .main_grad off a plain tensor (AttributeError). Save the shards instead, take main_grad from grad_buffer, re-gather in backward, and finalize through finalize_weight_grads. * Backward re-gathers at the shards' own precision, so bf16 shards under an fp8 recipe fail the dgrad GEMM's scaling-mode check. Re-quantize those; natively-fp8 shards are unaffected. Adds test_module_grouped_linear_distributed_weight.py, covering both paths in bf16 and MXFP8. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

zhongbozhu · 2026-09-17T21:07:39Z

@ksivaman can you take a look and approve?

ksivaman · 2026-09-22T15:15:02Z

/te-ci pytorch L0

ksivaman

LGTM

ksivaman · 2026-09-24T19:06:52Z

/te-ci pytorch L0

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

greptile-apps · 2026-09-25T02:02:19Z

+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() < (9, 0),
+    reason="Requires Hopper (SM90) or newer; the GroupedTensor grouped GEMM is unavailable.",
+)


Fallback Tests Are Skipped

This module-wide SM90 guard also skips every use_grouped_tensor=False case on pre-Hopper GPUs. The split-quantize fallback does not require the native GroupedTensor GEMM, so this removes distributed-weight regression coverage for a supported path on those systems. Please apply the capability skip only to cases that actually require the grouped-tensor path.

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!

ksivaman · 2026-09-25T02:06:36Z

/te-ci pytorch L0

fanshiqing requested a review from ksivaman as a code owner September 15, 2026 09:58

github-actions Bot added the org-contribution label Sep 15, 2026

fanshiqing linked an issue Sep 15, 2026 that may be closed by this pull request

GTP+TE integration NVIDIA/Megatron-LM#6619

Closed

greptile-apps Bot reviewed Sep 15, 2026

View reviewed changes

Comment thread tests/pytorch/test_module_grouped_linear_distributed_weight.py Outdated

Comment thread tests/pytorch/test_module_grouped_linear_distributed_weight.py Outdated

fanshiqing force-pushed the shiqingf/gtp-grouped-tensor-distributed-weight branch from b1f3a52 to f88b71f Compare September 15, 2026 10:28

fanshiqing mentioned this pull request Sep 15, 2026

fix(moe): select the MoE grouped-GEMM path from TE's distributed-weight support NVIDIA/Megatron-LM#7318

Merged

4 of 7 tasks

zhongbozhu reviewed Sep 15, 2026

View reviewed changes

Comment thread tests/pytorch/test_module_grouped_linear_distributed_weight.py

zhongbozhu reviewed Sep 15, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/module/grouped_linear.py

zhongbozhu suggested changes Sep 15, 2026

View reviewed changes

Comment thread tests/pytorch/test_module_grouped_linear_distributed_weight.py

Cover single_grouped_weight in the distributed-weight test

a5e31b7

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

zhongbozhu approved these changes Sep 17, 2026

View reviewed changes

Merge branch 'main' into shiqingf/gtp-grouped-tensor-distributed-weight

6d9cb9a

ksivaman previously approved these changes Sep 24, 2026

View reviewed changes

Merge branch 'main' into shiqingf/gtp-grouped-tensor-distributed-weight

e8c687e

Fix CI

8876ea9

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman dismissed their stale review via 8876ea9 September 25, 2026 01:59

greptile-apps Bot reviewed Sep 25, 2026

View reviewed changes

ksivaman approved these changes Sep 25, 2026

View reviewed changes

ksivaman merged commit 45de18a into NVIDIA:main Sep 25, 2026
23 of 24 checks passed

ksivaman mentioned this pull request Sep 25, 2026

[PyTorch] Keep DistributedWeight objects on ctx across saved-tensor hooks #3545

Open

1 of 13 tasks

Navigation Menu

[PyTorch] Support distributed weights in GroupedLinear's grouped-tensor path - #3517

[PyTorch] Support distributed weights in GroupedLinear's grouped-tensor path#3517
ksivaman merged 5 commits into
NVIDIA:mainfrom
fanshiqing:shiqingf/gtp-grouped-tensor-distributed-weight

fanshiqing commented Sep 15, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented Sep 15, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zhongbozhu commented Sep 17, 2026

Uh oh!

ksivaman commented Sep 22, 2026

Uh oh!

ksivaman left a comment

Uh oh!

ksivaman commented Sep 24, 2026

Uh oh!

greptile-apps Bot Sep 25, 2026

Uh oh!

ksivaman commented Sep 25, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

fanshiqing commented Sep 15, 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 Sep 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Findings

Summary

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zhongbozhu commented Sep 17, 2026

Uh oh!

ksivaman commented Sep 22, 2026

Uh oh!

ksivaman left a comment

Choose a reason for hiding this comment

Uh oh!

ksivaman commented Sep 24, 2026

Uh oh!

greptile-apps Bot Sep 25, 2026

Choose a reason for hiding this comment

Uh oh!

ksivaman commented Sep 25, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

fanshiqing commented Sep 15, 2026 •

edited

Loading

greptile-apps Bot commented Sep 15, 2026 •

edited

Loading