Skip to content

Add flag to fx.passes.split_module to normalize input names#157793

Merged
atalman merged 1 commit intorelease/2.8from
cherry-pick-157733-by-pytorch_bot_bot_
Jul 16, 2025
Merged

Add flag to fx.passes.split_module to normalize input names#157793
atalman merged 1 commit intorelease/2.8from
cherry-pick-157733-by-pytorch_bot_bot_

Conversation

@pytorchbot
Copy link
Copy Markdown
Collaborator

Stack from ghstack (oldest at bottom):

This is useful for vLLM, which runs AOTAutograd directly on graphs after
they have been split.

I created a new flag for this instead of reusing
keep_original_node_name (please let me know if you think I should reuse this).
The reasoning is:

  • The names of the placeholder nodes is different from the targets of
    the placehoder nodes. The targets are the actual input names.
  • Backwards compatibility: this API has been out for ~4 years, it
    looks public, and it has extensive public use. For example, this change
    would actually be BC-breaking to vLLM (they rely on the subgraph input
    names being different at the moment).

Test Plan:

  • new tests

cc @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv

This is useful for vLLM, which runs AOTAutograd directly on graphs after
they have been split.

I created a new flag for this instead of reusing
`keep_original_node_name` (please let me know if you think I should reuse this).
The reasoning is:
- The names of the placeholder nodes is different from the targets of
  the placehoder nodes. The targets are the actual input names.
- Backwards compatibility: this API has been out for ~4 years, it
  looks public, and it has extensive public use. For example, this change
  would actually be BC-breaking to vLLM (they rely on the subgraph input
  names being different at the moment).

Test Plan:
- new tests

Pull Request resolved: #157733
Approved by: https://github.com/ezyang

(cherry picked from commit b9afdd9)
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot Bot commented Jul 8, 2025

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/157793

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit e6a07f4 with merge base 3a7ff82 (image):

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@pytorch-bot pytorch-bot Bot added the release notes: fx release notes category label Jul 8, 2025
@zou3519 zou3519 added this to the 2.8.0 milestone Jul 16, 2025
Copy link
Copy Markdown
Contributor

@atalman atalman left a comment

Choose a reason for hiding this comment

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

lgtm

@atalman atalman merged commit 779e6f3 into release/2.8 Jul 16, 2025
106 of 107 checks passed
@Camyll
Copy link
Copy Markdown
Contributor

Camyll commented Aug 2, 2025

Validated tests on release/2.8

tvukovic-amd pushed a commit to ROCm/pytorch that referenced this pull request Aug 20, 2025
…157793)

Add flag to fx.passes.split_module to normalize input names (pytorch#157733)

This is useful for vLLM, which runs AOTAutograd directly on graphs after
they have been split.

I created a new flag for this instead of reusing
`keep_original_node_name` (please let me know if you think I should reuse this).
The reasoning is:
- The names of the placeholder nodes is different from the targets of
  the placehoder nodes. The targets are the actual input names.
- Backwards compatibility: this API has been out for ~4 years, it
  looks public, and it has extensive public use. For example, this change
  would actually be BC-breaking to vLLM (they rely on the subgraph input
  names being different at the moment).

Test Plan:
- new tests

Pull Request resolved: pytorch#157733
Approved by: https://github.com/ezyang

(cherry picked from commit b9afdd9)

Co-authored-by: rzou <zou3519@gmail.com>
@github-actions github-actions Bot deleted the cherry-pick-157733-by-pytorch_bot_bot_ branch September 1, 2025 02:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants