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

[fanshiqing]

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:

1. 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'

2. 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]

Confidence Score: 5/5

The PR appears safe to merge; no actionable correctness, security, or repository-rule violations remain.

Summary

This PR repairs distributed-weight handling across GroupedLinear’s grouped-tensor and split-quantize paths.

• Preserves original distributed parameters for gradient-buffer access and finalization.
• Re-materializes and, when necessary, re-quantizes weights for backward dgrad computation.
• Avoids retaining gathered, unsharded weights in persistent FP8 workspaces.
• Adds coverage for detached materialized weights and fused and unfused weight-gradient accumulation.
• Registers the new distributed-weight module tests in the PyTorch L0 suite.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    S[Distributed weight shards] -->|materialize for forward| G[Gathered weights]
    G -->|quantize when needed| F[GroupedLinear forward GEMM]
    S -->|save original shards| B[Autograd context]
    B -->|materialize for backward| BG[Backward gathered weights]
    BG -->|re-quantize when needed| D[dgrad GEMM]
    F --> W[wgrad GEMM]
    W --> GB[Distributed grad buffer]
    GB -->|finalize / reduce-scatter| S

Loading

_{Reviews (2) · Last reviewed commit: "[PyTorch] Support distributed weights in..."}

[zhongbozhu]

Nit: can we put it in the distributed folder of tests and following the programming pattern there?

[zhongbozhu]

will single weight be working in this case? Mcore can disable it for other reasons (like Muon), but TE should have single weight supported well

[zhongbozhu]

grouped tensor and single weights are orthogonal, we should include single weight as part of the test case, otherwise we should assert with clear error messages