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.

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

[2/2] Switch ao to use py_limited_api #1277

Merged
merged 1 commit into from
Dec 19, 2024
Merged

[2/2] Switch ao to use py_limited_api #1277

merged 1 commit into from
Dec 19, 2024

Conversation

janeyx99
Copy link
Contributor

@janeyx99 janeyx99 commented Nov 13, 2024

This PR enables torchao to build 1 python wheel across multiple python versions. How? Our previous PR #1276 replaced reliance on the torch_python/PYBIND11 API in favor of using torch.library to register custom ops. This means that ao's custom extensions are Python free now!

As a result, we can now safely make use of Python/setuptools py_limited_api flags to build a Python agnostic wheel by passing the right arguments into Extension and bdist_wheel. Though the core PR has landed: pytorch/pytorch#138088, torchao wheels are built on the stable release (2.5.1 for now) so torch_python.so will still be linked. I have tested locally that linking torch_python.so does not affect the wheel's python agnosticism, because torchao cuda extensions are python free and do not use anything from torch_python.

This will have repercussions on built + released wheels and will change their names to, for example, torchao-0.8.0-cp39-abi3-linux_x86_64.whl and should be installable on linux_x86_64 across multiple Python versions.

Why was this a draft?

I needed to figure out a plan to safely land this (aka test well) and now I'm confident we should land this. How did I test?

  1. build pytorch without my new change (so torchpython is unfortunately still linked)
  2. build the ao wheel with old pytorch
  3. install the wheel in the env (say 3.11)
  4. run test_ops.py and make sure everything passes
  5. NOW. SWITCH TO 3.9
  6. pip uninstall torch torchao
  7. python setup.py clean && python setup.py develop in pytorch to reinstall torch for 3.9 (and not 3.11)
  8. pip install the ao wheel built on 3.11 but here we are on 3.9
  9. run the tests, make sure they pass

The above shows that the torchao wheel built with this change is python agnostic.

Copy link

pytorch-bot bot commented Nov 13, 2024

🔗 Helpful Links

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

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

✅ No Failures

As of commit 1322f88 with merge base 2e032c6 (image):
💚 Looks good so far! There are no failures yet. 💚

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

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 13, 2024
@janeyx99 janeyx99 added the topic: for developers Use this tag if this PR is mainly developer facing label Nov 13, 2024
@janeyx99 janeyx99 changed the title [WIP] see what CI fails with this one lol [2/2][WIP] see what CI fails with this one lol Nov 13, 2024
@janeyx99 janeyx99 changed the title [2/2][WIP] see what CI fails with this one lol [2/2][DO NOT LAND] Switch ao to use py_limited_api Nov 13, 2024
@janeyx99 janeyx99 changed the title [2/2][DO NOT LAND] Switch ao to use py_limited_api [2/2] Switch ao to use py_limited_api Dec 18, 2024
@janeyx99 janeyx99 marked this pull request as ready for review December 18, 2024 23:43
@msaroufim msaroufim self-requested a review December 18, 2024 23:54
@@ -139,4 +140,7 @@ def get_extensions():
long_description_content_type="text/markdown",
url="https://github.com/pytorch-labs/ao",
Copy link
Contributor

Choose a reason for hiding this comment

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

can you remove -labs as well

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Will do as a separate PR, as that isn't related to this change

@janeyx99 janeyx99 merged commit 1beb6d0 into main Dec 19, 2024
46 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. topic: for developers Use this tag if this PR is mainly developer facing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants