ApptainerでPyTorch NGCコンテナを使う
はじめに
NVIDIAから提供されているPyTorch NGCコンテナをApptainerで使用する際のTipsを紹介します。
PyTorch NGCコンテナは、CUDAやcuDNNなどのNVIDIAライブラリが事前にインストールされており、GPUを活用したモデル学習を簡単に行えます。しかし、Dockerコンテナとして提供されているため、HPC環境で使用するにはApptainer向けに変換が必要です。また、追加のライブラリをインストールする手順がDockerを用いる場合と若干異なります。1
本記事では、これらの差異に着目して、ApptainerでのPyTorch NGCコンテナの使用方法を解説します。
前提
HPC環境でApptainerがインストールされていることを前提とします。インストールされていない場合は、管理者に依頼してください。
手順
1. Apptainerをロードする
自動的にApptainerがロードされる場合もありますが、手動でロードする必要がある場合は以下のようにコマンドを実行してください。なお、下記はHPC環境によって異なるため、適宜読み替えてください。
$ module load Apptainer/<version>
2. PyTorch NGCコンテナをダウンロードする
$ apptainer pull docker://nvcr.io/nvidia/pytorch:<version>
ここで、<version>
は使用したいPyTorch NGCのバージョンを指定します。例えば、PyTorch 25.04-py3を使用する場合は以下のようにします。バージョンは、NGCカタログページから確認できます。
$ apptainer pull docker://nvcr.io/nvidia/pytorch:25.04-py3
変換後のファイル名は、pytorch_<version>-py3.sif
の形式になります。例えば、上記のコマンドを実行すると、pytorch_25.04-py3.sif
というファイルが生成されます。
3. コンテナを適切なディレクトリに移動する
ダウンロードしたコンテナファイルは、デフォルトではカレントディレクトリに保存されます。個人の好みに応じて、プロジェクトディレクトリや特定のフォルダに移動することをお勧めします。私の場合は、以下のようにプロジェクトディレクトリ内のcontainers
フォルダに移動しています。
$ mkdir -p $PROJ_HOME/containers
$ mv pytorch_25.04-py3.sif $PROJ_HOME/containers
4. コンテナを実行する
続いて、ダウンロードしたコンテナを実行します。--nv
オプションを指定して、GPUサポートを有効にします。また、--bind
オプションを使用して、プロジェクトディレクトリをコンテナ内にマウントします。これにより、コンテナ内からプロジェクトファイルにアクセスできるようになります。
$ apptainer exec \
--bind $PROJ_HOME:$PROJ_HOME \
--nv $PROJ_HOME/containers/pytorch_25.04-py3.sif \
/bin/bash
HPC環境によっては、ホームディレクトリとLustreベースのファイルストレージを自動的にマウントする設定がされている場合があります。その場合、--bind
オプションは不要です。マウントされているか否かについては、使用するHPCのドキュメントを確認してください。
5. 必要なライブラリを追加インストールする
PyTorch NGCコンテナには、基本的なPyTorch環境が整っていますが、Transformersやその他のライブラリを使用する場合は、追加でインストールが必要です。追加のインストールには仮想環境を使用することをお勧めします。以下は、pip
を使用してTransformersライブラリをインストールする例です。
$ python3 -m venv --system-site-packages $PROJ_HOME/envs/test
$ source $PROJ_HOME/envs/test/bin/activate
$ unset PIP_CONSTRAINT
$ pip install transformers
このコマンドは、プロジェクトディレクトリ内のenvs/test
に仮想環境を作成し、そこにTransformersライブラリをインストールします。--system-site-packages
オプションを使用することで、コンテナ内のシステムパッケージも利用できます。
なお、PIP_CONSTRAINT
環境変数を解除することで、PyTorch NGCコンテナに設定されているバージョンロックを一時的に無視してインストールを行います。この点がDockerコンテナと異なる点です。/etc/pip/constraint.txt
にバージョンロックが設定されていますが、編集できない場合があるため、環境変数を解除する方法を採用しています。
インストールが完了したら、一旦コンテナ内から抜けてください。
$ exit
6. 実際にコンテナを使用する
インタラクティブジョブの場合
インタラクティブジョブを使用してコンテナを実行する場合、上記の4.に示したコマンドをそのまま使用できます。
$ apptainer exec \
--bind $PROJ_HOME:$PROJ_HOME \
--nv $PROJ_HOME/containers/pytorch_25.04-py3.sif \
/bin/bash
なお、コンテナ内でライブラリの環境変数を設定する必要がある場合は、以下を参考にしてください。ただし、パスは環境によっては異なる可能性があります。コンテナが使用するCUDAのバージョンは、公式のサポート表を参照してください。25.04-py3の場合は、CUDA 12.9が使用されているため、以下のように設定します。
export CUDA_HOME=/usr/local/cuda-12.9
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
export TMPDIR=/tmp
ジョブスクリプトの場合
ジョブスクリプトを使用してコンテナを実行する場合、以下のようにapptainer exec
コマンドを使用します。Slurmのジョブスクリプトの例を以下に示します。
#!/bin/bash
#SBATCH --job-name=pytorch_job
#SBATCH --output=pytorch_job.out
#SBATCH --error=pytorch_job.err
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
#SBATCH --gres=gpu:1
#SBATCH --time=01:00:00
#SBATCH --partition=gpu
module load Apptainer/<version>
PROJ_HOME=/path/to/your/project
apptainer exec \
--bind $PROJ_HOME:$PROJ_HOME \
--nv $PROJ_HOME/containers/pytorch_25.04-py3.sif \
python3 $PROJ_HOME/scripts/train.py
コマンドライン引数をたくさん指定したい場合は、シェルスクリプトを実行用(内部)とSlurmジョブスクリプト(ラッパー)用に分けることをおすすめします。以下がその例です。
実行用スクリプト
#!/bin/bash
source $PROJ_HOME/envs/test/bin/activate
export CUDA_HOME=/usr/local/cuda-12.9
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
export TMPDIR=/tmp
python3 $PROJ_HOME/scripts/train.py "$@"
Slurmジョブスクリプト
#!/bin/bash
#SBATCH --job-name=pytorch_job
#SBATCH --output=pytorch_job.out
#SBATCH --error=pytorch_job.err
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=4
#SBATCH --gres=gpu:1
#SBATCH --time=01:00:00
#SBATCH --partition=gpu
module load Apptainer/<version>
PROJ_HOME=/path/to/your/project
apptainer exec \
--bind $PROJ_HOME:$PROJ_HOME \
--nv $PROJ_HOME/containers/pytorch_25.04-py3.sif \
$PROJ_HOME/scripts/run.sh "$@"
このようにすることで、Slurmジョブスクリプトから実行用スクリプトを呼び出し、Pythonスクリプトを実行できます。
まとめ
本記事では、ApptainerでPyTorch NGCコンテナを使用する方法を解説しました。
PyTorch NGCコンテナはCUDA等の必要なライブラリがほとんど揃っているので、手早く実験環境を構築できます。使ったことがない方は、ぜひ試してみてください。
-
なお本記事は、[Tips] NGC PyTorchのversion lockを解除する方法に着想を得て、執筆しました。 ↩