Google Cloudは、同社のスーパーコンピューティングアーキテクチャ「AI Hypercomputer」の活用を支援するレファレンス実装や、ソフトウェアレイヤーのアップデートを発表した。
この記事は会員限定です。会員登録(無料)すると全てご覧いただけます。
Google Cloudは2024年10月26日(米国時間)、生成AI(人工知能)に最適化された同社のアーキテクチャ「AI Hypercomputer」の活用を支援するレファレンス実装やソフトウェアレイヤーのアップデートを発表した。
AI Hypercomputerは、パフォーマンスに最適化されたハードウェア、オープンソフトウェア、ニーズに応じた柔軟な利用プランを提供するスーパーコンピューティングアーキテクチャだ。
Google Cloudの主な発表内容は以下の通り。
「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のトレーニング性能がほぼ線形にスケーリングするという。
ほとんどの「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は述べている。
Copyright © ITmedia, Inc. All Rights Reserved.