[PyTorch] Fix distributed BF16 weights and trim NVFP4 saved state in the fused grouped MLP by xrennvidia · Pull Request #3545 · NVIDIA/TransformerEngine

xrennvidia · 2026-09-18T08:28:17Z

Description

Two fixes for the CuTeDSL fused grouped MLP (ops.fused.GroupedMLP_CuTeGEMM*),
both surfaced while running Megatron-Core MoE experts with GTP-sharded weights
and per-module quantization recipes.

The op assumed gathered distributed weights were already quantized: forward
quantized the shards before the gather, and backward read _columnwise_data
or the NVFP4 amax straight off the gathered tensor. That fails for a module
whose per-module recipe does not create native-quantized params, e.g. MXFP8
experts under a global NVFP4 recipe, whose BF16 shards gather as BF16.
Gather first, then quantize whatever is still unquantized with the op's
per-expert weight quantizers, mirroring module/linear.py. Native-quantized
shards pass through unchanged. The per-expert weight quantizers are saved on
the forward context and used in backward, so a later forward under a
different (or no) recipe cannot reset them out from under an earlier
backward. The quantizer lists are built lazily, only on the discrete-weight
and distributed-weight paths that consume them.
The save path kept NVFP4 rowwise input data and scales that no wgrad path
reads (the cuDNN grouped kernel, the single-group dense GEMM and the generic
grouped GEMM fallback all take the columnwise tiles). Prune them as already
done for MXFP8. Per decoder expert layer of a Nemotron hybrid proxy
(83k tokens, hidden 3072) this removes ~136 MiB from the ~1.9 GiB saved set
with unchanged numerics.

The related Linear / GroupedLinear / LayerNormLinear loss of the
DistributedWeight subclass under saved_tensors_hooks is fixed generically
in prepare_for_saving by #3730 and is intentionally not part of this PR.

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

ops/fused/grouped_mlp.py: gather distributed weights before quantizing;
quantize re-gathered BF16 weights with the per-expert quantizers, built
lazily and saved on the FC1/FC2 contexts for backward; move the FC2 backward
re-gather ahead of the NVFP4 amax computation; drop NVFP4 rowwise input data
and scales from the saved state.
tests/pytorch/test_grouped_mlp.py:
- TestGroupedMLPDistributedWeight: fused grouped MLP with BF16 and MXFP8
  fake distributed weights under per-module MXFP8, optionally inside a global
  NVFP4 autocast, bit-exact against a plain reference up to the fake hooks'
  scales; a pass-through unit check for already-quantized gathered weights;
  and a backward after a second forward under no recipe or NVFP4.
- TestGroupedMLPFusedOp.test_nvfp4_saves_columnwise_inputs_only: exactly
  one FP4 data and one scale tensor saved per quantized input on the fused
  NVFP4 path, for group sizes 1 and 4. The single-group dense path saves a
  GroupedTensor wrapper, whose buffers are inspected directly.

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-18T08:31:47Z

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no new blocking issues were found.

Summary

The PR fixes distributed BF16 weights and reduces saved NVFP4 input data in the fused grouped MLP.

Distributed BF16 expert weights are quantized after they are gathered.
NVFP4 backward saves only the input tiles its weight gradients use.

_{Reviews (25) · Last reviewed commit: "[PyTorch] Skip columnwise quantization o..." · Reviewed by Greptile}

fanshiqing · 2026-09-21T23:51:11Z

LGTM, thanks for the fix!

ksivaman · 2026-09-25T20:09:39Z

@xrennvidia Can you rebase this with main? #3517 has been merged

xrennvidia · 2026-09-28T12:29:16Z

@xrennvidia Can you rebase this with main? #3517 has been merged

Done.

greptile-apps · 2026-09-28T13:02:59Z

Want your agent to iterate on Greptile's feedback? Try greploops.

ksivaman · 2026-09-29T17:58:44Z

/te-ci pytorch L0

timmoon10 · 2026-10-21T06:46:31Z

+        fc1_weight_quantizers = [
+            fc1_op.get_quantizer("forward", 2 * idx + 1) for idx in range(num_groups)
+        ]
        fc1_grad_output_quantizer = fc1_op.get_quantizer("backward", 0)
        fc2_input_quantizer = fc2_op.get_quantizer("forward", 0)
        fc2_weight_quantizer = fc2_op.get_quantizer("forward", 1)
+        fc2_weight_quantizers = [
+            fc2_op.get_quantizer("forward", 2 * idx + 1) for idx in range(num_groups)
+        ]


I'm worried that iterating through all the weight quantizers will add non-trivial CPU overhead. I propose accessing the weight quantizers lazily when they are required (when quantizing discrete weights or with GTP) so we don't unncessarily slow down other code paths:

Suggested change

fc1_weight_quantizers = [

fc1_op.get_quantizer("forward", 2 * idx + 1) for idx in range(num_groups)

]

fc1_grad_output_quantizer = fc1_op.get_quantizer("backward", 0)

fc2_input_quantizer = fc2_op.get_quantizer("forward", 0)

fc2_weight_quantizer = fc2_op.get_quantizer("forward", 1)

fc2_weight_quantizers = [

fc2_op.get_quantizer("forward", 2 * idx + 1) for idx in range(num_groups)

]

fc1_grad_output_quantizer = fc1_op.get_quantizer("backward", 0)

fc2_input_quantizer = fc2_op.get_quantizer("forward", 0)

fc2_weight_quantizer = fc2_op.get_quantizer("forward", 1)

xrennvidia · 2026-10-21T07:55:59Z

Hi @timmoon10 thanks for the review, I have addressed your comments, please take a review and see if they look good to you. Thanks.

ksivaman · 2026-10-22T21:33:43Z

For the first bug, an alternate approach is #3730 that reduces CPU overhead that also tackles the issue in QuantizedTensor to address all of the TE modules. Currently (in this PR) the bug will still exist in megatron in the LayerNormLinear module. The relevant from @xrennvidia is cherry-picked over there as well.

ksivaman · 2026-10-22T22:01:34Z

+                materialize_weight_for_forward(fc1_weights) if fc1_is_dist else fc1_weights,
+                fc1_weight_quantizers,
+                rowwise=True,
+                columnwise=input_requires_grad,


Suggested change

columnwise=input_requires_grad,

columnwise=False,

When weight is distributed, this columnwise quantized data is a throwaway tensor?

Yes, you are right, GTP forward does not need the columnwise data, but the condition should be columnwise=input_requires_grad and not fc1_is_dist.

xrennvidia · 2026-10-23T06:56:26Z

For the first bug, an alternate approach is #3730 that reduces CPU overhead that also tackles the issue in QuantizedTensor to address all of the TE modules. Currently (in this PR) the bug will still exist in megatron in the LayerNormLinear module. The relevant from @xrennvidia is cherry-picked over there as well.

Hi @ksivaman , I think you are right. Your 3730 supersedes my fix. So I think I should drop my fix and let your PR take care of it. Thanks.

A distributed (e.g. GTP) weight whose shards are not natively quantized is re-gathered as BF16. The fused grouped-MLP op assumed gathered weights were already quantized: forward quantized the shards before the gather, and backward read _columnwise_data / the NVFP4 amax straight off the gathered tensor, which fails for a module whose per-module recipe (e.g. MXFP8 experts under a global NVFP4 recipe) does not create native-quantized params. Gather first, then quantize whatever is still unquantized with the op's per-expert weight quantizers, mirroring module/linear.py's handling of a re-gathered BF16 distributed weight. Native-quantized shards gather as quantized tensors and pass through unchanged. The FC2 backward re-gather moves ahead of the NVFP4 amax / global-scale computation so that path no longer reads un-gathered shards. The per-expert weight quantizer lists are built once in forward and saved on the FC1/FC2 contexts; backward quantizes the re-gathered weights with those rather than reading the live op. A later forward on the same module under a different recipe, or with quantization disabled, resets the ops' quantizers before the earlier backward runs (pipelined microbatches, or a module shared across recipe scopes), which would otherwise hit None or quantize in the wrong format. Regression tests: a fake routed-expert DistributedWeight group whose materialize hooks return BF16 tensors (or MXFP8 tensors for the native-quantized case) drives the CuTeDSL fused grouped MLP under per-module MXFP8, optionally inside a global NVFP4 autocast. Distinct power-of-two scales in the forward/backward hooks make the gathered weights observable: output, dgrad and the main_grad wgrads must match a plain reference bit-exactly up to those scales. A unit check confirms already-quantized gathered weights pass through the helper unchanged. A further case runs a second forward under no recipe or NVFP4 between an MXFP8 forward and its backward and requires bit-exact agreement with a reference. Fails on the unpatched op with AttributeError on _columnwise_data in backward; passes with the fix. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Xiaowei Ren <xren@nvidia.com>

…saved state Every wgrad path reads only the columnwise tiles of the saved FC1/FC2 inputs: the cuDNN grouped wgrad kernel, the single-group dense GEMM (general_gemm with the NT layout unwraps both operands columnwise) and the generic grouped-GEMM fallback (NVFP4 forces a TN layout by taking the columnwise data for a non-transposed A and a transposed B, which is the wgrad case). The save path already pruned the rowwise data and scales for MXFP8 but kept them for NVFP4, citing the generic fallback, so each quantized NVFP4 input was held (and, under activation offloading, copied to host and back) twice. Prune them for NVFP4 as well, unconditionally. The amax tensors are untouched. Per decoder expert layer of a Nemotron hybrid proxy (83k tokens, hidden 3072) this removes a 121 MiB data and a 15 MiB scale tensor from the ~1.9 GiB saved set, with unchanged numerics on both the cuDNN kernel and the forced generic fallback (NVTE_DISABLE_CUTEDSL_WGRAD_FUSED_GROUPED_MLP=1). Adds a saved_tensors_hooks test asserting exactly one FP4 data and one scale tensor is saved per quantized input on the fused NVFP4 path, for group sizes 1 and 4, with token counts chosen so the inputs cannot be confused with weights in the count, and registers it in the L0 unit-test runner. Grouped storage is saved as raw buffers; the single-group dense path (cuDNN frontend >= 1.27) saves a GroupedTensor wrapper whole, so the test inspects the wrapper's buffers directly and requires its rowwise data and scales to be None. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Xiaowei Ren <xren@nvidia.com>

Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com>

The check that the fused grouped MLP saves only the columnwise NVFP4 input tiles lives naturally with the other fused grouped MLP tests. Move it into TestGroupedMLPFusedOp as test_nvfp4_saves_columnwise_inputs_only, drop the standalone file, and remove its separate line from the L0 launcher, which already runs test_grouped_mlp.py. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Xiaowei Ren <xren@nvidia.com>

…st_grouped_mlp.py The DistributedWeight dispatch tests for the fused grouped MLP (BF16 and native-quantized shard gather, the quantizer pass-through unit check, and the backward-after-recipe-change regression) move into test_grouped_mlp.py as TestGroupedMLPDistributedWeight, with their fake DistributedWeight helpers. Drop the standalone file and its separate L0 launcher line, which already runs test_grouped_mlp.py. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Xiaowei Ren <xren@nvidia.com>

The per-expert weight quantizer lists were built unconditionally at the top of the fused grouped MLP forward, costing 2 * num_groups quantizer lookups on every call, including the single grouped weight path that never uses them. Build them only where they are consumed: the discrete-weight FC1/FC2 quantization and the distributed-weight gather-then-quantize sites. Save the full list on the backward context only for distributed weights, which is the one path that re-quantizes in backward; every other path keeps saving just the first expert's quantizer as before. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Xiaowei Ren <xren@nvidia.com>

Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com>

…ights in forward The forward GEMMs read only the rowwise weight data; the columnwise copy exists for the dgrad GEMM in backward, which reads it from the saved weight. A distributed weight's gathered copy is not saved (the shards are, and backward re-gathers and quantizes columnwise itself), so the columnwise tiles produced in forward were built and immediately discarded. Match the FC2 distributed site and skip them. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Xiaowei Ren <xren@nvidia.com>

ksivaman · 2026-10-23T16:52:15Z

/te-ci pytorch L0

ksivaman

Looks good, thanks

xrennvidia requested a review from ksivaman as a code owner September 18, 2026 08:28

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

greptile-apps Bot reviewed Sep 18, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/module/grouped_linear.py Outdated

xrennvidia force-pushed the xren/dist-weight-saved-tensor-hooks branch from 45d5cb7 to eea3da3 Compare September 18, 2026 09:16

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

GTP+TE integration NVIDIA/Megatron-LM#6619

Closed

This comment has been minimized.

Sign in to view

greptile-apps Bot reviewed Sep 28, 2026

View reviewed changes

Comment thread tests/pytorch/test_module_distributed_weight_saved_tensor_hooks.py Outdated

xrennvidia force-pushed the xren/dist-weight-saved-tensor-hooks branch from 9d8b650 to dddf418 Compare September 28, 2026 13:32

greptile-apps Bot reviewed Sep 28, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/module/grouped_linear.py Outdated

xrennvidia force-pushed the xren/dist-weight-saved-tensor-hooks branch 2 times, most recently from c594c31 to 7638a3c Compare October 18, 2026 06:57

xrennvidia requested a review from timmoon10 as a code owner October 18, 2026 07:05

xrennvidia changed the title ~~[PyTorch] Keep DistributedWeight objects on ctx across saved-tensor hooks~~ [PyTorch] Fix DistributedWeight handling under saved-tensor hooks and in the fused grouped MLP Oct 18, 2026

greptile-apps Bot reviewed Oct 18, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/ops/fused/grouped_mlp.py Outdated

xrennvidia force-pushed the xren/dist-weight-saved-tensor-hooks branch from e6f2809 to 250974f Compare October 18, 2026 07:35

xrennvidia mentioned this pull request Oct 18, 2026

[PyTorch] Quantize distributed BF16 weights in the fused grouped MLP #3695

Closed

13 tasks

ksivaman requested review from ksivaman and removed request for ksivaman October 20, 2026 20:47

timmoon10 reviewed Oct 21, 2026

View reviewed changes

ksivaman mentioned this pull request Oct 22, 2026

Fix DistributedWeight handling under saved-tensor hooks #3730

Open

9 of 13 tasks

ksivaman removed the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Oct 22, 2026

ksivaman reviewed Oct 22, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/ops/fused/grouped_mlp.py Outdated

xrennvidia and others added 6 commits October 22, 2026 23:57

Update transformer_engine/pytorch/ops/fused/grouped_mlp.py

9a8faad

Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com>

xrennvidia force-pushed the xren/dist-weight-saved-tensor-hooks branch from d0bff45 to 8ca9d08 Compare October 23, 2026 06:58

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

xrennvidia changed the title ~~[PyTorch] Fix DistributedWeight handling under saved-tensor hooks and in the fused grouped MLP~~ [PyTorch] Fix distributed BF16 weights and trim NVFP4 saved state in the fused grouped MLP Oct 23, 2026

xrennvidia and others added 2 commits October 23, 2026 15:29

Apply suggestion from @ksivaman

3faf7bf

Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com>

ksivaman approved these changes Oct 23, 2026

View reviewed changes

Conversation

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

Uh oh!

Confidence Score: 5/5

Summary

Uh oh!

Uh oh!

fanshiqing commented Sep 21, 2026

Uh oh!

ksivaman commented Sep 25, 2026

Uh oh!

xrennvidia commented Sep 28, 2026

Uh oh!

This comment has been minimized.

Uh oh!

greptile-apps Bot commented Sep 28, 2026

Uh oh!

Uh oh!

ksivaman commented Sep 29, 2026

Uh oh!

Uh oh!

Uh oh!

timmoon10 Oct 21, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

xrennvidia commented Oct 21, 2026

Uh oh!

ksivaman commented Oct 22, 2026

Uh oh!

ksivaman Oct 22, 2026

Choose a reason for hiding this comment

Uh oh!

ksivaman Oct 22, 2026

Choose a reason for hiding this comment

Uh oh!

xrennvidia Oct 23, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

xrennvidia commented Oct 23, 2026

Uh oh!

ksivaman commented Oct 23, 2026

Uh oh!

ksivaman left a comment

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

xrennvidia commented Sep 18, 2026 •

edited

Loading

greptile-apps Bot commented Sep 18, 2026 •

edited

Loading