はじめに
Turingの基盤AIチームに業務委託として所属している東京科学大学(Institute of Science Tokyo)の藤井です。
本記事では、PyTorchをsource buildした後に利用する際にCUDA Toolkitをmodule loadしておく必要があることの背景について解説を行います。
また、pip installによりbuild済みのbinaryをインストールしたときのPyTorchとsource buildした場合の違いについても解説を行います。
普段はSwallow Projectや横田研究室にて大規模モデルの分散並列学習や低精度学習について研究を行っていますので、そちらもご覧いただけますと幸いです。
本記事の概要を図にすると以下の通りになります。
pip install時の挙動
まず、source buildしたときとの比較のために、多くの方がよく行うpip install torchによりPyTorchをインストールした場合について解説します。
PyTorch 2.5.1+cu124をインストールした場合どうなっているのか実際に見てみましょう。
pip listにてインストールされている他のpackagesを観察すると以下のようなものが見当たるはずです。
nvidia-cublas-cu12 12.4.5.8
nvidia-cuda-cupti-cu12 12.4.127
nvidia-cuda-nvrtc-cu12 12.4.127
nvidia-cuda-runtime-cu12 12.4.127
nvidia-cudnn-cu12 9.1.0.70
nvidia-cufft-cu12 11.2.1.3
nvidia-curand-cu12 10.3.5.147
nvidia-cusolver-cu12 11.6.1.9
nvidia-cusparse-cu12 12.3.1.170
nvidia-cusparselt-cu12 0.6.2
nvidia-nccl-cu12 2.21.5
nvidia-nvjitlink-cu12 12.4.127
nvidia-nvtx-cu12 12.4.127
これらのpakcagesは、NVIDIA GPUを効率よく利用するために必要なものになります。
CUDA Toolkitや関連するライブラリがインストールされている状態です。
つまり、pip installによりPyTorchをインストールする限りにおいて、ローカルにCUDA Toolkit, cuDNN, nccl, cuBLASなどのライブラリを用意する必要はなく、pip installにより同時にinstallされるpackageに内包されているライブラリをPyTorchが動作する際に利用しています。
この同時にインストールされるライブラリは、PyTorchの依存関係としてPythonパッケージ(.whlファイル)としてPyPIなどからインストールされるものになります。
これらのpackageは、共有ライブラリ(.soファイル)とメタデータをPython環境にインストールし、PyTorchが正常に動作することを助けています。
これらを図示すると以下のようになります。
source build installの挙動
build時
PyTorchをbuildするにはCUDA Toolkitのヘッダファイル(cuda.hなど)が必要であるため、module load cuda/12.9などでLD_LIBRARY_PATHに設定されたパスを参照します。
そのため、build時にはCUDA Toolkit, cuDNN, ncclなど多数のLibraryが参照可能である必要があります。これは、pip installでinstall可能なPyTorchも同様の方法でbuildされているため、変わりません。
実行時
build時に利用していたライブラリへの参照はバイナリに埋め込まれますが、buildされたバイナリの中にそのライブラリ自体は含まれません。そのため、実行時にはライブラリの参照先が必要です。
pip install torchにてインストール可能なpackageの場合は、その参照先となるpackagesが同時にインストールされますが、source buildの場合はそうではありません。ここが決定的な違いとなります。
そのため、実行時にmodule load cuda/12.9 cudnn/9.7.0などを行わないとOSError: libcudart.so.12: cannot open shared object file: No such file or directoryのようにshared objectが見つからないということでエラーが発生します。これは先に説明したように、実行時にbuild済みのPyTorch binaryは、LD_LIBRARY_PATHやPATHから実行時に必要なライブラリを参照するためです。
上記の内容を図示すると以下のようになります。
おわりに
本記事では、pip installによりbinary形式のbuild済みPyTorchを利用するときと、source buildしたPyTorchを利用するときに裏で行われている動的リンクに焦点を当てて解説を行いました。
何気なく利用しているPyTorchの裏側の本当に表層的な部分を紹介したにとどまりますが、多くのリーサーチエンジニアによって支えられるソフトウェアスタックの上に、現在のMLライブラリがあることを少しでも実感いただけましたら幸いです。
PyTorchをsource buildする方法の一例についても紹介していますので、よろしければどうぞ
普段は、LLMの開発や、低精度学習による高速化を中心に研究開発を行っていますので、よろしければそちらもご覧ください。
Discussion