Zenn
Spiral.AIテックブログPublicationへの投稿
🌟

Graviton2 + NVIDIA GPUなインスタンスを使ったECS on EC2における機械学習ワークロードの実行

ken
2024/01/16に公開
AWS
ECS
#
graviton
#
g5g
tech

少し遅くなりましたが、あけましておめでとうございます。
先日弊社のMission/Challenge/Valueについて代表が語るnoteが公開されましたが、僕自身も引き続きひたすらいろいろなことにチャレンジしていきたい所存です。
そして早速、新年のチャレンジとしては面白いことができたのでお話ししたいと思います。

Graviton2 + NVIDIA GPUなインスタンス

どうも、AWS大好き芸人の@ken11です。AWSが好きすぎて毎日目黒経由で通勤してます。嘘です。

G5g

皆さんAWSのG5gインスタンスファミリーはご存じでしょうか？

https://aws.amazon.com/jp/ec2/instance-types/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が存在します。

https://aws.amazon.com/jp/bottlerocket/

じゃあ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"

3. ls -al /usr/src/kernels/とuname -rの結果を確認します

バージョンが違ったら(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のものを使っていきます。

https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags

こちらに一覧があるので、好きなものを選んでください。
注意事項として、先ほど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に興味がありましたら、お気軽にお声がけください。

脚注

  1. 今回お話しするG5gはGraviton2搭載です ↩︎

  2. 個人的には、だったらG4dngとかにしてほしかったみたいな気持ちはあります。G5系なのにAmpereちゃうんかーいみたいな ↩︎

  3. やたらと古いAMIならないこともないが…… ↩︎

  4. AWSオフィシャルのBottlerocketが満足にG5gを動かせないと知った瞬間にAWSさん本当にGraviton普及させる気ある？という気持ちになったのは内緒。 ↩︎

  5. 一応ソースコード追ってるんだけどなぜそういうことが起きるのか謎。GPUIDの確認に使われているnvmlが空を返しているとかそういうことなんだろうか？ ↩︎

Spiral.AIテックブログ
Spiral.AIテックブログPublication

実在する芸能人との会話ができる日本初のAIサービス「Naomi.AI」やカスタムChatGPTを作れる「Spiral.bot」を提供するSpiral.AI株式会社のテックブログです。

Discussion