[PyTorch] Fix distributed BF16 weights and trim NVFP4 saved state in the fused grouped MLP (#3545)

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

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>

* [PyTorch] Drop NVFP4 rowwise input copies from the fused grouped MLP 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>

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

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

* [PyTorch] Move the NVFP4 saved-state test into test_grouped_mlp.py

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>

* [PyTorch] Move the fused grouped MLP distributed-weight tests into test_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>

* [PyTorch] Build fused grouped MLP weight quantizer lists lazily

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>

* Apply suggestion from @ksivaman

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

* [PyTorch] Skip columnwise quantization of gathered distributed FC1 weights 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>

---------

Signed-off-by: Xiaowei Ren <xren@nvidia.com>
Signed-off-by: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Author: 4 people