AWS上でのファウンデーションモデルのトレーニングと推論のためのビルディングブロック

本記事は、AWS上でのファウンデーションモデルのトレーニングと推論に関わる機械学習エンジニアや研究者を対象としています。主にオープンソースソフトウェアスタックに基づき、AWSインフラが大規模分散トレーニングと推論におけるシステムのボトルネックとスケーリング特性にどう対応するかを分析します。

ファウンデーションモデルのスケーリング則は変化しています。従来は事前学習の計算量増加が性能向上を牽引していましたが、現在は事後学習（教師ありファインチューニング、強化学習）やテスト時計算（長考、検索、マルチサンプル戦略）も重要です。AWSのインフラは、緊密に結合されたアクセラレータコンピュート、高帯域低遅延ネットワーク、分散ストレージバックエンドを備え、これら3つのスケーリング領域をサポートします。

インフラ層では、AWSは複数世代のGPUインスタンスを提供します。P5インスタンスはH100 GPU、P5e/P5enはH200、P6はBlackwell B200およびB300を搭載します。ピークテンソルスループット、HBM容量と帯域幅、相互接続帯域幅が主要なスケーリング軸です。例えば、H100のBF16/FP16テンソルピークは0.9895 PFLOPS、B300は2.25 PFLOPSで、FP4では13.5 PFLOPSに達します。HBM容量はH100の80 GBからB300の288 GB HBM3eへ、帯域幅も3.35 TB/sから8 TB/sへ向上しています。

マルチGPUインスタンスはNVLink/NVSwitchによるノード内高速相互接続を提供し、オールリデュースなどの集団通信をホストネットワークスタックを経由せず実行します。ノード間のスケールアウトには、Elastic Fabric Adapter（EFA）を使用します。EFAはカーネルバイパスRDMAを提供し、SRDプロトコルを利用します。EFA v2（P5）、EFA v3（P5en、レイテンシ35%削減）、EFA v4（P6、集団通信性能18%向上）が利用可能です。

ストレージは階層化されています。ホットデータ用のローカルNVMe SSD、共有高スループットアクセス用のLustre並列ファイルシステム（Amazon FSx for Lustre管理）、永続ストレージ用のAmazon S3です。LustreはS3とのデータリポジトリ関連付けにより、トレーニングデータセットの遅延ロードとチェックポイントの自動エクスポートを実現します。

大規模クラスタでは、Amazon EC2 UltraClustersが数千のアクセラレータインスタンスを単一アベイラビリティゾーンに配置し、ペタビット級のノンブロッキングネットワークで相互接続します。通信集約型ワークロード（MoEモデルのエキスパート並列性など）では、NVLinkドメインのサイズが重要です。AWS UltraServersは、専用アクセラレータ相互接続で複数インスタンスのNVLinkドメインを拡張します。例えば、P6e-GB200 UltraServerはNVIDIA GB200 NVL72プラットフォームに基づき、単一NVLinkドメイン内に最大72基のBlackwell GPUと13.4 TBのHBM3eを集約します。その構成インスタンスp6e-gb200.36xlargeは4基のGPUと200 GB/sのEFA帯域幅を持ち、UltraServer全体では1,800 GB/sのEFA帯域幅を提供します。

また、記事はオープンソースソフトウェアスタックの重要性を強調します。リソース管理にはSlurmやKubernetes、モデル開発にはPyTorchやJAX、可観測性にはPrometheusやGrafanaが使用されます。AWSインフラはこれらのツールと統合され、ファウンデーションモデルのライフサイクル全体をサポートします。今後の連載記事では、各層の詳細な実装を取り上げる予定です。