StablityAI社の分散ML学習基盤のアーキテクチャの要約と独自補足
元の資料
aws re:Inventで発表されていた「How Stable Diffusion was built: Tips & tricks to train large AI models」が元資料。
34:00時点から、stable diffusionを学習させたときに使った、分散ML学習基盤の説明を行っていたので、要約と補足を加える。
内部的なものも含めて使っているAWSサービス一覧
compute
-
p4dインスタンス
A100GPU搭載の機械学習向けのインスタンス -
trainiumインスタンス
機械学習向けのtrainiumチップを搭載したインスタンス
network
-
Elastic fabric adaptor
インスタンス間の通信パフォーマンスを向上させる -
VPC
仮想ネットワーク
storage
-
S3
永久ストレージ -
FSx for Lustre
高速かつスケーラブルな共有ディスク（インスタンスストレージとS3の間に置いて使ったりする) -
EC2 instance store
揮発性のストレージ。インスタンスが止まるとデータが消える。
Orchestration
-
AWS ParallelCluster
ポッドやノードのオートスケールではなく、クラスタごとオートスケールできる -
AWS Cloud Formation
AWSのリソースをプログラミング言語やテキストファイルを使用して自動で構築できる
実際のアーキテクチャ
42:01あたりで、「(分散ML学習基盤を)自前で用意するのは非常に大変なため、stabilityAI社や他の会社は、AWS ParallellClusterを使って構築している」と発言している。
上記のAWS ParallellClusterは、並列計算を行うためのHPCクラスタを簡単なコンフィグファイルを用意すれば、構築できるというもの。
ParallellCluster使ってHPCをAWSで構築する際に必要なコンフィグファイルのイメージを掴んだり、その他の細かい情報のためには、こちら資料がわかりやすかった。
また、AWS ParallellClusterを使って機械学習モデルを学習させているこちらの資料が参考になりそうだった。
上記のストレージ部分の詳細
インスタンスストレージ、Fsx lustre、S3を使っている
checkpointはどのストレージでも使われているが、使い分けみたいなのは名言されていなかった（と判断している)
ユーザーの管理
44:07あたりで説明を行っており、AWS Parallel Clusterの認証のために、AWS Directory Serviceを使っているらしい。
その他
インストールされているソフトウェア
ユーザーが自由に入れられるものと、インスタンスが作成されるときにインストールされるものがあるらしい
ジョブキュー
優先順位度の高い学習ジョブとそのキュー、中断可能な前処理のような優先度の低いジョブとそのキューを使い、GPU利用率をできるだけ高める工夫を行っているらしい。
モニタリング
ジョブを走らせる前に、適切なGPUがついているか・十分なbandwidthがあるかなどをチェックしているらしい。
cloud watch,Grafanaなどを使って、GPU利用率・CPU利用率・ストレージスループットを見ている。
Discussion