[PyTorch] Keep DistributedWeight objects on ctx across saved-tensor hooks

[xrennvidia]

Description

_Linear and _GroupedLinear save the DistributedWeight parameter itself for
backward and gate their distributed-weight path on getting that object back
(is_distributed_weight(saved_weight)). That holds only while no
torch.autograd.graph.saved_tensors_hooks are installed: with hooks active,
SavedVariable::unpack rebuilds a fresh plain tensor from the unpack hook's
result, dropping the Python subclass and its is_distributed_weight marker.
The backward then takes the plain-parameter branch, whose weakrefs point at
the transient all-gathered weights, and fails with
"weight was removed while fuse_wgrad_accumulation=True".

Seen with Megatron-Core fine-grained activation offloading (which installs
such hooks around expert fc1/act) on bf16 GTP-sharded GroupedLinear experts.
Quantized weights were unaffected because their storage objects already stay
on ctx via prepare_for_saving/restore_from_saved, and the op-fuser path reads
the module's own weight attribute instead of the saved tensor.

Keep the DistributedWeight objects as ordinary ctx / LinearBwdArgs references
and prefer them over the saved tensors in backward.

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

Please list the changes introduced in this PR:

• Change A
• Change B

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, with the previous request for hook regression coverage fully addressed and manually resolved.

Summary

This PR preserves DistributedWeight Python objects across saved-tensor hooks so Linear and GroupedLinear backward passes continue using the distributed-weight path.

• Stores strong references to distributed weights outside save_for_backward.
• Prefers those references during backward before clearing them.
• Adds CUDA regression coverage for hooked and unhooked execution with fused and unfused weight-gradient accumulation.
• Registers the new regression suite in the PyTorch L0 test runner.

Diagram

sequenceDiagram
    participant Caller
    participant Forward
    participant Autograd
    participant Hooks as Saved-tensor hooks
    participant Backward
    participant DW as DistributedWeight

    Caller->>Forward: Run Linear / GroupedLinear
    Forward->>DW: Materialize weights for forward
    Forward->>Autograd: Save tensors
    Forward->>Forward: Keep original DistributedWeight reference
    Autograd->>Hooks: Pack and unpack saved tensors
    Hooks-->>Autograd: Return plain Tensor alias
    Autograd->>Backward: Invoke backward
    Backward->>Backward: Prefer retained DistributedWeight
    Backward->>DW: Materialize weights for backward
    Backward->>DW: Finalize distributed gradients

Loading

_{Reviews (3) · Last reviewed commit: "[PyTorch] Test DistributedWeight dispatc..."}