🔖

StablityAI社の分散ML学習基盤のアーキテクチャの要約と独自補足

2022/12/09に公開約2,100字

元の資料

https://www.youtube.com/watch?v=7I854do63Lg

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を使って機械学習モデルを学習させているこちらの資料が参考になりそうだった。
https://aws.amazon.com/blogs/machine-learning/building-an-interactive-and-scalable-ml-research-environment-using-aws-parallelcluster/

上記のストレージ部分の詳細

インスタンスストレージ、Fsx lustre、S3を使っている
checkpointはどのストレージでも使われているが、使い分けみたいなのは名言されていなかった(と判断している)

ユーザーの管理

44:07あたりで説明を行っており、AWS Parallel Clusterの認証のために、AWS Directory Serviceを使っているらしい。

その他

インストールされているソフトウェア

ユーザーが自由に入れられるものと、インスタンスが作成されるときにインストールされるものがあるらしい

ジョブキュー

優先順位度の高い学習ジョブとそのキュー、中断可能な前処理のような優先度の低いジョブとそのキューを使い、GPU利用率をできるだけ高める工夫を行っているらしい。

モニタリング

ジョブを走らせる前に、適切なGPUがついているか・十分なbandwidthがあるかなどをチェックしているらしい。
cloud watch,Grafanaなどを使って、GPU利用率・CPU利用率・ストレージスループットを見ている。

Discussion

ログインするとコメントできます