[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 regression or outstanding previous finding remains.

Summary

This PR extends module.GroupedLinear distributed-weight handling to the grouped-tensor path and preserves distributed-weight semantics across forward and backward.

• Saves original sharded parameters rather than transient gathered tensors.
• Re-materializes and, when necessary, re-quantizes weights for dgrad.
• Routes wgrad through distributed scratch buffers and finalization hooks.
• Avoids caching gathered weights and rejects delayed wgrad computation for distributed weights.
• Adds CUDA coverage for both grouped execution paths, detached materialized copies, fused and unfused wgrad, FP8 operation, and packed single-grouped weights.
• Registers the new suite in the L0 PyTorch test job.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[Original distributed shards] -->|materialize for forward| B[Transient gathered weights]
    B --> C[GroupedLinear forward GEMM]
    A -->|saved for backward| D[Autograd context]
    D -->|materialize for backward| E[Re-gathered weights]
    E -->|quantize when required| F[dgrad grouped GEMM]
    D -->|grad_buffer| G[wgrad grouped GEMM]
    G --> H[finalize_group_grads]
    H --> I[Shard gradients / main_grad]

Loading

_{Reviews (3) · Last reviewed commit: "Cover single_grouped_weight in the distr..."}

[zhongbozhu]

@ksivaman can you take a look and approve?