Graviton2 + NVIDIA GPUなインスタンスを使ったECS on EC2における機械学習ワークロードの実行
少し遅くなりましたが、あけましておめでとうございます。
先日弊社のMission/Challenge/Valueについて代表が語るnoteが公開されましたが、僕自身も引き続きひたすらいろいろなことにチャレンジしていきたい所存です。
そして早速、新年のチャレンジとしては面白いことができたのでお話ししたいと思います。
Graviton2 + NVIDIA GPUなインスタンス
どうも、AWS大好き芸人の@ken11です。AWSが好きすぎて毎日目黒経由で通勤してます。嘘です。
G5g
皆さんAWSのG5gインスタンスファミリーはご存じでしょうか?
これ、あまり知られていない気がするんですよね。
「G5ファミリーの亜種でしょ?」
これは正解であり、一方で誤解をはらんでいる可能性も高いと思われます。
G5ファミリーはNVIDIAのA10GというAmpere世代のGPUがついた、いわば「一般的なGPU付きインスタンス」です。
一方、今回ご紹介するG5gというのは、suffixのg
が示すようにAWS Gravitonシリーズ搭載[1]インスタンスというのが最大の特徴で、GPUに関してはNVIDIAのT4GというTuring世代の古いものがついています。
そんなちょっとややこしいインスタンス[2]がG5gです。
G5gを取り巻く状況
ではこのG5gインスタンスというのは誰もが気軽に使える状況でしょうか?
残念ながら2024年1月15日現在、このインスタンスの力を最大限活かして使うことは至難の業になっています。
というのも、このインスタンスファミリーの特徴であるGraviton2というのはCPUアーキテクチャが AArch64
です。
EC2を起動しようとすればすぐわかるのですが、この AArch64
に対応したAMIというのはとても選択肢が少なくなります。
そこで救済となるのがBottlerocketです。これはコンテナホスト用にAWSが開発しているLinuxベースのOSで、ECS用として、bottlerocket-aws-ecs-2-nvidia-aarch64
というAMIが存在します。
じゃあG5gインスタンスファミリーを使うときはこれを指定すればOKだね、めでたしめでたし。
…
……
………
とは残念ながらいかないわけです。(どうして)
issueも開いている状況ですが、現状このBottlerocketではNVIDIA Container Toolkitが正しく動作しません。(--gpus all
などのフラグを付けた瞬間にエラーして詰む)
先ほどもお話ししたようにAArch64ではあまりAMIの選択肢がありません。GPUサポートとなると古いものしか存在せず、2024年のいま使うには躊躇するものばかりです。頼みの綱だったBottlerocketに梯子を外された今、GPUサポートのものは実質的な選択肢がない状況[3][4]です。
この辺でおそらくG5gを(ECSで)使うことを諦める人がほとんどじゃないでしょうか?
はっきり言ってこの状況から先に進むのはコスパが悪いです。
ではなぜそれでも僕がチャレンジするのか?それは登ってみたとき広がる世界は小さいほど気持ちいいと言った人が過去にいたからとかそういうのももちろんあるわけですが、僕自身はGravitonシリーズに対してコスト・性能の両面から強い期待を抱いていて、近い将来にGraviton + InferentiaのようなインスタンスやGraviton + A10G or L4みたいなインスタンスも出てくると信じているからこそ、ここで一度Graviton + NVIDIA GPUの状況をしっかりと理解して活用方法を手の内化しておきたかったというのが理由です。
ECS on EC2 G5gの構築
というわけで、デフォルトで用意されている材料ではG5gを活用することが困難という現状がわかったので、ここからは実際に使えるECS on EC2 G5gを構築していきましょう。ほぼゼロからスタートという感じなので、ここから先は茨の道です。
AMIをつくる
いきなりそれか〜という感じですがそれはそうなんです、ECSで使うのが目的ですし、まずはボタン一つでGPUが使えるG5gインスタンスをつくることが肝心です。そのためにはAMI作成は必須の工程です。
過去にこういう記事もあるのですが、情報が古いので新たにつくることにします。
1. Amazon Linux 2023ではなくAmazon Linux 2を使います
最新の2023(実質Amazon Linux 3)を使いたかったところなんですが、試したところ2023では詰まる部分があったので今回はAmazon Linux 2を使います。
まずAMI作成のベースとなるAMIをAmazon Linux 2の中から選びましょう。ちなみにkernelが4系だとこれまた途中で詰まったので、5系が望ましいと思われます。
またこの際、amazon-ecs-agentが予め入っているものだとECSで使う際に楽なので、その中から選びます。
- Amazon Linux 2
- AArch64(ARM64)
- kernel5系
- amazon-ecs-agent
これを基準に選ぶと、amzn2-ami-ecs-kernel-5.10-hvm-2.0.20231219-arm64-ebs
(ami-03e11de03c83b999c
)などが候補となります。
これは2024年1月15日時点で最新のものです。
今回はこれをベースのAMIとして使います。選択したAMIで、おもむろにEC2でG5g.xlなどのG5gインスタンスを起動します。
2. 必要なライブラリをインストールする
sudo yum -y groupinstall "Development Tools"
ls -al /usr/src/kernels/
とuname -r
の結果を確認します
3. バージョンが違ったら(kernelが遅れていたら)バージョンアップをしてそろえます
sudo yum update kernel
4. 一旦再起動しましょう
sudo reboot
5. gccまわりを調整する
今回、主にNVIDIAのドライバーインストール時にgcc10
を使いたいです。デフォルトだとgcc
コマンドはバージョン7のものになっており、gcc10
を使うように変更します。
(なお、ここはもっとスマートな方法がある気がする。)
とりあえずデフォルトのgcc(7)をmvしてgcc10を使うようにsym linkを張って対応します。
sudo mv /usr/bin/gcc /usr/bin/gcc-7
sudo ln -s /usr/bin/gcc10-gcc /usr/bin/gcc
6. CUDA Toolkitをインストールする
さて、ここからいよいよNVIDIAまわりで必要なものをインストールしていきます。
先に整理しておきますが、DockerコンテナでNVIDIA GPUを使うために必要なものは主に以下のものがあります。
- NVIDIAのドライバー
- CUDA
- CuDNN
- NVIDIA Container Toolkit
これらを一通り入れていきます。なおこの作業にはこちらの記事をとても参考にさせていただきました、ありがとうございます。
ではまずCUDA Toolkitをインストールします。ドライバーと別々にインストールする記事が多いですが、CUDA Tookitを入れるとドライバーも一緒に入れられるのでまとめて入れてしまいます。
こちらからrunファイルを取得してきて実行します。
sudo sh cuda_あなたの選んだバージョン_linux_sbsa.run --kernel-source-path=/usr/src/kernels/カーネルバージョン.amzn2.aarch64
runファイルの実行時に--kernel-source-path
を指定する必要があるので、忘れないでください。また、OpenGLなどはインストール時にオフっても問題ありません。
これが完了すると
- NVIDIAのドライバー
- CUDA
- CuDNN
- NVIDIA Container Toolkit
こういう状況になります。
7. CuDNNをインストールする
続いてCuDNNをインストールしていきます。
こちらから好きなバージョンを選んでダウンロードしてきます。
ダウンロードしたらtarを解凍してファイルをコピーします。
tar -xf cudnn-linux-sbsa-選択したバージョン_cudaバージョン-archive.tar.xz
sudo cp -P cudnn-linux-sbsa-選択したバージョン_cudaバージョン-archive/include/* /usr/local/cuda/include/
sudo cp -P cudnn-linux-sbsa-選択したバージョン_cudaバージョン-archive/lib/* /usr/local/cuda/lib64/
sudo chmod a+r /usr/local/cuda/lib64/*
sudo ldconfig
これが完了すると
- NVIDIAのドライバー
- CUDA
- CuDNN
- NVIDIA Container Toolkit
こういう状況になります。
8. NVIDIA Container Toolkitをインストールする
次にNVIDIA Container Toolkitをインストールしていきます。
先に、yum-utils
が必要になるのでインストールします。
sudo yum -y install yum-utils
そのあとはこちらに従ってNVIDIA Container Toolkitをインストールしていきます。
9. NVIDIA Container Toolkitの設定
インストールしただけでは使えないので、設定をします。
こちらに従って設定をします。
これが完了すると
- NVIDIAのドライバー
- CUDA
- CuDNN
- NVIDIA Container Toolkit
こういう状況になります。
この状態で、いま動かしているG5gのEC2インスタンス上では、Dockerコンテナ上でもGPUが使える状態になっているはずです。
10. AMIにする
最後に、この状態でAMIを作成します。
作成前に、
sudo systemctl stop ecs
sudo rm -rf /var/lib/ecs/data/*
sudo rm -rf /var/log/ecs/*
sudo mkdir -p /etc/ecs/
sudo touch /etc/ecs/ecs.config
echo "ECS_ENABLE_GPU_SUPPORT=true" | sudo tee -a /etc/ecs/ecs.config
sudo /usr/libexec/amazon-ecs-init reload-cache
このような感じで現在動いているamazon-ecs-agentの状態をクリアにしてから、AWSコンソール上からAMI作成を実行します。AMI作成自体はEC2のコンソールからボタンをクリックするだけです。
お疲れ様です、まずはGraviton + GPUなG5gインスタンスで使えるECS用のAMIができあがりました。
実行するコンテナを用意する
次に、今回はECS上で動かすことが目的なので実際にECSで使うコンテナイメージの用意をします。
ん?コンテナが動けばなんでも動くでしょって?
それがそうでもない…
1. ベースとなるコンテナイメージ
悲しい事実なんですが、AArch64向けにCUDAサポートでビルドされたPyTorchは存在しないようなので、自分でビルドするかNGCのイメージを使うしかありません……
PyTorchのビルドは時間がかかるということもあり、今回はNGCのものを使っていきます。
こちらに一覧があるので、好きなものを選んでください。
注意事項として、先ほどAMI作成時にインストールしたドライバーとCUDAのバージョンが合ってないと動かなかったりするので、よくよく見て選んでください。(Tagにその辺が書かれていないのでとても面倒なんですが…)
2. (オプション)torchaudio
このNGCのイメージではtorch
とtorchvision
は含まれていますがtorchaudio
は入っていないので必要な場合は自分で追加しないといけません。
しかしそのままpip installすると依存関係でNGCにデフォルトで入っているtorch
を消してAArch64汎用torch(CUDAサポートなし)をインストールしようとしてくるうえに、それでインストールしたtorchaudioは動きません。
というわけでこちらのコメントを参考に、Dockerfile内で自分でtorchaudioをビルドしてインストールするします。
git clone https://github.com/pytorch/audio
cd audio
pip install -v -e . --no-use-pep517
ECS on EC2で起動する
最後の関門です!
ここまで来たらGPUサポートなG5g用AMIはあるし、そこで動くコンテナイメージもあるし、あとはECS上でそれらを使ってクラスター・サービス・タスク構築していくだけです!
が、トラップもあるので最後まで気が抜けません。
1. 普通にECS on EC2クラスターを作成する
本来なら先に起動テンプレートをつくったりするのが筋なんですが、面倒なのでまずは普通にクラスターをつくって起動テンプレート含め自動でつくってもらいましょう。
ここではAMIを通常のAmazon Linux2(ARM、GPUサポートなし)を選択、インスタンスはg5gを選択して作成します。
2. 起動テンプレートの編集
クラスターが作成されると、ASGと起動テンプレートができあがっていると思うので編集します。
まず、起動テンプレートで指定されているAMIを先ほど自分が作成したものに変更します。
<<重要>>
次に起動テンプレート内、高度な詳細
でユーザデータ
にnvidia-smi
を追記をします。
これはamazon-ecs-init
がなぜかそのままだとGPUを認識しているのに認識していない(なに言ってるかわからない)状態に陥り、GPUサポートコンテナインスタンスとして認識されなくなるからです。
もう少しちゃんと説明すると、/var/log/ecs/ecs-init.log
を見ると普通にGPUを認識しているのに、なぜかGPUIDsが入らないということが発生します。
このissueもおそらくおなじことが起きていると思います。
動作検証を重ねた結果、nvidia-smi
を一発撃ち込むと正常化するというのがわかったので、ユーザデータを使ってインスタンス起動時に撃ち込むようにしました。[5]
(なにそのハック感)
(これの調査に半日かかった)
3. ASGの起動テンプレートバージョンの更新
起動テンプレートを編集して新しいバージョンをつくったので、ASGの起動テンプレートバージョンを更新します。これを忘れて古いバージョンのインスタンスが起きてきて悩む人はかなり多いと思います。
古いインスタンスが起動していたらデタッチして終了しちゃいましょう。
新しい起動テンプレートのインスタンスが起きてくるのを待ちます。
4. 通常通り、GPU使用のタスク定義を書いてECSを使う
あとは普通に使うだけです。タスク定義でGPUリソースを指定するのをお忘れなく。
これでECS on EC2 G5gインスタンスでGPUを使ったワークロードを実行できます。
お疲れさまでした!
まとめ
というわけで今回はGraviton2とNVIDIA T4G搭載のAWS EC2 G5gインスタンスを使ってECS上で機械学習ワークロードを実行するまでの道のりをご紹介しました。
Bottlerocket使えなかった時点で心は折れていたんですけどね、なんかここでGraviton + GPUを諦めたら一生使わないよねとか思ったりして。
3日くらい試行錯誤した結果、なんとか実用までこぎ着けたのはよかったです。ただじゃあこれを大規模なワークロードで使うかっていうと、現状はメンテナンスがツラすぎて正直厳しいかなと思っています。
だって、バージョンアップのたびに自分でAMIつくるのツラくない?
あと、AArch64なのでライブラリまわりを自分でビルドしないといけないものが多いのがなかなかにツラいです。
とはいえいつでも誰かにおんぶに抱っこというのもよくないですから、今回このように地道に構築していけたのはよかったなと思います。
amazon ecs agent/ecs initなどの知見もたまりましたし。
今後もECSやEC2、Gravitonなどのチップをフル活用してサービスをつくっていけたらと思います!
こんな感じで、Spiral.AIはいろいろチャレンジングなこともたくさんできる環境です!
もしSpiral.AIに興味がありましたら、お気軽にお声がけください。
Discussion