Skip to content

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

Closed
zou3519 wants to merge 2 commits intogh/zou3519/1185/basefrom
gh/zou3519/1185/head
Closed

Add flag to fx.passes.split_module to normalize input names#157733
zou3519 wants to merge 2 commits intogh/zou3519/1185/basefrom
gh/zou3519/1185/head

Conversation

@zou3519
Copy link
Contributor

@zou3519 zou3519 commented Jul 7, 2025 •
edited by pytorch-bot bot
Loading

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

@zou3519
Add flag to fx.passes.split_module to normalize input names
031b2e4 
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

[ghstack-poisoned]
@pytorch-bot
Copy link

pytorch-bot bot commented Jul 7, 2025 •
edited
Loading

🔗 Helpful Links

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

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

✅ No Failures

As of commit 13030c1 with merge base 0f9c1b3 (image):
💚 Looks good so far! There are no failures yet. 💚

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 7, 2025
@zou3519
Update on "Add flag to fx.passes.split_module to normalize input names"
13030c1 
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

[ghstack-poisoned]
zou3519 added a commit that referenced this pull request Jul 7, 2025
@zou3519
Add flag to fx.passes.split_module to normalize input names
49cad78 
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

ghstack-source-id: 1c9cef1
Pull Request resolved: #157733
@zou3519 zou3519 requested review from BoyuanFeng and anijain2305 July 7, 2025 22:45
ezyang
ezyang approved these changes Jul 8, 2025
View reviewed changes
Copy link
Contributor

@ezyang ezyang left a comment

Choose a reason for hiding this comment

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

You might be able to juice the docstring a little more with an LLM, consider it.

@zou3519 zou3519 added the ciflow/trunk Trigger trunk jobs on your pull request label Jul 8, 2025
@zou3519
Copy link
Contributor Author

zou3519 commented Jul 8, 2025

@pytorchbot merge

@pytorchmergebot
Copy link
Collaborator

pytorchmergebot commented Jul 8, 2025

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@zou3519
Copy link
Contributor Author

zou3519 commented Jul 8, 2025

@pytorchbot cherry-pick --onto release/2.8 -c special

@pytorch-bot
Copy link

pytorch-bot bot commented Jul 8, 2025

❌ 🤖 pytorchbot command failed:

@pytorchbot cherry-pick: error: argument -c/--classification: invalid choice: 'special' (choose from 'regression', 'critical', 'fixnewfeature', 'docs', 'release')

usage: @pytorchbot cherry-pick --onto ONTO [--fixes FIXES] -c
                               {regression,critical,fixnewfeature,docs,release}

Try @pytorchbot --help for more info.

@zou3519
Copy link
Contributor Author

zou3519 commented Jul 8, 2025

@pytorchbot cherry-pick --onto release/2.8 -c fixnewfeature

pytorchbot pushed a commit that referenced this pull request Jul 8, 2025
@zou3519 @pytorchbot
Add flag to fx.passes.split_module to normalize input names (#157733)
e6a07f4 
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)
@pytorchbot
Copy link
Collaborator

pytorchbot commented Jul 8, 2025

Cherry picking #157733

The cherry pick PR is at #157793 and it is recommended to link a fixnewfeature cherry pick PR with an issue. The following tracker issues are updated:

Details for Dev Infra team Raised by workflow job

This was referenced Jul 8, 2025
Merged
Open
atalman pushed a commit that referenced this pull request Jul 16, 2025
@pytorchbot @zou3519
Add flag to fx.passes.split_module to normalize input names (#157793)
779e6f3 
Add flag to fx.passes.split_module to normalize input names (#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: #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 gh/zou3519/1185/head branch August 8, 2025 02:20
tvukovic-amd pushed a commit to ROCm/pytorch that referenced this pull request Aug 20, 2025
@pytorchbot @zou3519 @tvukovic-amd
Add flag to fx.passes.split_module to normalize input names (pytorch#…
71c68bc 
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ezyang ezyang ezyang approved these changes

@anijain2305 anijain2305 Awaiting requested review from anijain2305

@BoyuanFeng BoyuanFeng Awaiting requested review from BoyuanFeng

Assignees

No one assigned

Labels

ciflow/trunk Trigger trunk jobs on your pull request fx Merged release notes: fx release notes category

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@zou3519 @pytorchmergebot @pytorchbot @ezyang @facebook-github-bot