Skip to content

[Pytorch] Enable TE Sequential Op to consume extra_outputs from a previously run Op - #3320

Open
vthumbe1503 wants to merge 13 commits into
NVIDIA:mainfrom
vthumbe1503:enable_extra_out_consumption
Open

[Pytorch] Enable TE Sequential Op to consume extra_outputs from a previously run Op#3320
vthumbe1503 wants to merge 13 commits into
NVIDIA:mainfrom
vthumbe1503:enable_extra_out_consumption

Conversation

@vthumbe1503

Copy link
Copy Markdown
Collaborator

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

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

vthumbe1503 and others added 11 commits July 28, 2026 23:05
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
…h error handling tests

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@vthumbe1503
vthumbe1503 marked this pull request as ready for review August 5, 2026 23:58
@vthumbe1503
vthumbe1503 requested a review from timmoon10 as a code owner August 5, 2026 23:58
@vthumbe1503
vthumbe1503 requested a review from ptrendx August 6, 2026 00:00
@vthumbe1503 vthumbe1503 changed the title Enable TE Sequential Op to consume extra_outputs from a previously run Op [Pytorch] Enable TE Sequential Op to consume extra_outputs from a previously run Op Aug 6, 2026
@greptile-apps

greptile-apps Bot commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR adds named internal extra-tensor channels to PyTorch fusible operations, including forward routing, backward gradient accumulation, validation, documentation, and tests. The stale-routing fix now makes channel bindings immutable after fuser construction, but its lock also applies to temporary fusers used for standalone operation calls.

  • Adds producer-to-consumer channel resolution with fan-out support.
  • Keeps unbound extra tensors in the public Sequential interface.
  • Routes and accumulates internal channel gradients during backward.
  • Adds channel validation, mutation locking, tests, and usage documentation.

Confidence Score: 4/5

The PR should not merge until standalone BasicOperation execution stops permanently preventing later channel configuration.

A direct BasicOperation call creates and discards a temporary fuser, but construction irreversibly locks the operation, causing later public channel setters to fail even though no retained routing snapshot exists.

Files Needing Attention: transformer_engine/pytorch/ops/fuser.py and transformer_engine/pytorch/ops/op.py

Important Files Changed

Filename Overview
transformer_engine/pytorch/ops/fuser.py Adds internal channel routing and gradient propagation, but permanently locks operations even when the fuser is temporary.
transformer_engine/pytorch/ops/op.py Adds the public channel-binding API and irreversible mutation guard used by both persistent and transient fusers.
transformer_engine/pytorch/ops/basic/grouped_linear.py Initializes the conditional extra-input count before BasicOperation sizes its channel metadata.
tests/pytorch/test_fusible_ops.py Adds broad channel routing and gradient coverage but does not cover configuring an operation after a standalone call.
docs/examples/op_fuser/op_fuser.rst Documents channel topology, public versus internal slots, fan-out, fuser boundaries, and fused-operation responsibilities.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A["BasicOperation.forward()"] --> B["Temporary OperationFuser"]
  B --> C["Capture channel routing"]
  C --> D["Permanently lock BasicOperation"]
  D --> E["Run operation and discard fuser"]
  E --> F["Later channel setter"]
  F --> G["RuntimeError"]
Loading

Reviews (2): Last reviewed commit: "address review comment" | Re-trigger Greptile

Comment thread transformer_engine/pytorch/ops/op.py
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Comment on lines +511 to +512
for op in self._basic_ops:
op._lock_extra_channels()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

P1 Transient fusers permanently lock channels

When a BasicOperation is executed directly before its channels are configured, BasicOperation.forward creates and discards a temporary OperationFuser, but this constructor permanently locks the operation. A later set_extra_input_channel or set_extra_output_channel call therefore raises RuntimeError even though no retained fuser depends on the captured routing.

Knowledge Base Used: PyTorch Fusible-Operation Framework

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant