Google Geminiのトレーニング：TPU、複数データセンター、宇宙線のリスク

研究論文、公開コメント、Google CloudのCEOとの対話から読み取れること

先週のGeminiのローンチは、世界で最も人気のある生成AIプラットフォームとしてのOpenAIのGPTモデルに対し、AI企業としての主導権を再び示そうとするGoogleにとって最大の取り組みとなりました。

そのモデルファミリーを構築するために、同社は大規模なインフラ構築に着手し、NvidiaのGPUに依存せずにモデルのトレーニングが可能であることを証明しようとしました。

Googleは、3つのGeminiモデルをどのように構築したかについての具体的な詳細を提供することを避けましたが、本記事ではDCDがインタビュー、研究論文、および公開コメントから学んだことをまとめました。

最大のモデルであるGemini Ultraは、複数のデータセンターにまたがってトレーニングされており、PaLM-2のような以前の単一施設モデルよりも進歩しています。

Google CloudのCEOであるThomas Kurian氏は、DCDに対し次のようにコメントしています。「私たちは、厳密に何か所であるかは公表していませんが、複数のサイトと、そのサイト内の複数のクラスタにわたってトレーニングされています」

「我々はマルチホストと呼ばれる技術を使ってトレーニングを分散させています。通常、我々がトレーニングを分散させる理由は、例えば片側で電源トラブルやその他の問題が発生した場合でも、レジリエンスを確保できるようにするためです。それはまた、スペースや電力を考慮し、より大きなマシンのクラスタを展開することも可能にします」

マルチホストはGoogle Cloud上で顧客に公開される予定であり、これにより間もなく、ひとつのランで共に稼働する施設のレイテンシの限界（つまり距離の制限）をよりよく理解できるようになるはずです。

これを実現するためにGoogleは、Googleのクラスタ内（同一クラスタに属する2つのオブジェクト間の距離）ネットワークとクラスタ間（異なる2つのクラスタに属する2つのオブジェクト間の距離）ネットワークを使い、TPU SuperPods（4096 基のチップ構成）を連結しているとGoogle Deepmindの論文は述べています。同社はSuperPod内のモデル並列性と複数のSuperPod間のデータ並列性を活用しています。

その中心にあるのが、Googleのネットワーキング・プラットフォーム「Jupiter」です。

これは、データセンターのスパインに代わる社内の光スイッチングネットワーク（現在はOCSとして知られているが、以前はMission Apolloと呼ばれていた）に依存しています。

従来のネットワーク・トポロジーでは、信号は電気と光の間を行ったり来たりしていました。「それはすべてホップ・バイ・ホップであり、電気に変換しては光に戻し、といった具合に、ほとんどの作業を電子的な領域に残したままでした」と、Googleのシステム＆サービス・インフラストラクチャー・チーム・リードのAmin Vahdat氏は今年初めに語っていました。

「これはコストとエネルギーの両面で高くつきます」

同社のカスタムメイドのOCSでは、「可能な限りデータを光領域に残し」、小さなミラーを使って光源からの光ビームをリダイレクトし、光クロスコネクトとして宛先ポートに直接送信しています。

これにより、コストだけでなくレイテンシーも劇的に低減され、マルチホストも可能になります。

Ultraでは、OCSを使用して、4x4x4チップキューブを任意の3Dトーラスのトポロジーに約10秒で動的に再構成しました。「我々は、ホットスタンバイとローリングメンテナンスを可能にするため、SuperPodごとに少数のキューブを保持することにしました」

Geminiの各バージョン（Nano、Pro、Ultra）は、GoogleのTPU（TPUv5eとTPUv4を混在させて使用）上でトレーニングされています。

Proモデルは数週間でトレーニングされ、「Ultraのリソースのほんの一部を活用している」とGoogle Deepmindのレポートは述べています。

「Gemini Ultraのトレーニングには、TPUv4アクセラレータの大規模なフリートを使用した」と、同レポートは付け加えています。

TPUv5eは、大規模言語モデルのトレーニングにおいて、TPUv4と比べて価格性能で2.3倍向上しています（Google調べ）が、正式に発表されたのは今年11月のことです。Google内部ではもっと早くから利用可能だったはずですが、Ultraに対応できるほどの量ではなかったのかもしれません。

しかし、ChatGPTに匹敵するモデルがGoogleのハードウェアでトレーニングされたという事実は、AI開発に対するNvidiaの牙城を崩す重要な一歩につながると、同社は主張しています。

「私は最近、”NvidiaのGPUなしではMLは実現しない “とか、”MLはNvidiaのCUDAでしかできない “と主張する誤った情報の投稿や記事を複数目にしてきました」と、GoogleのCloud TPU成長責任者であるMax Sapozhnikov氏はTwitter/Xでコメントしています。

「このような通説が解明されるのは素晴らしいことです。GPUやCUDAに依存することなく、TPU上で多くの素晴らしいMLが行われています」。

「Anthropic、Midjourney、Salesforce、その他多くの企業は、すでにTPU上でスタックを構築しており、ハードウェアのコストと電力効率、XLAコンパイラの最適化のメリットを活用している」 と彼は付け加えています。

Anthropic社は、AWSがGoogleとともに同社に投資した後、AmazonのTrainiumハードウェアも使用しています。

Googleは最新のTPUのサーマル・デザイン・ポイント（TDP）をまだ公表していませんが、2021年の研究論文によると、TPUv3のTDPは450W（DSAメモリーシステムの電力に加え、サーバーホストの電力のシェアを含めると660W）で、前チップの280W/460Wから上昇していたことが明らかになりました。

同社は2018年からTPUを液冷化しています。「我々は、非常に高密度なフットプリントを有し、液冷のような高度な機能を備えた大規模システムを展開しており、これによりシステムから大幅に高いスループットを得ることができる」と、Kurian氏は述べています。

液冷は高密度データセンターの温度上昇には不可欠ですが、意図しない副作用もあります： それは「宇宙線」です。

今年初め、NTTと北海道大学の研究者は論文で、「半導体を水で冷却した場合、熱中性子数が著しく増加することが予想される」との発表を行いました。

スケールの大きなハードウェアはすべて放射線の影響を受けやすく、プロセスノードが小さくなればなるほど、そのリスクは高まります。これはデータを破壊する可能性があり、さらに懸念されるのは無言のデータ破壊につながることです。半導体放射線の研究者であるAndrew Keller氏が昨年DCDに語ったところによると、チップは「まだビットを処理しているが、そのデータが間違っている-しかもそれが間違っていることに気づいていない」可能性があるといいます。

Keller 氏は、コロラド州デンバーにある10万基のFPGAが稼動するデータセンターが、どれくらいの頻度で放射線の影響を受けるかを調べようとしました。その結果、平均30分ごとにコンフィギュレーション・メモリの問題が発生し、0.5～11日ごとにSDCが発生することが判明しました。

GoogleのTPUが放射線の影響をどの程度受けるかは不明ですが、同社は研究論文の最初のリリースでTPUについて言及しています。「このような大規模なハードウェアアクセラレータでは、宇宙線などの外的要因によって、純正のマシンが故障することはよくあることです」不思議なことに、この論文は静かに更新され、この言及は削除されました。DCDはGoogleにその理由を尋ねています。

GoogleはGemini Ultraについて、「SDCイベントは1～2週間ごとにトレーニングに影響を与える」と予想していると述べています。欠陥のあるハードウェアを迅速に検出し、排除するためには、アイドル状態のマシンやホットスタンバイのSDCスキャナをプロアクティブに組み合わせ、決定論的リプレイを利用して不正な計算を分離するいくつかの新しいテクニックが必要でした」

同社は、「完全決定論的インフラストラクチャ」によって、Ultraモデルの開発中に根本原因（ハードウェア障害を含む）を迅速に特定することができ、これは安定したトレーニングに向けた重要な要素であったと付け加えています。

最後に、この論文では、同社が「永続的なクラスタ・ストレージに重みを定期的にチェックポイントするという従来のアプローチ」を避けたことが明らかにされています。その代わりに、同社は 「モデル状態の冗長なインメモリーコピーを利用し、予定外のハードウェア障害が発生した場合は、無傷のモデルレプリカから直接迅速にリカバリーする」としています。

これは、PaLMやPaLM-2のトレーニング実行とは異なり、「かなり大きなトレーニングリソースが使用されているにもかかわらず、復旧時間の大幅なスピードアップ」につながりました。

グッド・プット（単位時間当たりにネットワークから特定の宛先に配信される有用な情報ビットの数）は、「最大規模のトレーニング・ジョブの場合、85～97％に増加した」と同社は述べています。

この記事は海外Data Centre Dynamics発の記事をData Center Cafeが日本向けに抄訳したものです。