Google Cloud、生成AI構築に最適化したアーキテクチャ「AI Hypercomputer」の活用を支援するレファレンス実装を公開：大規模トレーニングのパフォーマンス最適化とモニタリングを支援

Google Cloudは、同社のスーパーコンピューティングアーキテクチャ「AI Hypercomputer」の活用を支援するレファレンス実装や、ソフトウェアレイヤーのアップデートを発表した。

[＠IT]

　Google Cloudは2024年10月26日（米国時間）、生成AI（人工知能）に最適化された同社のアーキテクチャ「AI Hypercomputer」の活用を支援するレファレンス実装やソフトウェアレイヤーのアップデートを発表した。

　AI Hypercomputerは、パフォーマンスに最適化されたハードウェア、オープンソフトウェア、ニーズに応じた柔軟な利用プランを提供するスーパーコンピューティングアーキテクチャだ。

　Google Cloudの主な発表内容は以下の通り。

「MaxText」が「A3 Mega VM」をサポート

　「MaxText」は、Google CloudのTPU（Tensor Processing Unit）およびGPUでの大規模言語モデル（LLM）の学習や推論を支援するレファレンス実装だ。NVIDIAの「H100 Tensor Core GPU」を搭載したインスタンス「A3 VM」と比較してGPU間ネットワーク帯域幅が2倍向上した「A3 Mega VM」で、MaxTextを使用できるようになった。

　Google Cloudによると、A3 Mega VMとMaxTextを使用した場合、クラスタ内のVM数に応じて、LLMのトレーニング性能がほぼ線形にスケーリングするという。

A3 Mega上での「Llama2-70b」（MaxText）の事前学習（Pre-training）に関するパフォーマンス比較　青色が理想的なスケーリングで青色が実際の結果を示している（提供：Google Cloud）

MoEのレファレンス実装とカーネルをオープンソースに

　ほとんどの「Mixture of Expert」（MoE）のユースケースでは、少数の「エキスパート」を活用して応答を生成することが重視される一方、特定のユースケースにおいては多様なエキスパートを活用して豊かな応答を生成することが求められる。

　こうした柔軟性を提供するために、MaxTextを拡張し、「制限付き」と「無制限」のMoE実装を追加した。これにより、モデルのアーキテクチャとして最適な実装を選択できる。制限付きMoEモデルは予測可能なパフォーマンスを提供し、無制限モデルはパフォーマンスを最適化するためにリソースを動的に割り当てる。

　また、MoEトレーニングを加速させるため、Cloud TPUにおける「ブロック疎行列乗算」に最適化された「Pallas」カーネルをオープンソース化した。Pallasは、「JAX」の拡張機能であり、GPUやTPUなどのXLA（Accelerated Linear Algebra）デバイス向けに生成されたコードを細かく制御できる。これらのカーネルは、「PyTorch」とJAXの両方で使用でき、MoEのトレーニングタスクを支援するという。

大規模トレーニングのモニタリングを支援するレファレンス実装

　大規模トレーニングジョブのモニタリングやダッシュボード構築を支援するレファレンス実装をGitHubで公開した。このレファレンス実装を使用すると、Google Cloudプロジェクト内でクラウドモニタリングダッシュボードを作成できる。このダッシュボードでは、平均または最大CPU使用率などの統計指標が表示され、アーキテクチャ上の異常を検知して対策を講じることができるという。

　「トレーニングタスクにおいて、一体となって動作することが求められる大規模クラスタは、MLOpsにおいて複雑さが増す。クラウドリソースを最大限活用し、AIシステムのメトリクスを改善させるには、適切な指標を使用してトレーニングジョブをモニタリングすることが不可欠だ」と、Google Cloudは述べている。