Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[torch.compile] PostGradPassManager, Inductor code caching fix, fix_functionalization pass refactor + tests #10273

Merged
merged 19 commits into from
Nov 21, 2024
Merged

[torch.compile] PostGradPassManager, Inductor code caching fix, fix_functionalization pass refactor + tests #10273

merged 19 commits into from
Nov 21, 2024

Conversation

ProExpertProg
Copy link
Contributor

@ProExpertProg ProExpertProg commented Nov 12, 2024 •
edited
Loading

This PR extracts the post_grad passes into PostGradPassManager. By overriding pickling of the class, we avoid failing to pickle custom passes, which is used in the code cache.

It also extracts pass config into PassConfig and adds an option to disable the reshapes pass. It adds utilities for registering custom inductor passes, and extracts fix_functionalization into a proper pass.

@github-actions GitHub Actions
Copy link

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

  • Add ready label to the PR
  • Enable auto-merge.

🚀

@ProExpertProg ProExpertProg force-pushed the luka/fix-inductor-pickle branch 3 times, most recently from 48bd06a to 9f9aba5 Compare November 12, 2024 23:27
@ProExpertProg ProExpertProg changed the title Extract post grad passes into manager Extract post grad passes into manager & extract fix functionalization to pass Nov 12, 2024
@ProExpertProg ProExpertProg force-pushed the luka/fix-inductor-pickle branch from 9f9aba5 to b051eb2 Compare November 13, 2024 21:52
@ProExpertProg ProExpertProg changed the title Extract post grad passes into manager & extract fix functionalization to pass [torch.compile] Extract post grad passes into manager & extract fix functionalization to pass Nov 13, 2024
@ProExpertProg ProExpertProg changed the title [torch.compile] Extract post grad passes into manager & extract fix functionalization to pass [torch.compile] PostGradPassManager, Inductor code caching fix, fix_functionalization pass refactor + tests Nov 13, 2024
@ProExpertProg ProExpertProg force-pushed the luka/fix-inductor-pickle branch from af751fc to 234ea39 Compare November 14, 2024 16:01
@mgoin mgoin added the ready ONLY add when PR is ready to merge/full CI is needed label Nov 14, 2024
@ProExpertProg ProExpertProg force-pushed the luka/fix-inductor-pickle branch from bbb6300 to 5af5547 Compare November 14, 2024 19:15
bnellnm
bnellnm reviewed Nov 14, 2024
View reviewed changes
vllm/compilation/config.py Outdated Show resolved Hide resolved
@mergify Mergify
Copy link

mergify bot commented Nov 17, 2024

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ProExpertProg.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify bot added the needs-rebase label Nov 17, 2024
@youkaichao youkaichao mentioned this pull request Nov 19, 2024
Merged
@ProExpertProg ProExpertProg force-pushed the luka/fix-inductor-pickle branch from 5af5547 to 8561146 Compare November 19, 2024 21:39
@mergify mergify bot removed the needs-rebase label Nov 19, 2024
@youkaichao
Copy link
Member

edge cases still:

image

@ProExpertProg
Copy link
Contributor Author

Huh, it worked for me, maybe I didn't check the final changes

@ProExpertProg ProExpertProg force-pushed the luka/fix-inductor-pickle branch 2 times, most recently from eecc7d3 to d0c0c45 Compare November 20, 2024 22:51
ProExpertProg and others added 19 commits November 20, 2024 17:52
@ProExpertProg
Extract post grad passes into manager. Handles passes, config, pickling.
5a1e3f2 
Add config option to disable reshape pass.

Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Refactor fix_functionalization into pass
887fc54 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Move is_func to inductor_pass.py
ba89bf0 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Final functionalization.py refactor - each case is now as low as 2 lines
658535a 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Add test for fix_functionalization
0b5c746 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Pickling strategy following the torch==2.6 uuid approach
7765deb 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Move pass manager to own file
16d2fb0 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Add PassManager caching test
0a32e90 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Move config into tests, explicitly disable caching in TestBackend
83aa179 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Add logging to fix_functionalization
b2a4492 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Fix PassConfig docstring
20a2212 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Fix log message V1
1d007c8 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
PR comments:
9c54e7c 
- rename functionalization.py
- add detailed comment about caching

Also:
- rename TODOs
- format

Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Add rotary embedding to func test
b97864f 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Add logging & timing to passes
343f3a3 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Inspect functions for UUID, add UUID to non-managed passes, clean up,…
d66893e 
… move VllmInductorPass into own file

Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Fix assert
86d92bb 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Fix config
894c63d 
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg
Fix class inspection
ee289bc 
Signed-off-by: Luka Govedic <[email protected]>
@ProExpertProg ProExpertProg force-pushed the luka/fix-inductor-pickle branch from d0c0c45 to ee289bc Compare November 20, 2024 22:52
@youkaichao youkaichao merged commit 8b0fe06 into vllm-project:main Nov 21, 2024
49 checks passed
@youkaichao youkaichao deleted the luka/fix-inductor-pickle branch November 21, 2024 05:45
@youkaichao
Copy link
Member

@ProExpertProg great job!

tlrmchlsmth pushed a commit to neuralmagic/vllm that referenced this pull request Nov 23, 2024
@ProExpertProg @tlrmchlsmth
[torch.compile] Inductor code caching fix (vllm-project#10273)
fc1de1b 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
Signed-off-by: Tyler Michael Smith <[email protected]>
mfournioux pushed a commit to mfournioux/vllm that referenced this pull request Nov 28, 2024
@ProExpertProg @mfournioux
[torch.compile] Inductor code caching fix (vllm-project#10273)
37168df 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
Signed-off-by: Maxime Fournioux <[email protected]>
sleepwalker2017 pushed a commit to sleepwalker2017/vllm that referenced this pull request Dec 13, 2024
@ProExpertProg
[torch.compile] Inductor code caching fix (vllm-project#10273)
666b5e6 
Signed-off-by: luka <[email protected]>
Signed-off-by: Luka Govedic <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@WoosukKwon WoosukKwon WoosukKwon left review comments

@SageMoore SageMoore SageMoore left review comments

@bnellnm bnellnm bnellnm approved these changes

@youkaichao youkaichao youkaichao approved these changes

@robertgshaw2-redhat robertgshaw2-redhat Awaiting requested review from robertgshaw2-redhat

@njhill njhill Awaiting requested review from njhill njhill is a code owner

@ywang96 ywang96 Awaiting requested review from ywang96 ywang96 is a code owner

@comaniac comaniac Awaiting requested review from comaniac comaniac is a code owner

@alexm-redhat alexm-redhat Awaiting requested review from alexm-redhat

Assignees
No one assigned
Labels
ready ONLY add when PR is ready to merge/full CI is needed
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
6 participants
@ProExpertProg @youkaichao @SageMoore @WoosukKwon @bnellnm @mgoin