ChainerMNとAWS EKSを使ってクラウド上で分散学習

ChainerMNとAWS EKSを使ってクラウド上で分散学習:

Chainer/CuPy 2018の23日目です。

はじめに

ようやく・・・ようやく書き始められる！という気持ちでいます。既に12/25を超えて大分年末ですが、やっとこ書いていきます。遅れてすみません。。本日は、ChainerMNを使って、クラウド上で分散学習をしてみようという記事です。実験するにあたり、札束を募集していたのですが、全く集まりませんでした。仕方ないので株で儲けようと思ったら、日経平均が１日で1000円も下がるしで大変でした。

さて、ChainerMNとは、Chainerで分散学習するための機能です。しかし、実際の所、分散学習ができる環境なんてあんまりないよねー、と多分そこまで積極的に使われていないのではないでしょうか。AWS上に多数のインスタンスを立てて、ChainerMNを動かすという取り組みもありますが、ぶっちゃけ準備と片付けが面倒なので、なかなか普段から使おうとはならないと思います。

一方、最近のクラウド関連の技術の発展は凄いもので、Kubernetes（k8s）という、簡単にデプロイしたシステムをスケールしてくれる技術があります。また、その上で動く機械学習向けの環境を用意してくれるKubeflowという奴がいたりと、分散学習の環境の準備を容易化してくれそうな雰囲気を出しています。さらに、AWSのサービスとしてk8sを動かすAmazon EKSという奴も最近リリースされており、これならAWS上でサクッと分散学習できるんじゃね？って事で色々試してみました。

さらに、EKSでの分散学習に加えて、EFS上にデータを乗せてみようという取り組みもやってみます。AWS EFSは、一言で言えばファイル共有のシステムです。分散学習する際に、学習データをどこに用意するかは結構難しい問題です。数十台のインスタンスに学習データをコピーしても良いけど、流石にちょっと大変だよね・・・ってなりますよね。EFSを使うと、数十〜数百台の大規模な分散環境から利用できたり、容量が無制限である、ボリュームをマウントして使えるなど、ご利益が沢山。という事で、今回はこの機能も使ってみましょう。

まとめると、

• 学習データはAmazon EFSに入れて共有
• 分散環境はAmazon EKSで自動で作成
• ChainerMNを使って簡単に分散学習
• インスタンスを変えた場合と、ノード数を変えた場合で実験

をやっていきます！本記事では、まずツール類の準備、続いて学習コード・データの準備をし、分散環境の構築と実験について述べます。機械学習では普段使わない要素も結構あるので、一番最後に用語をまとめました。

ツール類の準備

インストールする必要のあるツール群は以下になりますので、簡単に入れ方を説明します。

• Docker
• AWS CLI
• eksctl
• kubectl
• ksonnet

Docker

これがないと始まりません。Dockerのインストールは他にも沢山記事があると思うので、ここでは省略。

AWS CLI

これがないと始まりません、その２。AWSのコマンドを叩くためのツールです。これも、沢山記事があるので、ここでは省略。なお、クレデンシャル（AWSにアクセスするためのAPIキーなど）を設定する必要がありますので、これも設定しておきましょう。クレデンシャルの発行は、AWSのコンソール画面からIAMの設定に進み、ユーザーを追加します。AdministratorAccessのアクセス権限をつけてあげましょう。この辺参照。で、作成したユーザの画面から、認証情報→アクセスキーの作成でクレデンシャルを発行しましょう。この辺参照

eksctl, kubectl, ksonnet

eksctlはクラスタを作成するためのツールです。クラスタを作成するためには、AWSコンソール内のAmazon EKSをポチッとするだけでは作ることはできず、VPCやらAuto ScalingやらSecurity Groupやらの設定をしてあげなくてはなりません。eksctlは、この辺の作業を一発でやってくれます。kubectlはk8sをいじるためのツールで、ksonnetはkubeflowを使うためのツールです。

これらは特に設定等はないので、インストール方法のみ書きます。

$ brew install weaveworks/tap/eksctl 
$ brew install kubectl 
$ brew install ksonnet/tap/ks 

Deep Learning編

クラスターを作ってすぐ動かせるように、まずは先にDeep Learningするために必要な諸々（データやコード）を先にやってしまいましょう。本来ならば、EC2のCPUモードx2とかで動作を確認しながら開発を行う所ですが、ソースコードとデータは既にあるものとして進めます。今回は、SSD:Single Shot Multibox Detectorを対象に実験してみましょう。SSDはChainerCVにサンプルがあるので、それを流用します。

データの準備

SSDではPascalVOC2007と2012のデータセットを使って学習します。まずは、このデータをEFSに保存してあげましょう。

EFSにデータを保存するためには、簡単なコピー機能みたいなのはないので、EC2のインスタンスを立ててEFSを接続して、データを準備していきましょう。EFSの作成と、EFSの接続の方法は

• AWS EFS をEC2にマウントしてみる
• Amazon Elastic File System（Amazon EFS）に触ってみた

を参考にしてください。データのコピーに使うだけなのでEC2はどのインスタンスでも構いません。ハマりポイントとしてはセキュリティグループの設定だと思います。EFSの各ゾーンにアタッチしたセキュリティグループの設定で、タイプNFS、ソースをEC2のセキュリティグループにします（上記リンクの２つ目）。

sudo yum install -y amazon-efs-utils 
sudo mkdir efs 
sudo mount -t efs fs-3c8f3194:/ efs 

続いて、PascalVOCのデータをEFS上に保存しましょう。ここでは、~/efs/にマウントされているものとします。

$ cd efs 
$ wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar 
$ wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar 
$ wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar 
$ tar xvf VOCtrainval_11-May-2012.tar 
$ tar xvf VOCtrainval_06-Nov-2007.tar 
$ tar xvf VOCtest_06-Nov-2007.tar 

とすると、~/efs/VOCdevkitにPascalVOCのデータが保存されたと思います。保存されたら、EC2はもう使わないので落としておきましょう。

ソースコードの準備

学習コードはChainercvのssdのサンプルを少しだけ改変して使うことにしましょう。以下のtrain_multi.pyとtrain.pyを使います。

https://github.com/chainer/chainercv/tree/master/examples/ssd 

さらに、マウントしたEFSからデータを読み込むようにするため、

- VOCBboxDataset(year='2007', split='trainval') 
- VOCBboxDataset(year='2012', split='trainval') 
+ VOCBboxDataset(data_dir='/mnt/VOCdevkit/VOC2007', year='2007', split='trainval') 
+ VOCBboxDataset(data_dir='/mnt/VOCdevkit/VOC2012', year='2012', split='trainval') 
-- 
- VOCBboxDataset(year='2007', split='test') 
+ VOCBboxDataset(data_dir='/mnt/VOCdevkit/VOC2007', year='2007', split='test') 

のように修正します。

Dockerイメージの準備

EKSでの学習はDockerの中で行われるため、分散に必要なライブラリ群（OpenMPIやNCCL）を揃えたDocker Imageが必要です。以下のように作りました。ハマりポイントとしては、NCCLをaptでインストールするとCUDA10.0バージョンがインストールされてしまいます。かと言って、CUDA10.0のベースイメージを使うと、EKS側のドライバの問題に引っかかりました。色々考えるのが面倒だったので、もう全部ビルドしてやります。やはり何事も自分でビルドするのが正義（ただしOpenCVはpipで入れた）。

FROM nvidia/cuda:9.2-cudnn7-devel 
 
ARG OPENMPI_VERSION="3.1.3" 
ARG NCCL_VERSION="v2.3.7-1" 
 
# Install basic dependencies and locales 
RUN apt-get update && apt-get install -yq --no-install-recommends --allow-change-held-packages \ 
      locales wget sudo ca-certificates ssh build-essential devscripts debhelper git \ 
      openssh-client openssh-server \ 
      devscripts debhelper fakeroot \ 
      python3-dev \ 
      python3-pip \ 
      python3-wheel \ 
      python3-setuptools && \ 
    rm -rf /var/lib/apt/lists/* /var/cache/apt/archives/* && \ 
    echo "en_US.UTF-8 UTF-8" > /etc/locale.gen && locale-gen 
 
# Install OpenMPI with cuda 
RUN cd /tmp && \ 
  wget -q https://www.open-mpi.org/software/ompi/v${OPENMPI_VERSION%\.*}/downloads/openmpi-$OPENMPI_VERSION.tar.bz2 && \ 
  tar -xjf openmpi-$OPENMPI_VERSION.tar.bz2 && \ 
  cd /tmp/openmpi-$OPENMPI_VERSION && \ 
  ./configure --prefix=/usr --with-cuda && make -j2 && make install && rm -r /tmp/openmpi-$OPENMPI_VERSION* && \ 
  ompi_info --parsable --all | grep -q "mpi_built_with_cuda_support:value:true" 
 
RUN cd /tmp && \ 
  git clone -b $NCCL_VERSION https://github.com/NVIDIA/nccl.git && \ 
  cd nccl && \ 
  make -j src.build && \ 
  make pkg.debian.build && \ 
  dpkg -i build/pkg/deb/libnccl* 
 
# Install ChainerMN 
RUN pip3 install --no-cache-dir mpi4py 
 
# Set default NCCL parameters 
RUN echo NCCL_DEBUG=INFO >> /etc/nccl.conf 
 
# Install OpenSSH for MPI to communicate between containers 
RUN mkdir -p /var/run/sshd 
 
# Allow OpenSSH to talk to containers without asking for confirmation 
RUN cat /etc/ssh/ssh_config | grep -v StrictHostKeyChecking > /etc/ssh/ssh_config.new && \ 
    echo "    StrictHostKeyChecking no" >> /etc/ssh/ssh_config.new && \ 
    mv /etc/ssh/ssh_config.new /etc/ssh/ssh_config 
 
RUN pip3 install chainer==5.1.0 cupy==5.1.0 chainercv==5.1.0 
RUN pip3 install opencv-python chainercv 
 
COPY train_multi.py / 
COPY train.py / 

このイメージをビルドしましょう。

$ docker build -t my-image . 

AWS ECRへの登録

次にECRに登録します。色々場合によって異なりますが、基本は

• 【AWS】初めてのECR

のようにやればOKです。

ここまでで、学習データの作成、学習コードの作成、コンテナの作成ができました。次は、分散環境を構築し、作成したコンテナを流して学習してみましょう。

分散環境の構築編

ということで、分散環境を構築していきます。といっても、決まったコマンドをポチポチ入れていくだけなので、上から順に実行していきましょう。ここでは説明も入れながらコマンドと僕がハマったポイントを書いていきます。コマンドにするとダラダラ長いですが、コマンド書くだけなので、ハマりさえしなければ、とても簡単。ハマりさえしなければだけど。

クラスタの作成

まずは、クラスタを作成しましょう。以下のコマンドで、EKSの作成〜ノードの立ち上げまでやってくれます。ここでは、my-eksという名前、p2.xlargeインスタンスを２個使います。さらにGPUをONにします。

$ eksctl create cluster --name my-eks --nodes 2 --node-type p2.xlarge --region us-east-1 
$ kubectl apply -f https://raw.githubusercontent.com/NVIDIA/k8s-device-plugin/v1.10/nvidia-device-plugin.yml 

無事に立ち上がり、GPUが認識されているか確認してみましょう。

$ kubectl get nodes "-o=custom-columns=NAME:.metadata.name,MEMORY:.status.allocatable.memory,CPU:.status.allocatable.cpu,GPU:.status.allocatable.nvidia\.com/gpu" 

こんな感じで返ってきたら勝ちです。

NAME                                           MEMORY       CPU   GPU 
ip-192-168-26-143.us-east-1.compute.internal   62772568Ki   4     1 
ip-192-168-58-220.us-east-1.compute.internal   62772568Ki   4     1 

ハマりポイント

• ノードを立ち上げられるゾーンを指定しないといけない場合があります。英語ですが頑張って読みましょう。エラーメッセージ中に、利用可能なゾーンが書かれていると思います。例えば以下のようにすると動く場合があります。

$ eksctl create cluster --name gpu-cluster --nodes 2 --node-type p2.xlarge --region us-east-1 --zones=us-east-1a,us-east-1b,us-east-1c,us-east-1d 

• p2.xlargeを始めとしたEKS-GPUインスタンスはデフォルトでは使えない設定になっているらしいです。失敗したらコンソールのエラーメッセージにインスタンスをONにするためのリンクが貼られているので、ポチりましょう。
• CLIにアクセス権限が設定されていないと立ち上がりません。エラーメッセージを読みましょう。
• 失敗したら、AWSのコンソールのCloudFormationを確認して、詳しい状況を見たり、手動削除したりしましょう。CloudFormationが立ち上げた諸々のサービスが残っていると厄介だったりするので、ちゃんとDELETE_COMPLETEされるかどうかを確認すると良いです。
• サクッと書いてますが、ここに進むまでですら、僕は１日くらいかけました・・・。

EFSのアクセス権限の設定

EKSからEFSにアクセスできるようにします。EFSのVPCとゾーンを設定します。コンソール上でEFS内のファイルシステムアクセスの管理を選び、VPCとして作成したEKSのVPCを設定します。VPCの名前はごちゃごちゃしてますが、EKSを作った時に設定した名前が含まれているものを選びましょう。

次に、アベイラビリティゾーンとして、画面に表示されている+ボタンを押して追加します。セキュリティグループとして、デフォルトのものを削除し、現在立ち上がってるEC2インスタンスに設定されているセキュリティグループを設定します。

クラスタにKubeflowをインストール

色々なやり方があるらしいですが、僕は下記のようにしました。kubeflowがgithubにアクセスするらしく、GITHUB_TOKENを設定しなくてはいけないらしいです。TOKENは、ここにやり方が書いてあります。

export GITHUB_TOKEN=XXXXXXXXXXXXXXXXXXXXXXXXXXXXXX 
export KUBEFLOW_VERSION=0.2.5 
curl https://raw.githubusercontent.com/kubeflow/kubeflow/v${KUBEFLOW_VERSION}/scripts/deploy.sh | bash 

これを実行すると、何やら色々インストールされます。以下のようになれば成功です。

$ kubectl get pods 
NAME                                        READY     STATUS    RESTARTS   AGE 
ambassador-59cb5ccd89-bgbpc                 2/2       Running   0          1d 
ambassador-59cb5ccd89-wkf76                 2/2       Running   0          1d 
ambassador-59cb5ccd89-wvdz9                 2/2       Running   0          1d 
centraldashboard-7d7744cccb-g6hcn           1/1       Running   0          1d 
spartakus-volunteer-8bf586df9-xdtqf         1/1       Running   0          1d 
tf-hub-0                                    1/1       Running   0          1d 
tf-job-dashboard-bfc9bc6bc-h5lql            1/1       Running   0          1d 
tf-job-operator-v1alpha2-756cf9cb97-rkdtv   1/1       Running   0          1d 

初期化と名前空間の設定

kubeflowを動かすための各種設定をしましょう。ks initで初期化し、名前空間を作ります。

APP_NAME=kubeflow-chainer-ssd; ks init ${APP_NAME}; cd ${APP_NAME} 
NAMESPACE=chainer-ssd; kubectl create namespace ${NAMESPACE} 
ks env set default --namespace ${NAMESPACE} 

EFSの登録

EFSを登録しましょう。登録には、Persistent Volumeの設定と、Persistent Volume Claimの設定が要るらしいです。dist_pv.yamlを作成します。serverのアドレスは、EFSのコンソールを確認すると書いてありますので、それを使ってください。

dist_pv.yaml

apiVersion: v1 
kind: PersistentVolume 
metadata: 
  name: nfs-data 
  namespace: chainer-ssd 
spec: 
  accessModes: 
    - ReadWriteMany 
  capacity: 
    storage: 85Gi 
  mountOptions: 
    - hard 
    - nfsvers=4.1 
  nfs: 
    path: / 
    server: fs-XXXXXXXX.efs.us-XXXXXXX.amazonaws.com 
  persistentVolumeReclaimPolicy: Retain 
  storageClassName: nfs-external 

dist_pvc.yaml

apiVersion: v1 
kind: PersistentVolumeClaim 
metadata: 
  annotations: 
    volume.beta.kubernetes.io/storage-class: nfs-external 
  name: nfs-external 
  namespace: chainer-ssd 
spec: 
  accessModes: 
    - ReadWriteMany 
  resources: 
    requests: 
      storage: 82Gi 

これらを登録してあげましょう。

kubectl create -f dist_pv.yaml 
kubectl create -f dist_pvc.yaml 

EFSが登録されたかは、下記で確認できます。

$ kubectl get pv -n ${NAMESPACE} 
NAME       CAPACITY   ACCESS MODES   RECLAIM POLICY   STATUS   CLAIM                              STORAGECLASS   REASON   AGE 
nfs-data   85Gi       RWX            Retain           Bound    chainer-ssd/nfs-external   nfs-external            58s 

鍵の登録

クラスタのノード間が通信できるように作成した鍵を送りましょう。

SECRET=openmpi-secret; mkdir -p .tmp; yes | ssh-keygen -N "" -f .tmp/id_rsa 
kubectl create secret generic ${SECRET} -n ${NAMESPACE} --from-file=id_rsa=.tmp/id_rsa --from-file=id_rsa.pub=.tmp/id_rsa.pub --from-file=authorized_keys=.tmp/id_rsa.pub 

クラスタにOpenMPIのインストール

OpenMPIを使えるようにします。

VERSION=master 
ks registry add kubeflow github.com/kubeflow/kubeflow/tree/${VERSION}/kubeflow 
ks pkg install kubeflow/openmpi@${VERSION} 

学習コードの登録

パラメータを設定します。WORKERSはノード数、GPUはノーどあたりのGPU数、IMAGEは学習コードをpushしたリポジトリを指定します。現状では、唯一ここだけは利用者が自身の環境に応じていれてやる必要があります。

$ WORKERS=2 
$ GPU=1 
$ IMAGE=XXXXXXXXX.dkr.ecr.us-east-1.amazonaws.com/XXXXXXXXXXXXXXX 

実行コマンドは下記です。色々設定してますが、正直わからん。

$ COMPONENT=chainer-ssd 
$ EXEC="mpiexec -n ${WORKERS} --hostfile /kubeflow/openmpi/assets/hostfile --allow-run-as-root --display-map --tag-output --timestamp-output -mca btl_tcp_if_exclude lo,docker0 --mca plm_rsh_no_tree_spawn 1 -bind-to none -map-by slot -mca pml ob1 -mca btl ^openib sh -c 'NCCL_SOCKET_IFNAME=eth0 NCCL_MIN_NRINGS=8 NCCL_DEBUG=INFO python3 /train_multi.py'" 

さて、ここまでに設定したパラメータを登録してあげましょう。

$ ks generate openmpi ${COMPONENT} --image ${IMAGE} --secret ${SECRET} --workers ${WORKERS} --gpu ${GPU} --exec "${EXEC}" 

EFSのマウント

もう少しです。最後にEFSのマウント設定をします。

$ ks param set ${COMPONENT} volumes '[{ "name": "efs-pvc", "persistentVolumeClaim": { "claimName": "nfs-external" }}]' 
$ ks param set ${COMPONENT} volumeMounts '[{ "name": "efs-pvc", "mountPath": "/mnt"}]' 

ようやく実行

実行はシンプル、下記のようにやります。

$ ks apply default 

ログの確認

Pod（実行する本体）が立ち上がったかどうかを下記で確認します。

$ kubectl get pod -n ${NAMESPACE} -o wide 

もしRunnningになっていれば、ログを確認して見ましょう。

$ kubectl logs -n ${NAMESPACE} -f ${COMPONENT}-master 

また、Runningにならない場合は、

$ kubectl describe pod -n ${NAMESPACE} 

でメッセージを確認してみましょう。

クラスタの削除

片付けるときは、クラスタを削除しましょう。最初にEFSの管理画面でアベライビティゾーンを全て外して保存、それから少し待って再度管理画面でVPCをEKS以外に変更します。あとは以下のコマンドで構築したクラスタを削除するだけ！

$ eksctl delete cluster --name my-eks 

実験結果

さて、32台くらいで学習しようと思ったんですが、ここで大変なことに気づいた。EC2インスタンスって同時に立ち上げられる台数を制限されているんですね。。。台数の緩和には数日は掛かるとのことだったので断念（追記：制限緩和を申請中！）。

実験は、まずは以下の環境で実施しました。インスタンス毎の違いを確認しました。p3.2xlargeが大分速いです。これは、GPU以外に、ネットワークの速度がp3はp2に比べ大分速いことも効いていると思います。なお、EFSを使わない場合との比較もしたかったのですが、時間が足りず断念。

Instance	# of Instance	iters/sec
p2.xlarge	2	0.3347
p3.2xlarge	2	1.5555

続いて、並列効果を確認するために、2ノードと4ノード、8ノーどで比較してみました。分散の効果がそこまでは出てないですね。おそらくネットワークの速度が遅いからだと思います。高速なネットワークを使えるp3は台数制限の関係で利用できなかったので、宿題にしよう。

Instance	# of Instance	iters/sec
p2.xlarge	2	0.3347
p2.xlarge	4	0.4565
p2.xlarge	8	0.8147

まとめ

というわけで、AWSを使って（簡単？）に分散学習をしてみよう、をやってみました。導入までは中々大変ですし、機械学習では普段触らないようなAWSの諸々やコンテナ技術などがありますが、一通りやってフローにしてしまえば、すごく簡単に分散できると思いました。数字変えれば、そのまま512台でも学習できます。パフォーマンスについては、まだ調べたい所もありますが、その辺は来年に検証していきたいと思います。現状ではEFSのマウントと、GPU数・WORER数・ECRは人出でいれてやる必要がありますが、この辺はスクリプト組んでやれば、ボタン一つで分散もできそうだなと思います。それでは良いお年を。

参考：本記事で用いられる技術集

結構沢山の技術を使うので、都度説明すると読みづらくなるので、簡単に一言だけでも紹介します。

Docker

仮想化のためのツール。簡単に言えば、パソコンの中に、別のパソコン環境を作るみたいなものです。元のパソコンは個人の趣味趣向に従った暖かい環境が整ってますが、他の人の環境でプログラムを動かそうとすると、ライブラリのバージョン違いなどで、サクッとは動いてくれません。上記別のパソコン環境を定義してあげることで、どんな環境でも同じようにプログラムが動作してくれる、という奴です。分散する上でのコアの技術ですね。

MPI

MPIとは分散環境を利用するための標準的な規格。分散環境上でデータをやり取りをする際には、この規格に合った通信用のライブラリ、例えばOpenMPIなど、を使う。

NCCL

NCCLは，NVIDIAが提供しているマルチGPU向けの集合通信用のライブラリ。これをインストールした上でChainer/Cupyをインストールして分散できるようにする。

Kubernetes

コンテナ化したアプリケーションのデプロイ、スケーリング、および管理を行うための、コンテナオーケストレーションシステムである。要は沢山の計算環境で同じプログラムを実行し、管理する仕組みですね。

Kubeflow

k8s上で機械学習をするための仕組み。詳しくは僕も分かっていないけど、例えばMPIがクラスタ内の他のノードのIPアドレスを知りたいみたいな場合に、kubeflowが教えてくれる・・・感じなのかな？

Amazon EKS

Kubernetesクラスタのインストールと運用を自分自身で行うことなく、Kubernetes を AWS で簡単に実行できるためのサービス。

Amazon EFS

AWS上で使用できる共有ファイルストレージです。安全で耐久性があり、スケーラブルな共有ファイルサービスを要求に応じて EC2 インスタンスに提供できる。

Amazon ECR

Dockerコンテナを保存しておく場所。上記にあるように、k8sはDockerコンテナを動かすわけですが、そのために今回はDockerコンテナをECRにおいて置くことにしました。

Amazon IAM Role

AWSの各種サービスの権限管理の仕組み。AWSはセキュリティ上、デフォルトではサービス間の連携はできません。連携には、各サービスに対して、「こっちのサービスにアクセスして良いよ」みたいな権限を与えてあげる必要があります。この仕組みを管理するのがIAM Roleです。

Amazon EC2

簡単に言えばPCみたいなもの。Linuxのインスタンス（１台のPCが１インスタンスと考えれば大体イメージとしては大丈夫）を借りられます。普通のPCと同様に利用できます。

その他使われているAWSサービス

CloudFormation, VPC, AutoScalingなどなど。今回はツールを使ってEKSの立ち上げを行うのですが、その後ろでは様々なAWSサービスが使われています。