🍣

Tsubame4.0でコンテナ環境を作る

2024/06/13に公開

ディープラーニング

Tsubameミニキャンプが開催され、東工大のTSUBAME4.0を触る機会がありました。

TSUBAMEでコンテナを使ったDeepLearning環境構築方法とTipsを紹介したいと思います。
未検証ですが、他の国内スパコンでも使えると予想しています。

TSUBAME4.0とは

Tsubame4.0は東工大が有するスーパーコンピュータです。Tsubame4.0はNVIDIAの最新GPUであるH100を960枚(4枚/ノード x 240ノード)備えたスーパーコンピュータです。

その他概要は、公式のユーザマニュアルなどご覧ください。

コンテナを使った環境構築

今回は、コンテナを起動し、その中に必要なライブラリを書き込んでいくという方法を採用します。スパコンは実験環境で使うことが多いと思いますので、一旦はこれで十分かと思います。

コンテナは、上記で述べたnvidia謹製のpytorchコンテナを使いたいと思います。

TSUBAMEでは、コンテナの代表であるdockerを使用することができません。その代わりapptainerというコンテナ環境があります。apptainerはオープンソースのコンテナプラットフォームであり、Linux財団のプロジェクトとして進められています。apptainerからもdockerでbuildされたコンテナを使うことは可能です。

以下の手順はTSUBAMEのユーザマニュアルに記載されている通りですが、多少変更を加えています。

まず、コンテナをbuildします。今回は、pytorch_gpu という名前で作成します^[1]。

$ APPTAINER_CACHEDIR=/path/to/cache apptainer \
    build -s /path/to/pytorch_gpu \
    docker://nvcr.io/nvidia/pytorch:24.05

-s はサンドボックスモードというものです。サンドボックスモードにしない場合は、コンテナ内で上書することができません。つまり、後述のshell起動時に -wを付けるとエラーにあなります。
APPTAINER_CACHEDIR は設定をしなかった場合に ${HOME}/.apptainer/cache になります。TSUBAMEでは、${HOME}以下の容量が限られているため、APPTAINER_CACHEDIRはグループディスクなどに設定することをお勧めします。コンテナの保存先(今回は/path/to/nvidia_pytorch)も同様です。

buildが完了したら、マウント用のdirを用意します。

mkdir /path/to/pytorch_gpu/gs /path/to/pytorch_gpu/app

shell起動します。--nvccli はexperimentということで、cli関連が不要な場合は--nvとしてください。 --nvの場合は、nvidia-smiなどがコンテナに含まれません。

$ apptainer shell -f -w -B /gs -B /home -B /apps --nvccli \
  /gs/bs/tge-mc2406/ug03712/container/pytorch_gpu_nv
INFO:    User not listed in /etc/subuid, trying root-mapped namespace
INFO:    Using fakeroot command combined with root-mapped namespace
INFO:    Setting 'NVIDIA_VISIBLE_DEVICES=all' to emulate legacy GPU binding.
WARNING: Skipping mount /etc/localtime [binds]: /etc/localtime doesn't exist in container
15:4: not a valid test operator: 12.3
15:4: not a valid test operator: 12.4
21:4: not a valid test operator: (
21:4: not a valid test operator: 545.23.08
Apptainer>

実行

インタラクティブジョブでの実行

今回は、huggingface transformersのexampleにあるrun_glue.pyを実行したいと思います。
なお、 .bashrc などにHF_HOMEを設定しておき、コンテナに入ったら $source .bashrc を実行することをお勧めします。

まずインタラクティブジョブを実行し計算ノードにはいります。

$ qrsh -g [TSUBAMEグループ] -l [資源タイプ]=[個数] -l h_rt=[経過時間]

計算ノードに入ったら、apptainerコンテナを起動します。

$ apptainer shell -f -w -B /gs -B /home -B /apps --nvccli \
  /gs/bs/tge-mc2406/ug03712/container/pytorch_gpu_nv

次にhuggingface transformersのexampleを実行する環境設定をします。

$ pip install transformers
$ git clone https://github.com/huggingface/transformers.git
$ cd transformers/examples/pytorch/text-classification
$ git checkout tags/vXX.XX.XX
$ pip install -r requirements.txt

git checkout するのは、transformersのバージョンとexampleのバージョンを合わせるためです^[2]。

そして、実行します。

$ export TASK_NAME=mrpc
$ python run_glue.py \
  --model_name_or_path google-bert/bert-base-cased \
  --task_name $TASK_NAME \
  --do_train \
  --do_eval \
  --max_seq_length 128 \
  --per_device_train_batch_size 32 \
  --learning_rate 2e-5 \
  --num_train_epochs 3 \
  --output_dir /path/to/your/gs/$TASK_NAME/

最後まで実行されると以下のように表示されます。

...
***** train metrics *****
  epoch                    =        3.0
  total_flos               =   674108GF
  train_loss               =     0.4134
  train_runtime            = 0:00:24.24
  train_samples            =       3668
  train_samples_per_second =    453.797
  train_steps_per_second   =     14.228
06/13/2024 06:27:15 - INFO - __main__ - *** Evaluate ***
[INFO|trainer.py:804] 2024-06-13 06:27:15,160 >> The following columns in the evaluation set don't have a corresponding argument in `BertForSequenceClassification.forward` and have been ignored: idx, sentence1, sentence2. If idx, sentence1, sentence2 are not expected by `BertForSequenceClassification.forward`,  you can safely ignore this message.
[INFO|trainer.py:3719] 2024-06-13 06:27:15,161 >> ***** Running Evaluation *****
[INFO|trainer.py:3721] 2024-06-13 06:27:15,162 >>   Num examples = 408
[INFO|trainer.py:3724] 2024-06-13 06:27:15,162 >>   Batch size = 8
100%|███████████████████████████████████████████████████████████████████████████████████████████████| 51/51 [00:00<00:00, 169.03it/s]
***** eval metrics *****
  epoch                   =        3.0
  eval_accuracy           =     0.8529
  eval_combined_score     =      0.874
  eval_f1                 =     0.8951
  eval_loss               =     0.3659
  eval_runtime            = 0:00:00.30
  eval_samples            =        408
  eval_samples_per_second =   1323.542
  eval_steps_per_second   =    165.443

バッチジョブでのコンテナ実行^[3]

用意したコンテナを用いて、バッチジョブでも実行できます。
バッチジョブでは、シェルファイルを用意して、実行を行います。

まず、実行ファイルを用意します。今回は、job.sh という名前で用意します。ファイル名は任意です。

job.sh

#!/bin/sh
#$ -cwd
#$ -l gpu_h=1
#$ -l h_rt=0:30:00
#$ -o logs/
#$ -e elogs/


GROUP_DISK=/gs/bs/tge-mc2406/ug03712/

export TASK_NAME=mrpc

nvidia-smi --query-gpu=timestamp,name,gpu_bus_id,utilization.gpu,utilization.memory,memory.used,memory.free --format=csv -l 1 > ./gpu-usage.log &

apptainer run --nv -w -B /gs -B /apps -B /home ${GROUP_DISK}/container/pytorch_gpu/ \
  python run_glue.py \
    --model_name_or_path google-bert/bert-base-cased \
    --task_name $TASK_NAME \
    --do_train \
    --do_eval \
    --max_seq_length 128 \
    --per_device_train_batch_size 32 \
    --learning_rate 2e-5 \
    --num_train_epochs 3 \
    --output_dir /local/${JOB_ID}/$TASK_NAME/

cp -r /local/${JOB_ID}/${TASK_NAME}/ /gs/bs/tge-mc2406/ug03712/$}TASK_NAME}

先ほど作った環境を利用するために、 apprainer run でコンテナを指定して実行しています^[4]。
スクリプトの冒頭にある、#$ のブロックは後述する qsub コマンドのオプション指定です。
上記のnvidia-smiと書かれている行は、gpuが使われているかどうかを確認するために入れています。gpu使用率が、 gpu-usage.log に出力されます^[5]。
- gpuのパフォーマンスが下がるため、通常通り実行したい場合は、この行は不要です。
最終行の cpは、上記スクリプトがローカルスクラッチ領域^[6]という実行中のみ使用可能な領域を指定しているため、出来上がったリソースを保存するための操作です。

次に job.sh を先ほどのtext-classification以下に移動します。

$ mv job.sh /path/to/transformers/example/pytorch/text-classification/

そして、実行します。

$ qsub -g <your_group_name> job.sh

実行すると、 ./log/job.sh.o<job_id> に標準出力が、./elog/job.sh.e<job_id>に標準エラーが出力されます。標準先は、job.sh内の $# -o を標準エラー先は、 $# -e を変えると変更できます。

コンテナを使うメリット

最後に、私が思いつくコンテナを使うメリットについて挙げたいと思います。

すでにいろいろインストール済みの環境を使えるので、インストールの手間が省ける。
他の場所で仮想環境を構築することが可能。
スパコン環境でも、自由にライブラリをインストールできる（たまにjavaがないなどあります）。
コンテナ内にインストールすることで、ログインノードの容量逼迫に悩む必要がなくなる（意外と大事）。

1. は、Tsubameに用意されているmoduleを読み込む場合で十分な場合があります。
1. については、例えばnvidia謹製のpytorchコンテナを使うことが可能です。nvidiaのエンジニアがチューニングをしているとのことで、より高速な実行環境を得られます^[7]。なお、flash-attentionもこのコンテナを使用することを推奨しています。また、近年の国内スーパーコンピュータは、同じようなソフトウェアアセットを用意しているため、ほかのスパコンでも同じ環境が使えるメリットがありえます。（未検証）
1. については、グループディスクにサンドボックスモードでコンテナをbuildすることで、インストールしたライブラリをグループディスクに置くことができます。もちろん、pythonの仮想環境をグループディスクにおいてもよいのですが、それよりやや楽という印象があります。

まとめ

apptaienrコンテナを作成して、ライブラリをインストールし、何かしらを実行する一連の手順を紹介しました。nvidiaのライブラリ関連は整合性を含めて注意する必要があるので、コンテナ環境でできると楽だと思います。

最後に、このような機会を提供してくださった、運営の方々に感謝いたします。

脚注

公式ドキュメントでは、サンドボックスモードでbuildする場合も --nv と -B を指定していましたが、build時は指定しなくても、gpuが動作しました。 ↩︎
もちろん、cloneしたtransformersをインストールしてもよいです。 ↩︎
この節は通した動作は時間切れで確認できていません。インタラクティブジョブで、apptaiern runが動作することは確認しています。 ↩︎
はじめ、 -w を付けていなかったため、エラーも出さず、まったく動作しないという事態になりました。解決方法はレトリバの方に教えていただきました。 ↩︎
こちらの資料に掲載されています。 ↩︎
ローカルスクラッチ領域はSSDなので、高速なディスクアクセスが必要な場合に活用できます。 ↩︎
nvidiaのコンテナを使うことで出会ったバグについて公式がサポートしれくれる可能性があります。 ↩︎

TSUBAME4.0とは

コンテナを使った環境構築

実行

インタラクティブジョブでの実行

バッチジョブでのコンテナ実行[3]

コンテナを使うメリット

まとめ

Discussion

バッチジョブでのコンテナ実行^[3]