[PyTorch] Fix DistributedWeight handling under saved-tensor hooks and in the fused grouped MLP

[xrennvidia]

Description

Three fixes for distributed (e.g. GTP-sharded) weights in the PyTorch frontend,
all surfaced while running Megatron-Core MoE experts with fine-grained
activation offloading and per-module quantization recipes.

1. Linear and GroupedLinear lose the DistributedWeight subclass in
   backward when torch.autograd.graph.saved_tensors_hooks are installed:
   autograd unpacks the saved parameter as a fresh plain tensor, so backward
   takes the plain-parameter branch and fails with
   "weight was removed while fuse_wgrad_accumulation=True". Keep the
   DistributedWeight objects on the backward args and prefer them over the
   saved tensors, on both the split-quantize and native grouped-tensor paths.
2. The fused grouped-MLP op assumed gathered distributed weights were already
   quantized, which fails for a module whose per-module recipe does not create
   native-quantized params (e.g. MXFP8 experts under a global NVFP4 recipe).
   Gather first, then quantize whatever is still unquantized with the op's
   per-expert weight quantizers, mirroring module/linear.py. 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.
3. The fused grouped-MLP save path kept NVFP4 rowwise input data and scales
   that no wgrad path reads. 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.

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

• module/linear.py, module/grouped_linear.py: keep DistributedWeight
  objects on LinearBwdArgs / GroupedLinearBwdArgs /
  GroupedLinearGraphSafeBwdArgs and use them in backward.
• ops/fused/grouped_mlp.py: gather distributed weights before quantizing;
  quantize re-gathered BF16 weights with the per-expert quantizers, save those
  quantizers 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.
• New tests, all registered in qa/L0_pytorch_unittest/test.sh:
  • tests/pytorch/test_module_distributed_weight_saved_tensor_hooks.py:
    Linear and both GroupedLinear paths with identity saved-tensor hooks,
    using an in-repo fake DistributedWeight implementer. The grouped-tensor
    case skips when the native path is unsupported on the device / cuBLASLt.
  • tests/pytorch/test_ops_grouped_mlp_distributed_weight.py: 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 plus a backward after a second forward under no
    recipe or NVFP4..
  • tests/pytorch/test_ops_grouped_mlp_nvfp4_saved_tensors.py: 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]

Confidence Score: 5/5

[High impact] The PR appears safe to merge; no new actionable issue was established.

Summary

This PR keeps distributed-weight objects available during backward and fixes gathered-weight quantization in the fused grouped MLP.

• Linear layers keep distributed weights available during backward.
• The fused grouped MLP quantizes gathered weights with each expert’s recipe.
• The fused grouped MLP saves only the input layout used for weight gradients.

The earlier coverage thread was manually resolved without explanation. The current tests cover that concern and skip unsupported native grouped execution rather than counting a fallback as native coverage. The earlier recipe-change concern is addressed by saving the quantizers on the forward context.

Review was based on code inspection; GPU tests were not run.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Distributed weight shards] --> B[Gather full weights]
  B --> C{Already quantized?}
  C -->|Yes| D[Use gathered weights]
  C -->|No| E[Quantize with expert quantizers]
  E --> D
  D --> F[Fused grouped MLP forward]
  F --> G[Save expert quantizers and columnwise inputs]
  G --> H[Backward gathers weights again]
  H --> I[Quantize plain weights with saved quantizers]
  I --> J[Compute gradients]

Loading

_{Reviews (19) · Last reviewed commit: "[PyTorch] Drop NVFP4 rowwise input copie..." · Reviewed by Greptile}

[fanshiqing]

LGTM, thanks for the fix!

[ksivaman]

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

[xrennvidia]

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

Done.

[greptile-apps]

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

[ksivaman]

/te-ci pytorch L0