FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Support paged stashing for GroupedLinear activations by lhb8125 · Pull Request #3423 · NVIDIA/TransformerEngine · GitHub

/ TransformerEngine Public

Support paged stashing for GroupedLinear activations - #3423

Merged
vthumbe1503 merged 2 commits into
NVIDIA:mainfrom
lhb8125:denliu/paged-stash-grouped-tensor
Sep 11, 2026
ConversationCommits (2)ChecksFiles changed
Merged

Support paged stashing for GroupedLinear activations#3423
vthumbe1503 merged 2 commits into
NVIDIA:mainfrom
lhb8125:denliu/paged-stash-grouped-tensor

Conversation

lhb8125 commented Aug 25, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

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

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Aug 25, 2026
lhb8125 mentioned this pull request Aug 25, 2026
Open
6 tasks
lhb8125 marked this pull request as ready for review August 25, 2026 10:57
lhb8125 requested a review from ksivaman as a code owner August 25, 2026 10:57
lhb8125 added 2 commits August 25, 2026 03:57
feat: mark GroupedLinear activations for paged stashing
90b0fd2 
Signed-off-by: hongbinl <hongbinl@nvidia.com>
Clarify paged stash grouped tensor markers
e4e4eed 
Signed-off-by: hongbinl <hongbinl@nvidia.com>
lhb8125 force-pushed the denliu/paged-stash-grouped-tensor branch from fb11b5c to e4e4eed Compare August 25, 2026 10:57

greptile-apps Bot commented Aug 25, 2026

Copy link
Copy Markdown
Contributor

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

lhb8125 mentioned this pull request Aug 25, 2026
Open

lhb8125 commented Sep 4, 2026

Copy link
Copy Markdown
Contributor Author

@vasunvidia @ksivaman Could you review this PR? thanks!

vasunvidia commented Sep 8, 2026

Copy link
Copy Markdown
Collaborator

Looks good to me

zhongbozhu commented Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Hi @timmoon10 @vthumbe1503 , can you guys take a look and merge this PR? This is for sync-free moe + grouped linear path

vthumbe1503 approved these changes Sep 10, 2026
View reviewed changes

vthumbe1503 left a comment •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality

Changes for this PR LGTM.

I have another PR here https://github.com/NVIDIA/TransformerEngine/pull/3350/changes#diff-3399df04d9593b14f5d1d886f72cf4fd752677dc02394cb9ad8d8719afe19193R308
that changes this method altogether

So it turns out, along with BF16 groupedtensor, this wont work for nvfp4 grouped tensor as well. Since the columnwise data is transposed, paged stashing would read the data incorrectly in case of NVFP4. But that would also need change at the MCore level. cc: @vasunvidia

I wouldnt block this PR for nvfp4 use-case. Changes for BF16/FP16 grouped tensor's paged stashing looks good to me.

Comment thread transformer_engine/pytorch/utils.py
tensor.columnwise_data is not None
), "Columnwise data is not set for grouped tensor"

if not hasattr(tensor, "columnwise_data"):

vthumbe1503 Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality

can we change this to

Suggested change
if not hasattr(tensor, "columnwise_data"):
if not isinstance(tensor, GroupedTensorStorage):

vthumbe1503 commented Sep 10, 2026

Copy link
Copy Markdown
Collaborator

/te-ci pytorch

vthumbe1503 merged commit 02f8e75 into NVIDIA:main Sep 11, 2026
23 of 28 checks passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Show hidden characters
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

vthumbe1503 vthumbe1503 approved these changes
ksivaman Awaiting requested review from ksivaman ksivaman is a code owner
+1 more reviewer
zhongbozhu zhongbozhu approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

community-contribution PRs from external contributor outside the core maintainers, representing community-driven work.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

Back | FazBrowse Home | New Git URL