Fix DistributedWeight handling under saved-tensor hooks by ksivaman · Pull Request #3730 · NVIDIA/TransformerEngine

ksivaman · 2026-10-22T19:22:38Z

Description

An alternate approach to fixing the bug outlined in #3545 that tackles all modules at once. Cherry-picked parts of tests from @xrennvidia . The main advantages of this approach are tackling CPU overhead and having a central solution that covers all modules that use our restore_from* functions.

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

Compared to [PyTorch] Fix distributed BF16 weights and trim NVFP4 saved state in the fused grouped MLP #3545 , this makes a centralized change that touches all modules by fixing the core issue once in QuantizedTensor
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

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

prepare_for_saving and restore_from_saved run for every tensor every TE module saves, so the DistributedWeight check sits on a hot path for all users, not just GTP. Test None and torch.Tensor first, read the marker with an inline getattr (DistributedWeight implementers must be Tensor subclasses), and only query torch.compiler.is_compiling() for a flagged weight. On restore, test the common None/Tensor case first. On a CPU microbenchmark this brings the non-distributed cost back to roughly main's, e.g. 4.5 -> 5.2 us (was 10.0 us) per prepare_for_saving for a 32-expert bf16 GroupedLinear. Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Port the Linear and GroupedLinear saved-tensor hook tests from NVIDIA#3545 to the prepare_for_saving fix, and extend them: - Run each module case both under identity saved_tensors_hooks and under te.checkpoint(use_reentrant=False), whose pack hook discards saved tensors and returns recomputed ones in backward. - Cover LayerNormLinear, which also gates its backward on the saved weight being a DistributedWeight. - Add CPU-only checks that prepare_for_saving keeps a DistributedWeight by reference and that backward gets the same object back with no hooks, identity hooks, and torch.utils.checkpoint. Co-authored-by: Xiaowei Ren <xren@nvidia.com> Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

greptile-apps · 2026-10-22T19:26:48Z

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no actionable issues remain.

Summary

Keeps the original DistributedWeight object alongside its saved tensor so backward can still call its methods after saved-tensor hooks unpack it.

Distributed weights keep their identity through saved-tensor hooks.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    W[DistributedWeight] --> P[prepare_for_saving]
    P --> S[Tensor saved through autograd]
    P --> R[Original object kept by reference]
    S --> U[Unpack saved tensors]
    U --> B[restore_from_saved]
    R --> B
    B --> M[Provider gathers weight for backward]

_{Reviews (3) · Last reviewed commit: "Merge branch 'main' into fix_dweights_fo..." · Reviewed by Greptile}

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman · 2026-10-22T21:28:43Z

/te-ci pytorch L0

xrennvidia

LGTM

timmoon10

This fix is esoteric and I don't see how it addresses the bugs in #3545. Can you provide more explanation?

timmoon10 · 2026-10-23T20:22:37Z

+    reference and restored in place of the saved tensor: saved-tensor hooks
+    (e.g. activation offloading) unpack saved tensors as plain tensors, which
+    drops the subclass and its distributed-weight hooks."""


Can you explain in more detail? You're saying that something (in CPU offloading or Mcore or PyTorch?) replaces the contents of ctx.saved_tensors with a copy of the original tensors, so that they are not identical objects?

Yes pytorch autograd is the culprit. With a hook installed, autograd wraps the return object from the hook into a fresh torch.Tensor, so any attributes set on it are lost. This method works around autograd by saving a reference outside and using that instead. @xrennvidia's #3545 implemented a similar change (now removed) at the module level (linear and groupedlinear).

timmoon10 · 2026-10-23T20:24:29Z

+            # Skipped under torch.compile so traced graphs are unchanged.
+            if (
+                getattr(tensor, "is_distributed_weight", False)
+                and not torch.compiler.is_compiling()


Does this work? If the old logic was invalid for eager exectuion, I'm skeptical that it's valid under torch.compile. It might be better to error out if GTP and torch.compile are used together.

I'm not sure I follow your comment about the invalid logic? The modules force eager execution for torch compile with is_distributed_weight anyway. Although I agree that GTP + torch.compile is untested, so I've opted to error out here.

timmoon10 · 2026-10-23T20:40:11Z

    metadata: Dict[str, Any]


+class _SavedDistributedWeight(NamedTuple):


If I understand correctly (maybe false), the root problem is that we need the restored tensor to be identical to the saved tensor (not just identical data, but restored is saved) because the saved tensor might have weird monkeypatched attrs (like GTP callbacks). If that's the case, this fix is more general than DistributedWeight and it could apply to any class that requires identical tensor objects in the forward and backward.

This is the correct assessment

Maybe prematurely generalizing, but it may be clearer to rename this to something like _SavedTensorRef.

ksivaman · 2026-10-23T21:47:28Z

This fix is esoteric and I don't see how it addresses the bugs in #3545. Can you provide more explanation?

With #3545 in its current state, this PR is orthogonal to it.

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

timmoon10

Overall LGTM, with minor suggestions.

timmoon10 · 2026-10-23T22:19:08Z

+        pytest.skip("requires CUDA")
+
+
+def test_saved_tensor_hooks_unpack_subclass_as_plain_tensor():


Do we need to test this? This is more of a PyTorch contract. Even if they break it, it shouldn't affect us since we unconditionally restore the original tensor object for any tensor with is_distributed_weight=True.

ksivaman and others added 4 commits October 22, 2026 11:29

Fix DistributedWeight handling under saved-tensor hooks

619359f

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Merge branch 'NVIDIA:main' into fix_dweights_for_thooks

083c404

greptile-apps Bot reviewed Oct 22, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/quantized_tensor.py Outdated

Still use the autograd checker for the dist weight

5cb15c3

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

ksivaman requested review from timmoon10 and xrennvidia October 22, 2026 21:28

ksivaman mentioned this pull request Oct 22, 2026

[PyTorch] Fix distributed BF16 weights and trim NVFP4 saved state in the fused grouped MLP #3545

Open

7 of 13 tasks

xrennvidia approved these changes Oct 23, 2026

View reviewed changes

timmoon10 reviewed Oct 23, 2026

View reviewed changes

ksivaman added 2 commits October 23, 2026 15:01

Error for distweight + torch.compile

36db2ca

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Merge branch 'main' into fix_dweights_for_thooks

2b15ccf

timmoon10 approved these changes Oct 23, 2026

View reviewed changes

		metadata: Dict[str, Any]


		class _SavedDistributedWeight(NamedTuple):

		pytest.skip("requires CUDA")


		def test_saved_tensor_hooks_unpack_subclass_as_plain_tensor():

Conversation

ksivaman commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Diagram

Uh oh!

Uh oh!

ksivaman commented Oct 22, 2026

Uh oh!

xrennvidia left a comment

Choose a reason for hiding this comment

Uh oh!

timmoon10 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ksivaman commented Oct 23, 2026

Uh oh!

timmoon10 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ksivaman commented Oct 22, 2026 •

edited

Loading

greptile-apps Bot commented Oct 22, 2026 •

edited

Loading