🌊

Apple silicon搭載Mac上のDockerの中でpytorch+cpuが使いたくなった時に自力でwheelをbuildする方法

2023/12/17に公開

Docker

tl;dr

https://github.com/pytorch/pytorch をcloneする
.devcontainer/Dockerfileの一行目FROM mcr.microsoft.com/vscode/devcontainers/miniconda:0-3を自分が欲しいpython versionのものに変更する。( https://github.com/devcontainers/images/tree/main/src/miniconda の historyなどを参照する)
- pythorchの特定のversionを使いたい場合はgit tagが切られているので戻すことができるが、古いcommitには.devcontainerディレクトリが存在しない場合があるので、適宜トラック外にしておく
vscodeのdevcontainer環境から.devcontainer/cpu/devcontainer.jsonを参照する形でdevcontainerを起動する
devcontainer環境でUSE_CUDA=0 USE_CUDNN=0 python setup.py bdist_wheel する
- なお筆者の環境（M1 Pro + 32GB RAM）ではBuildに1時間半ほどかかかる
dist/配下にtorch-1.13.1+git49444c3-cp310-cp310-linux_aarch64.whlのようなファイルが生成される

そもそも何でそんなことをしたいのか

pytorchの公式whlはhttps://download.pytorch.org/whl/torch/にあるが、cpu+linux+aarch64のbuildでものはだいたいのversionで提供されていない。
普段Apple silicon搭載Macで開発をしている際、開発環境をdevcontainer(docker)で構築しているため、linux_aarch64環境になる。
- Apple silicon上で頑張ってx86_64にする諸々の方法はボトルネックがあり遅かったり不安定だったり(QEMU/Use Rosetta for x86/amd64 emulation on Apple Silicon)するので、微妙な気持ちになる。
本番環境はクラウドなのでlinux+x86_64で動く前提でDockerfileが書かれており、imageサイズを削減するためにcpuのみimageを作ることが多く、poetryなどでwhlを指定してcpuのみbuildが採用されていたりする。(下図)

[tool.poetry.dependencies]
torch = {version = "1.13.1", source = "pytorch"}
[[tool.poetry.source]]
name = "pytorch"
url = "https://download.pytorch.org/whl/cpu"
priority = "explicit"

つまり結果として、Apple silicon Mac上のdocker環境でpoetry installなどとすると以下のように該当のwhlが見つからないということでエラーがでる。

  • Installing torch (1.13.1+cpu): Failed

  RuntimeError

  Unable to find installation candidates for torch (1.13.1+cpu)

  at /usr/local/python/3.10.13/lib/python3.10/site-packages/poetry/installation/chooser.py:73 in choose_for
       69│ 
       70│             links.append(link)
       71│ 
       72│         if not links:
    →  73│             raise RuntimeError(f"Unable to find installation candidates for {package}")
       74│ 
       75│         # Get the best link
       76│         chosen = max(links, key=lambda link: self._sort_key(package, link))
       77│ 

Cannot install torch.

要するにdocker上でなければ(macosx_arm64は提供されているのでcpu指定を外す必要はあるが)成功するし、x86_64のdocker上でも成功するので、「linux_aarch64でcpu buildとかそんな例外環境知らんがな」と放置されているのである。

筆者は環境を持っていないが、nativeなlinux_aarch64環境でも同様のエラーが出ると思われる。graviton2などのARMベースのサーバーも増えてきているので、そういう環境でも同様の現象が起こっている可能性がある。

ということで、localでbuildして、そのwhlを指定することで回避することができる。(具体的な操作はtl;drを参照)
なおpoetryで開発環境と本番環境を揃えたいという気持ちはすでに破綻しているが気にしないことにする。

具体的なwhlファイルをpoetryで指定する方法は以下の通り。

[tool.poetry.dependencies]
torch = { path = "/poetry_wheels/torch-1.13.1+git49444c3-cp310-cp310-linux_aarch64.whl"}

ちなみにlinux_aarch64上でも cpu build指定しなければいいじゃんという話があるが、今度は下記のようにCUDA依存があるので、linux_aarch64 on docker on apple silicon macでは詰みである。

 • Installing nvidia-cuda-nvrtc-cu11 (11.7.99): Failed

  RuntimeError

  Unable to find installation candidates for nvidia-cuda-nvrtc-cu11 (11.7.99)

  at /usr/local/python/3.10.13/lib/python3.10/site-packages/poetry/installation/chooser.py:73 in choose_for
       69│ 
       70│             links.append(link)
       71│ 
       72│         if not links:
    →  73│             raise RuntimeError(f"Unable to find installation candidates for {package}")
       74│ 
       75│         # Get the best link
       76│         chosen = max(links, key=lambda link: self._sort_key(package, link))
       77│ 

Cannot install nvidia-cuda-nvrtc-cu11.

補足

上述の方法でbuildしたwhlは、一部shared objectが足りないというエラーが出る。
個別にインストールすれば良いのだが、公式配布されているwhlは起こらないようなので、なんでそうなっているのか不明である。

GitHubで編集を提案

tl;dr

そもそも何でそんなことをしたいのか

補足

Discussion