uvでPyTorchをインストールする方法の試行

もりりん

2024/08/21にリリースされたuvのv0.3.0で、Pythonのバージョン管理もできるようになったので、poetryやryeで組んであるプロジェクトを全てuvに移行しようと思った。

もりりん

とりあえず以下の環境で動いてほしい

x86_64のLinuxマシン
arm64のmacOSマシン

もりりん

新しいプロジェクトの作成

mkdir new_uv
cd new_uv
uv init
uv python pin 3.11
uv run python -V

Python 3.11.9

とりあえず何も考えずに入れてみる。

uv add torch
uv run python -c "import torch; print(torch.__version__);"
# x86_64 Linux
uv run python -c "import torch; print(torch.__version__);"
/home/mjun/workspace/playground/new_uv/.venv/lib/python3.11/site-packages/torch/_subclasses/functional_tensor.py:258: UserWarning: Failed to initialize NumPy: No module named 'numpy' (Triggered internally at ../torch/csrc/utils/tensor_numpy.cpp:84.)
  cpu = _conversion_method_template(device=torch.device("cpu"))
2.4.0+cu121

# arm64
/Users/mjun/workspace/playground/new_uv/.venv/lib/python3.11/site-packages/torch/_subclasses/functional_tensor.py:258: UserWarning: Failed to initialize NumPy: No module named 'numpy' (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/torch/csrc/utils/tensor_numpy.cpp:84.)
  cpu = _conversion_method_template(device=torch.device("cpu"))
2.4.0

numpyが入ってないみたい。

uv add 'numpy<2.0.0'

を実行。

もりりん

uv.lockの中身を見てみる。

[[package]]
name = "torch"
version = "2.4.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
    { name = "filelock" },
    { name = "fsspec" },
    { name = "jinja2" },
    { name = "networkx" },
    { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
    { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
    { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
    { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
    { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
    { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
    { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
    { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
    { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
    { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
    { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
    { name = "sympy" },
    { name = "triton", marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" },
    { name = "typing-extensions" },
]
wheels = [
    { url = "https://files.pythonhosted.org/packages/80/83/9b7681e41e59adb6c2b042f7e8eb716515665a6eed3dda4215c6b3385b90/torch-2.4.0-cp311-cp311-manylinux1_x86_64.whl", hash = "sha256:e743adadd8c8152bb8373543964551a7cb7cc20ba898dc8f9c0cdbe47c283de0", size = 797262052 },
    { url = "https://files.pythonhosted.org/packages/84/fa/2b510a02809ddd70aed821bc2328c4effd206503df38a1328c9f1f957813/torch-2.4.0-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:7334325c0292cbd5c2eac085f449bf57d3690932eac37027e193ba775703c9e6", size = 89850473 },
    { url = "https://files.pythonhosted.org/packages/18/cf/f69dff972a748e08e1bf602ef94ea5c6d4dd2f41cea22c8ad67a607d8b41/torch-2.4.0-cp311-cp311-win_amd64.whl", hash = "sha256:97730014da4c57ffacb3c09298c6ce05400606e890bd7a05008d13dd086e46b1", size = 197860580 },
    { url = "https://files.pythonhosted.org/packages/b7/d0/5e8f96d83889e77b478b90e7d8d24a5fc14c5c9350c6b93d071f45f39096/torch-2.4.0-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:f169b4ea6dc93b3a33319611fcc47dc1406e4dd539844dcbd2dec4c1b96e166d", size = 62144370 },
    { url = "https://files.pythonhosted.org/packages/bf/55/b6c74df4695f94a9c3505021bc2bd662e271d028d055b3b2529f3442a3bd/torch-2.4.0-cp312-cp312-manylinux1_x86_64.whl", hash = "sha256:997084a0f9784d2a89095a6dc67c7925e21bf25dea0b3d069b41195016ccfcbb", size = 797168571 },
    { url = "https://files.pythonhosted.org/packages/9a/5d/327fb72044c22d68a826643abf2e220db3d7f6005a41a6b167af1ffbc708/torch-2.4.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:bc3988e8b36d1e8b998d143255d9408d8c75da4ab6dd0dcfd23b623dfb0f0f57", size = 89746726 },
    { url = "https://files.pythonhosted.org/packages/dc/95/a14dd84ce65e5ce176176393a80b2f74864ee134a31f590140456a4c0959/torch-2.4.0-cp312-cp312-win_amd64.whl", hash = "sha256:3374128bbf7e62cdaed6c237bfd39809fbcfaa576bee91e904706840c3f2195c", size = 197807123 },
    { url = "https://files.pythonhosted.org/packages/c7/87/489ebb234e75760e06fa4789fa6d4e13c125beefa1483ce35c9e43dcd395/torch-2.4.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:91aaf00bfe1ffa44dc5b52809d9a95129fca10212eca3ac26420eb11727c6288", size = 62123112 },
]

uvはRyeと違ってuniversalな依存関係の解決に対応しているので、異なるplatformのwheelのurlも書いてあるみたい。

もりりん

次は、torchのバージョンを切り替えてみる。uvが管理するPythonにパッケージをインストールする方法は2つある。
uv add [package]
uv pip install [package]
uvでパッケージ管理するなら上を、手動でパッケージを入れるには下を使う。flash-attenとかは下じゃないと入らない気がする。(uv pip install flash-attn --no-build-isolation)

もりりん

以下を実行してみる。

uv add torch --index-url https://download.pytorch.org/whl/cpu
uv run python -c "import torch; print(torch.__version__);"
# x86_64
2.4.0+cpu
# arm64
uv add torch --index-url https://download.pytorch.org/whl/cpu
Resolved 11 packages in 2.63s
error: distribution torch==2.4.0+cpu @ registry+https://download.pytorch.org/whl/cpu can't be installed because it doesn't have a source distribution or wheel for the current platform

macだとインストールできない。この時のx86_64側のpyproject.tomlはこんな感じ。

[project]
name = "new-uv"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
    "torch>=2.4.0",
    "numpy<2.0.0",
]

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

もりりん

が関係してそう

もりりん

まあ目標はmacOS, Linuxで使いまわせるpyproject.toml, uv.lockを作ることなので、気を取り直してpyproject.tomlを編集する。

これを見ると、[tool.uv]にfind-linksかextra-index-urlを追加すればよさそう。
find-linksにしてみる。

[project]
name = "new-uv"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
    "torch==2.4.0",
    "numpy<2.0.0",
]

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

[tool.uv]
find-links = ["https://download.pytorch.org/whl/cu118/torch"]

uv lock
uv sync
uv run python -c "import torch; print(torch.__version__);"
# x86_64
2.4.0+cpu
# arm64
uv sync
Resolved 23 packages in 667ms
error: distribution torch==2.4.0+cu118 @ registry+https://download.pytorch.org/whl/cu118 can't be installed because it doesn't have a source distribution or wheel for the current platform

macOSでダメなのはいいとして、x86はバージョン変わらないな......

もりりん

--reinstallつけてやってみる。

# x86_64
uv sync --reinstall
uv run python -c "import torch; print(torch.__version__);"
2.4.0+cu121

cu118が入ってこないで、cu121が来た。
local identifer(2.4.0+cu121の+cu121の部分)をちゃんと入れてみる。

もりりん

それなら直接バージョンを書いてしまうことにする。

[project]
name = "new-uv"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
    "torch==2.4.0+cu118",
    "numpy<2.0.0",
]

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

[tool.uv]
extra-index-url = ["https://download.pytorch.org/whl/cu118"]
# or
# extra-index-url = ["https://download.pytorch.org/whl/cu118/torch"]
# or
# find-links = ["https://download.pytorch.org/whl/cu118/torch"]

uv sync --reinstall
uv run python -c "import torch; print(torch.__version__);"
2.4.0+cu118

ちゃんと入った！

もりりん

あとはこれをmacで実行してみる

もりりん

uvではpoetryのようなPEP508に基づくmarkerが使える。これでインストールするPyTorchを切り替えてみる。

https://peps.python.org/pep-0508/

以下のようにpyproject.tomlを書いてみる。

[project]
name = "new-uv"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
    "torch==2.4.0+cu118; sys_platform == 'linux'",
    "torch==2.4.0; sys_platform == 'darwin'",
    "numpy<2.0.0",
]

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

[tool.uv]
# extra-index-url = ["https://download.pytorch.org/whl/cu118/torch"]
find-links = ["https://download.pytorch.org/whl/cu118/torch"]

uv sync
uv run python -c "import torch; print(torch.__version__);"
# x86_64
2.4.0+cu118
# arm64
2.4.0

いけた！！！

もりりん

ちなみにこれだとダメなので注意。

[tool.uv]
extra-index-url = ["https://download.pytorch.org/whl/cu118"]

もりりん

同じようにtorchvisionをextra-index-urlにしてやってみた

[project]
name = "new-uv"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
    "torch==2.4.0+cu118; sys_platform == 'linux'",
    "torch==2.4.0; sys_platform == 'darwin'",
    "torchvision==0.19.0+cu118; sys_platform == 'linux'",
    "torchvision==0.19.0; sys_platform == 'darwin'",
    "numpy<2.0.0",
]

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

[tool.uv]
extra-index-url = [
    "https://download.pytorch.org/whl/cu118/torch",
    "https://download.pytorch.org/whl/cu118/torchvision",
]

uv sync --reinstall
  × No solution found when resolving dependencies for split (sys_platform == 'linux'):
  ╰─▶ Because there is no version of torch{sys_platform == 'linux'}==2.4.0+cu118 and your project depends on torch{sys_platform == 'linux'}==2.4.0+cu118, we can conclude that your
      project's requirements are unsatisfiable.

解決できない。

もりりん

find-linksにして再トライ

[project]
name = "new-uv"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
    "torch==2.4.0+cu118; sys_platform == 'linux'",
    "torch==2.4.0; sys_platform == 'darwin'",
    "torchvision==0.19.0+cu118; sys_platform == 'linux'",
    "torchvision==0.19.0; sys_platform == 'darwin'",
    "numpy<2.0.0",
]

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

[tool.uv]
# extra-index-url = [
#     "https://download.pytorch.org/whl/cu118/torch",
#     "https://download.pytorch.org/whl/cu118/torchvision",
# ]
find-links = [
    "https://download.pytorch.org/whl/cu118/torch",
    "https://download.pytorch.org/whl/cu118/torchvision",
]

uv sync --reinstall

これはいけた。

もりりん

最終的なpyproject.tomlはこれ

[project]
name = "new-uv"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
    "torch==2.4.0+cu118; sys_platform == 'linux' and platform_machine == 'x86_64'",
    "torch==2.4.0; sys_platform == 'darwin' or (sys_platform == 'linux' and platform_machine == 'aarch64')",
    "torchvision==0.19.0+cu118; sys_platform == 'linux'",
    "torchvision==0.19.0; sys_platform == 'darwin' or (sys_platform == 'linux' and platform_machine == 'aarch64')",
    "numpy<2.0.0",
]

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

[tool.uv]
find-links = [
    "https://download.pytorch.org/whl/cu118/torch",
    "https://download.pytorch.org/whl/cu118/torchvision",
]

もりりん

uv v0.4.23以降で、複数のindex-urlを指定できるようになったので、その方法に切り替えました。
上記までの方法だと、CUDA 12.4を指定してmacOSでPyTorchが入りませんでした。

[project]
name = "new-uv"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
    "torch==2.5.0+cu124; sys_platform == 'linux' and platform_machine == 'x86_64'",
    "torch==2.5.0; sys_platform == 'darwin' or (sys_platform == 'linux' and platform_machine == 'aarch64')",
]

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

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

[tool.uv.sources]
torch = [
    { index = "torch-cuda", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'"},
    { index = "torch-cpu", marker = "sys_platform == 'darwin' or (sys_platform == 'linux' and platform_machine == 'aarch64')"},
]

もりりん

このスクラップは4ヶ月前にクローズされました