PyTorchビルドからコントリビュートのための注意

(この情報は2025/05/17の情報)

PyTorchをビルドする

人にはPyTorchをビルドして編集しなければいけない時がある, 国際学会投稿のためとか, ボスに命令されたからなどである。

Clone

まずcloneして現在の環境でinstallできることが確認できる(pipとかでinstallしてみる)PyTorchにcheckoutする。自分はenrootコンテナのpytorch:24.10-py3に合わせた。

git clone --recursive https://github.com/pytorch/pytorch 
git checkout e000cf0ad9 # pytorch:24.10-py3用のリポジトリはこれ
git submodule sync
git submodule update --init --recursive
# ビルド前に一旦前のビルド結果が残ってないかを確認する
python setup.py clean
rm -rf build/

Env

環境変数を設定する, 自分の環境ではcondaではなくコンテナ内のPythonを使っていることに注意
あとは諸事情で環境がArm CPUだったので, 推奨設定のUSE_PRIORITIZED_TEXT_FOR_LD=1をつけている。正直ここはPyTorchのGitHubを眺めながら設定して欲しい

export USE_CUDA=1
export CUDA_HOME=/usr/local/cuda
export CMAKE_PREFIX_PATH=$(dirname $(which python))/../
export USE_PRIORITIZED_TEXT_FOR_LD=1  

Build

whlをビルドしてpipでいれる。初回はメモリ80GBで2時間以上かかるが, それ以降は小さな変更ならば10分程度で終わります。build/を消すとまた最初からなので注意

python setup.py bdist_wheel
pip install dist/torch-2.5.0a0+e000cf0-cp310-cp310-linux_aarch64.whl

その他の注意

デバッガやGitHubで検索しつつ変更箇所を特定するようになるとは思いますが, 一部ビルド時に自動生成されるファイルなどがあります。torchgen/gen.pyなどがその自動生成用スクリプトで, おそらくですがbuild/の各デバイスに対する関数はこれで生成されています。
そのため, このようなファイルは一度ビルドした環境でないと追えないのでGitHubとかで追うことは難しいので注意してください。

自動生成されたものは大体こういう感じです。build/aten/src/ATen以下のcppファイルなどが生成されています。

主にデバイスの置き場所に関するassertionを出すcheck_and_update_common_deviceはこの自動生成経由でしか使われていないっぽいので最初に挙動を追う時に苦労しました。