スパコンで機械学習のすすめ

機械学習をする上でどうしてもネックになってくるのがGPUですよね。 GCPやAWSを使うのも一つの手ですが、研究で使うのであれば一番安く大量のGPUを使うことができる環境はおそらくスパコンです。 とはいえ、なんとなく難しそうといって敬遠されがちな感じがするので、実際はなんも難しくないよという話をします。

スパコンといっても色々ありますが、私が使ったことがあるのがReedbushとABCIなので、この２つに関して話を進めていきますが、他のスパコンもおそらく同様にして使うことができると思います。

スパコン概要

Reedbush

https://www.cc.u-tokyo.ac.jp/supercomputer/reedbush/service/

東大にあるスパコンです。 Reedbush-U, Reedbush-H, Reedbush-Lがあり、UはCPUのみ、Hは１ノードにTesla P100が２つ、Lは１ノードにTesla P100が４つのっています。 大学のスパコンというだけあって、卒論・修論シーズンにはちょっと混雑します。

ABCI

産総研にあるスパコンです。日本語名は「AI橋渡しクラウド」でダサい。１ノードにTesla V100が４つ乗っています。ハイパフォーマンスコンピュータ関係の人だけでなく、機械学習などを専門にする人にもぜひ使って欲しいという意気込みで作られているだけ合って、ドキュメント（https://portal.abci.ai/docs/ja/)もきれいだし、jupyter notebookの使用も想定されていたりして（https://abci.ai/ja/how_to_use/data/HowTo_JupyterNotebook.pdf）かなり使いやすいと思います。 また、ImageNetなどのよく使うデータセットに関しては共有ストレージに予め用意しておこうかという話もあり、これからどんどん使いやすくなっていくと思います。

環境構築

公開鍵の登録

スパコンの利用申請をして申請が通ると、メールや郵便でusernameとPasswordが送られてきます。 そのusenameとpasswordを使ってポータル（Reedbush https://reedbush-www.cc.u-tokyo.ac.jp , ABCI https://portal.ABCI.ai/user/ ）にアクセスし、公開鍵を登録することで、スパコンにsshできるようになります。

sshとスパコンの構成の理解

ABCIの場合

挙動から予想すると、こういう構成になっていると考えられます。 スパコン部分とストレージを共有しているノードにアクセスするためには、一旦アクセスサーバーに接続してトンネリングする必要があります。 具体的なコマンドは下のようになります。

$ ssh -L 10022:es:22 -l username as.abci.ai

# 別のターミナルで
$ ssh -p 10022 -l username localhost

Reedbushの場合

Reedbushの場合はこのようになっていると考えられます。 トンネリングなどは必要なくそのままログインサーバーにアクセスすることができます。 cdwというコマンドで スパコンにマウントされているディレクトリに移動することができます。

$ ssh <username>@reedbush-u.cc.u-tokyo.ac.jp
$ cdw  # /lustre/<groupname>/<username> に移動

そして、この/lustre/<groupname>/<username>/以下がスパコン本体のホームディレクトリとストレージを共有しています。

ライブラリの導入

いくつかのよく使われるライブラリは既にスパコン内にインストールされています。 module availableとすると、その一覧をみることができます。

cudaやcudnnなど必要なパッケージは、ここから入れることができます。 module load cuda9/9.1.85 などとすれば、自分の環境にパスが通りつかえるようになります。 tensorflowやkerasなどの機械学習ライブラリについても、ここから入れることができますが、バージョンが古かったり、Pytorchなどは現状ここから入れることができなかったりするので、python系のライブラリに関しては、全てpyenv環境を用いてインストールするのが良いかなと思います。

pyenvの導入

pyenvを使わなくても、先述したmodule loadを使って、pythonやpythonライブラリの一部は使うことができますが、pyenvを使えば自分のマシンと同様にして自分の好きなライブラリを好きなように使うことができるのでおすすめです。

ReedbushもABCIもあらかじめgitがインストールされているので、通常と同様にgit cloneでpyenvを導入することができます。 一点注意が必要なのは、Reedbushの場合は、 /home/<username>/以下ではなく/lustre/<grupname>/<username>/以下に入れるようにしてください。/home/<username>/に入れてしまうと、スパコン本体から読むことができなくなってしまいます。

pyenvが無事に導入できたら、あとは普通に、好きなpythonのバージョンで使いたいライブラリをインストールしてください。virtualenvなどを使いたければそのへんも自由にセットアップしてください。

実行

動かしたいプログラムについては、スパコン特有のものは何もありませんので割愛します。 普段と同じように作成してください。 そうして作成したプログラムを実行していきます。

ここでもう一度スパコンの構成図を思い出してほしいのですが、アクセスしているのはスパコン本体ではなくログインノードです。 ですから、そのままここで python hoge.py などとして実行しても、スパコン本体ではなくログインノードで実行されるだけです。 もちろん、syntax errorやlmport errorなどの初歩的なエラーで落ちないか確かめるうえでは、ログインノードで試しに走らせて見るというのは意味のあることですが。

スパコン本体で実行するには、インタラクティブジョブとバッチジョブの２種類のやり方があります。 インタラクティブジョブはその名の通りインタラクティブに実行できるので、直感的でわかりやすいです。パラメータなどをいじりながら様子をみながら実行したいなどの場合には、こっちが向いていると思います。しかしながら、一般的にスパコンにやらせたいような規模のおおきなプログラムに関してはバッチジョブで実行するのが向いているので、今回はバッチジョブのやり方を説明します。

バッチジョブを実行するには Reedbushの場合は以下のようなバッチスクリプト(run.sh)をかくことになります。

#! /bin/sh
#PBS -q h-small
#PBS -l select=1:mpiprocs=1:ompthreads=36
#PBS -W group_list=<groupname>
#PBS -l walltime=02:00:00

. /etc/profile.d/modules.sh
module load cuda9

export PYENV_ROOT="$PBS_O_WORKDIR/.pyenv"
export PATH="$PYENV_ROOT/bin:$PATH"
if command -v pyenv 1>/dev/null 2>&1; then
  eval "$(pyenv init -)"
fi
eval "$(pyenv virtualenv-init -)"

cd $PBS_O_WORKDIR/hoge
python fuga.py

こうすると/lustre/<groupname>/<username>/hoge/fuga.pyをスパコンで実行することができます。

まず、２行目の#PBS -q h-smallでジョブクラスを指定しています。初めに述べたようにReedbushにはU, H, Lの３種類のマシンがあります。この例では、”h-small”なのでReedbush-Hを用いたノード数の少ないジョブクラスを指定しています。 その他のジョブクラスについては https://www.cc.u-tokyo.ac.jp/supercomputer/reedbush/service/job.php を参照してください。

３行目の#PBS -l select=1:mpiprocs=1:ompthreads=36は順番に、使うノード数・１ノードあたりのプロセス数・スレッド数を表しています。この例では１ノード１プロセス３６スレッドということです。

#PBS -l walltime=02:00:00は実行時間です。ここで設定した時間を超えてしまうと、そのジョブはkillされてしまいます。ですので、想定される実行時間よりも少し余裕を持って設定したほうがいいです。ちなみにtokenは実際に実行した時間分しか消費されません。 ただ、ここで設定された時間をもとにスケジューリングされるので、必要以上に長い時間を書いてしまうと、なかなかジョブが実行されなくなってしまう場合もあります。なので程よい時間に設定しましょう。

module loadコマンドで何かをインストールした場合はジョブの中でも同様にmodule loadする必要があります。

. /etc/profile.d/modules.sh
module load cuda9

今回はcuda9をインストールしています。

export PYENV_ROOT="$PBS_O_WORKDIR/.pyenv"
export PATH="$PYENV_ROOT/bin:$PATH"
if command -v pyenv 1>/dev/null 2>&1; then
  eval "$(pyenv init -)"
fi
eval "$(pyenv virtualenv-init -)"

pyenv（＋pyenv-virtualenv）環境にパスを通しています。これで、スパコン本体からもpyenvで構築した環境を用いることができるようになります。 ちなみに、$PBS_O_WORKDIRはジョブを投げた際のカレントディレクトリを指すので、このスクリプトは/lustre/<groupname>/<username>/から投げられることを想定しています。他のところから投げたい場合には、適宜パスを変更してください。

そして最後に、

cd $PBS_O_WORKDIR/hoge
python fuga.py

でhogeディレクトリにはいりfuga.pyを実行しています。

ABCIの場合も書き方は少し違いますが、かくことはおおよそ同じなので、例だけおいておきます。

#!/bin/bash
#$-l rt_F=1
#$-g <groupname>
#$-l h_rt=02:00:00
#$-cwd

source /etc/profile.d/modules.sh
module load cuda/9.0
module load cudnn

export PYENV_ROOT="/home/<username>/.pyenv"
export PATH="$PYENV_ROOT/bin:$PATH"
if command -v pyenv 1>/dev/null 2>&1; then
  eval "$(pyenv init -)"
fi
eval "$(pyenv virtualenv-init -)"

python hoge/fuga.py

バッチスクリプトができたら、あとはこれをシステムに投げるだけです。 ReedbushもABCIも qsub run.sh というコマンドでジョブを投げることができます。

投げたジョブの状態を確認するには、 Reedbushでは

rbstat

ABCIでは

qstat

というコマンドを実行します。

実行した結果の、標準出力・標準エラー出力はそれぞれ、 run.sh.o<job id>, run.sh.e<job id> というファイルに書き込まれます。

いかがでしたか？

スパコンを使うのはそれほど難しくないことがわかっていただけたと思います。

今回説明した例では、１ノード、１プロセスしか使っていないので、スパコンの真価を発揮できているとは言えません。なので、また気が向いたら今度はマルチノード・マルチプロセス・マルチGPUの場合についても書こうと思います。