Issues creating a cuda-enabled pytorch environment with UV

[pstjohn]

I'm having some issues using uv lock & uv sync in a package that needs to use cuda-enabled pytorch wheels.

Here's a minimal-ish reproduction:

https://github.com/pstjohn/uv-torchvision-repro
The workspace stuff is there since I'd like to eventually get this working in a monorepo-like environment (#6935)

The desired behavior would be that uv lock would match what I get with
pip install --extra-index-url https://download.pytorch.org/whl/cu124 torch torchvision, which would be (at the time of writing) to install torch 2.4.1 and torchvision 0.19.1 both in their linux_x86_64 / cu124 / cp310 variants.

But instead I get the following error:

0.197 error: distribution torchvision==0.15.0 @ registry+https://download.pytorch.org/whl/cu124 
can't be installed because it doesn't have a source distribution or wheel for the current platform

If I use uv pip install --extra-index-url https://download.pytorch.org/whl/cu124 torch torchvision, I get torch==2.4.1+cu124 and torchvision==0.2.0. (i.e., seemingly not pulling torchvision from the cu124 index).

Any thoughts on where I might be going wrong? Thanks!

[pstjohn]

Just some additional data:

with torchvision == 0.19.1,

$ uv lock --extra-index-url https://download.pytorch.org/whl/cu124
  × No solution found when resolving dependencies:
  ╰─▶ Because there is no version of torchvision==0.19.1 and subpackage-b depends on torchvision==0.19.1, we can conclude that subpackage-b's
      requirements are unsatisfiable.
      And because your workspace requires subpackage-b, we can conclude that your workspace's requirements are unsatisfiable.

but with torchvision == 0.19.1+cu124, it resolves correctly.

[charliermarsh]

Have you taken a look at https://docs.astral.sh/uv/pip/compatibility/#local-version-identifiers?

[pstjohn]

Yeah that seems to explain it. It's odd that specifying torchvision==0.19.1 doesn't accept 0.19.1+cu124, though.

Having to pin the versions exactly to get cuda-enabled wheels isn't great for us. As a library we'd just like to depend on torch, but in setting up a python environment we'll need to be able to install the cuda-compiled dependency versions (which you typically do for different cuda versions just by changing the index-url, https://pytorch.org/get-started/locally/#start-locally)

I'm wondering if an option to --ignore-local-version-identifiers in uv lock would be enough to get this working? But there might be edge cases to those local identifiers I'm not following

[msarahan]

If I'm reading the docs right, it's that torchvision==0.19.1 doesn't match 0.19.1+cu124 because the whole spec is ignored, due to the fact that it has a local version identifier.

This general problem is something that I would like to solve by improving metadata. That's quite a rabbit hole, but if you'd like to learn more:

• https://discuss.python.org/t/what-to-do-about-gpus-and-the-built-distributions-that-support-them/7125/64 - Rehash just before PyCon 2024
• https://discuss.python.org/t/implementation-variants-rehashing-and-refocusing/54884 - Rehash after a long discussion with Oscar Benjamin and others
• https://discuss.python.org/t/selecting-variant-wheels-according-to-a-semi-static-specification/53446/111 - my latest proposal, which has side-tracked me to a PEP for index priority. @charliermarsh you may be interested in this. It is deeply influenced by uv's design.

[vajranaga]

If trying to add cuda-enabled pytorch to a project, I found the following process worked for me (Windows 11, but should work in Linux):

1. Install torch into the venv first with: > uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 in the project directory. This should be done first before other packages, as it will search the pytorch index for the related dependencies and will fail, likely, like this:

uv add torch --index-url https://download.pytorch.org/whl/cu124
  × No solution found when resolving dependencies:
  ╰─▶ Because jupyterlab was not found in the package registry and your project depends on jupyterlab>=4.2.5, we can
      conclude that your project's requirements are unsatisfiable.
  help: If this is intentional, run `uv add --frozen` to skip the lock and sync steps.

2. Add torch to the project dependencies with > uv add torch torchvision torchaudio.
   When I tried this first (before installing with uv pip install) I got the following error:

uv add torch --index-url https://download.pytorch.org/whl/cu124
Resolved 11 packages in 597ms
error: distribution torch==2.4.1 @ registry+https://download.pytorch.org/whl/cu124 can't be installed because it doesn't have a source distribution or wheel for the current platform

3. Add remainder of dependencies as required (jupyterlab, notebook, pandas, matplotlib, etc)...

This is what worked for me today. Verified with

Python 3.12.6 (tags/v3.12.6:a4a2d2b, Sep  6 2024, 20:11:23) [MSC v.1940 64 bit (AMD64)] on win32
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.is_available()
True
>>> exit()

[charliermarsh]

https://discuss.python.org/t/selecting-variant-wheels-according-to-a-semi-static-specification/53446/111 - my latest proposal, which has side-tracked me to a PEP for index priority. @charliermarsh you may be interested in this. It is deeply influenced by uv's design.

@msarahan -- That's really interesting. Is there anything we can do on our side to support the PEP, or the variant-selection proposal?

[ngbrown]

@vajranaga:

2. Add torch to the project dependencies with > uv add torch torchvision torchaudio.

Doing this without any index specifier doesn't appear to actually capture the CUDA related sources into the uv.lock file. So this would tell me that it's not reproducible for the next time the environment is setup.

Meanwhile, if I use uv add torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 as that step, I get the following error:

  × No solution found when resolving dependencies:
  ╰─▶ Because there is no version of torch==2.4.1 and torchaudio==2.4.1+cu121 depends on torch==2.4.1, we can conclude
      that torchaudio==2.4.1+cu121 cannot be used.
      And because only the following versions of torchaudio are available:
          torchaudio<2.4.1
          torchaudio==2.4.1+cu121
      and your project depends on torchaudio>=2.4.1, we can conclude that your project's requirements are
      unsatisfiable.
  help: If you want to add the package regardless of the failed resolution, provide the `--frozen` flag to skip
        locking and syncing.

Even though on the previous step I had used the cu121 path on the uv pip install step and the output included:

 + torch==2.4.1+cu121
 + torchaudio==2.4.1+cu121
 + torchvision==0.19.1+cu121

The following works and it produces what looks like would be a working uv.lock file:

uv add torch==2.4.1+cu121 torchaudio==2.4.1+cu121 torchvision==0.19.1+cu121 --index-url https://download.pytorch.org/whl/cu121

Infact, it looks like it's the only thing that is needed:

> py -3.12 -m venv .\.venv\
> .\.venv\Scripts\activate.ps1
> pip install uv
> uv add torch==2.4.1+cu121 torchaudio==2.4.1+cu121 torchvision==0.19.1+cu121 --index-url https://download.pytorch.org/whl/cu121

Then delete the .venv and do the process over, but use uv sync and it works fine pulling source urls from the uv.lock file.

[ashuguptahere]

Is there a way to add torch, torchaudio and torchvision (with gpu and without gpu) in single pyproject.toml file like we can do in poetry for example:

[tool.poetry.dependencies]
python = "^3.12"

[tool.poetry.group.cuda]
optional = true

[tool.poetry.group.cuda.dependencies]
torch = { version = "^2.4.1", source = "pytorch-cuda", markers = "extra=='cuda' and extra!='cpu'" }
torchaudio = { version = "^2.4.1", source = "pytorch-cuda", markers = "extra=='cuda' and extra!='cpu'" }
torchvision = { version = "^0.19.1", source = "pytorch-cuda", markers = "extra=='cuda' and extra!='cpu'" }

[tool.poetry.extras]
cpu = ["torch", "torchvision"]
cuda = ["torch", "torchvision"]

[[tool.poetry.source]]
name = "pytorch-cuda"
priority = "explicit"
url = "https://download.pytorch.org/whl/cu118"

[[tool.poetry.source]]
name = "pytorch-cpu"
priority = "explicit"
url = "https://download.pytorch.org/whl/cpu"

On a side note, after running below uv command, I am unable to add any other python package like, numpy, ultralytics etc.

uv add torch==2.4.1+cu121 torchaudio==2.4.1+cu121 torchvision==0.19.1+cu121 --index-url https://download.pytorch.org/whl/cu121

error:

  × No solution found when resolving dependencies for split (python_full_version == '3.8.*' and platform_machine == 'arm64' and platform_system == 'Darwin' and sys_platform ==
  │ 'win32'):
  ╰─▶ Because there is no version of torch==2.4.1+cu121 and your project depends on torch==2.4.1+cu121, we can conclude that your project's requirements are unsatisfiable.
  help: If you want to add the package regardless of the failed resolution, provide the `--frozen` flag to skip locking and syncing.

[charliermarsh]

No, we don't support resolutions with conflicting extra groups right now, though we're adding first-class support for it in the future.

[YoniChechik]

I've tried to do a frash project with uv init and then adding torch + torchvision with cuda. tried many of the examples here but same errors appeared. Can someone help wit hthe basic commands?

p.s.:
torch is such a popular package. I suggest that the docs will explicitly instruct new users how to install it.

[zanieb]

@YoniChechik see #6523

[vajranaga]

I found that the instructions here:
https://docs.astral.sh/uv/concepts/dependencies/#index

Worked for my rye (with uv) projects too. I just had to manually add those lines for which cuda/cpu index I was using, then uv sync and everything worked fine.

Specifically the explicit = true in the [[tools.uv.index]] block.

[extrange]

EDIT: Just realized my uv.lock specifies the index URL probably due to previously using it, so what I wrote below is not valid.

Contrary to the original poster, I actually can get Pytorch (with the CUDA 12.4 runtime) working with just uv add torch torchvision torchaudio.

[project]
name = "whisper"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
    "torch>=2.5.1",
    "torchaudio>=2.5.1",
    "torchvision>=0.20.1",
]

$ python -c 'import torch; print(torch.__version__)'
2.5.1+cu124

nvidia-smi for my system:

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.90.07              Driver Version: 550.90.07      CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  Tesla T4                       On  |   00000000:00:1E.0 Off |                    0 |
| N/A   35C    P8              8W /   70W |       1MiB /  15360MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
                                                                                         
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+

[hongjiaherng]

I couldn't make it work with the approach of @extrange. And the solution of @vajranaga is indeed working for me!

I found that the instructions here: https://docs.astral.sh/uv/concepts/dependencies/#index

Worked for my rye (with uv) projects too. I just had to manually add those lines for which cuda/cpu index I was using, then uv sync and everything worked fine.

Specifically the explicit = true in the [[tools.uv.index]] block.

Here's my pyproject.toml. I added these manually then run uv sync.

[project]
name = "example"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
    "torch==2.5.1+cu124",
    "torchaudio==2.5.1+cu124",
    "torchvision==0.20.1+cu124",
]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.uv.sources]
torch = { index = "pytorch-cu124" }
torchvision = { index = "pytorch-cu124" }
torchaudio = { index = "pytorch-cu124" }

[[tool.uv.index]]
name = "pytorch-cu124"
url = "https://download.pytorch.org/whl/cu124"
explicit = true

[YouJiacheng]

I have no idea why this can work.
It's expected on Linux

[zanieb]

@YouJiacheng it looks like the dependency on numpy is optional? What happens if you uv add numpy too?

[YouJiacheng]

@zanieb I knew I can install numpy (and it will be fine), my question is, why uv add torch install torch==2.5.1+cu124 for me.
It's expected on Linux.

[zanieb]

Okay, please explain what you're expecting to see when reporting an issue — I can't just guess it from a screenshot. Please read https://docs.astral.sh/uv/pip/compatibility/#local-version-identifiers

[YouJiacheng]

solved!
this behavior is expected on Linux.