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

[uv] use torch cpu instead of cuda by default #320

Merged
merged 1 commit into from
Dec 20, 2024
Merged

Conversation

rmehri01
Copy link
Contributor

Why

  • by default, the CUDA version of torch is installled which takes up a bunch of space, yet we don't support GPUs on replit
    • we also have similar overrides for tensorflow for this reason:
      var tf_variants = [...]string{"tensorflow-cpu", "tensorflow-cpu-aws", "tensorflow-gpu", "tensorflow-intel", "tf-nightly", "tf-nightly-cpu", "tf-nightly-cpu-aws", "tf-nightly-gpu", "tf-nightly-intel"}
  • this leads to deployments failing because their image size is too large, let's try using the cpu version instead

What changed

  • unfortunately we can't do the override in the same way as tensorflow since the cpu version of torch isn't available on pypi
  • instead, add extra indexes that will be checked before pypi, described here: https://docs.astral.sh/uv/configuration/indexes/

Test plan

  • make upm on a linux machine
  • make a python project that uses uv
  • do upm add torch and see the following output, notably torch==2.5.1+cpu:
--> uv add torch --index https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match
Resolved 11 packages in 513ms
Installed 10 packages in 135ms
 + filelock==3.16.1
 + fsspec==2024.10.0
 + jinja2==3.1.4
 + markupsafe==3.0.2
 + mpmath==1.3.0
 + networkx==3.4.2
 + setuptools==75.6.0
 + sympy==1.13.1
 + torch==2.5.1+cpu
 + typing-extensions==4.12.2

@rmehri01 rmehri01 added the bug Something isn't working label Dec 19, 2024
@rmehri01 rmehri01 requested a review from a team as a code owner December 19, 2024 02:48
@rmehri01 rmehri01 requested review from ryantm and removed request for a team December 19, 2024 02:48
@rmehri01 rmehri01 force-pushed the rmehri01/torch-cpu branch 2 times, most recently from af6c426 to ddf7ed2 Compare December 20, 2024 01:41
Copy link
Collaborator

@blast-hardcheese blast-hardcheese left a comment

Choose a reason for hiding this comment

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

Thanks!

@rmehri01 rmehri01 merged commit 8f4d06f into main Dec 20, 2024
5 checks passed
@rmehri01 rmehri01 deleted the rmehri01/torch-cpu branch December 20, 2024 02:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants