[Pytorch] Enable TE Sequential Op to consume extra_outputs from a previously run Op - #3320
[Pytorch] Enable TE Sequential Op to consume extra_outputs from a previously run Op#3320vthumbe1503 wants to merge 13 commits into
Conversation
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>
for more information, see https://pre-commit.ci
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
for more information, see https://pre-commit.ci
Greptile SummaryThe 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.
Confidence Score: 4/5The 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
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"]
Reviews (2): Last reviewed commit: "address review comment" | Re-trigger Greptile |
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
| for op in self._basic_ops: | ||
| op._lock_extra_channels() |
There was a problem hiding this comment.
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
Description
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: