Support paged stashing for GroupedLinear activations

[lhb8125]

Description

Enable paged stashing integrations to identify the activation storage saved by the device-initiated PyTorch GroupedLinear autograd function.

• Mark input_to_save immediately before prepare_for_saving / save_for_backward.
• Extend mark_grouped_tensor to support unquantized BF16/FP16 rowwise storage.
• Preserve the existing quantized columnwise data and scale-inverse markers.
• Add unit coverage for rowwise and columnwise storage.

This is a follow-up to #3224. The paired Megatron-LM integration is NVIDIA/Megatron-LM#6828.

Validation

• Black 24.4.2 with the repository pre-commit arguments
• Ruff
• Full TE license check
• Four-rank GB300 Megatron-LM integration test using HybridEP, MXFP8, device-initiated GroupedLinear, and GPU-only paged stash: TestPagedStashingGroupedTensor::test_forward_backward_without_op_fuser
  • all four ranks passed
  • activation metadata captured
  • stash/reload completed without overflow
  • output and input-gradient parity passed

[greptile-apps]

Greptile Summary

This PR integrates device-initiated GroupedLinear activations with paged stashing while preserving the existing quantized-storage markers.

• Marks the selected activation immediately before flattening and saving the autograd state.
• Extends mark_grouped_tensor to mark unquantized rowwise BF16/FP16 storage.
• Continues marking quantized columnwise data and its inverse-scale metadata.
• Adds focused tests for plain, unquantized rowwise, and quantized columnwise storage.

Confidence Score: 5/5

The PR appears safe to merge with no concrete correctness, security, or compatibility defects identified.

The marker is attached to the same physical activation and scale tensors subsequently flattened for autograd, and the new branches consistently handle plain, unquantized rowwise, quantized columnwise, and absent saved inputs.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Marks the chosen backward activation immediately before its storage is flattened and passed to autograd.
transformer_engine/pytorch/utils.py	Extends marker placement to plain tensors and unquantized rowwise grouped storage while retaining quantized columnwise metadata handling.
tests/pytorch/test_grouped_tensor.py	Adds focused assertions for marker placement across plain, rowwise-only, and quantized columnwise representations.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
A[GroupedLinear activation] --> B{Saved representation}
B -->|Plain tensor| C[Mark activation]
B -->|Unquantized grouped| D[Mark rowwise data]
B -->|Quantized grouped| E[Mark columnwise data and scale inverse]
C --> F[prepare_for_saving]
D --> F
E --> F
F --> G[PyTorch save_for_backward]
G --> H[Paged stash and reload]

Loading

_{Reviews (1): Last reviewed commit: "Clarify paged stash grouped tensor marke..." | Re-trigger Greptile}