Trillium による LLM トレーニング効率の改善 — 性能分析
Mohan Pichika
Group Product Manager
Vaibhav Singh
Group Product Manager
※この投稿は米国時間 2024 年 11 月 14 日に、Google Cloud blog に投稿されたものの抄訳です。
生成 AI モデルの急速な発展により、ハードウェア アクセラレータの性能と効率についてかつてないほど要求が高まっています。Google は先月、第 6 世代の Tensor Processing Unit（TPU）である Trillium を発表し、次世代モデルの需要に応えています。Trillium は、極めて大規模なトレーニングでの利用を目的に、チップからシステム、そして Google データセンターのデプロイまで、大きな規模でのパフォーマンスを重視して構築されています。
このたび、Trillium について初の MLPerf トレーニング ベンチマーク結果を発表しました。MLPerf 4.1 トレーニング ベンチマークでは、Trillium は前世代の Cloud TPU v5p との比較で 1 ドルあたり 1.8 倍のパフォーマンスを発揮し、（スループットの）スケーリング効率が 99% にもなることが示されています。
このブログ記事では、Trillium のパフォーマンスを簡潔に分析し、これまでで最もパフォーマンスに優れた TPU として傑出している理由を示します。システム比較指標の概観から始めて、まず従来のスケーリング効率を取り上げます。スケーリング効率に加えて考慮すべき重要指標として収束スケーリング効率を取り入れます。これらの 2 つの指標と 1 ドルあたりのパフォーマンスを評価し、Trillium の Cloud TPU v5p に対する比較を提示します。結論として、情報に基づいてクラウド アクセラレータを選択するためのガイダンスを示します。
従来のパフォーマンス指標
アクセラレータ システムの評価や比較は、ピーク スループットや実効スループット、スループットのスケーリング効率など複数の尺度で行われます。これらの指標はそれぞれ評価基準として有用ですが、収束時間が考慮されていません。
ハードウェア仕様とピーク パフォーマンス
従来の比較は、ピーク スループット、メモリ帯域幅、ネットワーク接続性などハードウェア仕様に重点を置いていました。こうしたピーク値では、理論的な限界は明らかになりますが、実世界のパフォーマンスはアーキテクチャ設計やソフトウェア実装に大きく依存するため、その予測にはうまく作用しません。現在の ML ワークロードは数百から数千ものアクセラレータで実行されるのが通常で、特定のワークロードに向けた適切な規模のシステムでの実効スループットが重要な指標になります。
使用率パフォーマンス
システム パフォーマンスは、モデルの実効 FLOP 使用率（EMFU）やメモリ帯域幅の使用率（MBU）など、ピーク能力に対する実スループットを測る使用率指標でも定量化できます。しかし、これらのハードウェア効率指標は、トレーニング時間やモデル品質のようなビジネス価値の指標には直接変換できません。
スケーリング効率とトレードオフ
システムのスケーラビリティは、強いスケーリング（固定のワークロードに対するシステムサイズによるパフォーマンス改善）と弱いスケーリング（ワークロードとシステムサイズが比例的に拡大するときの効率）の両方で評価されます。どちらの指標も有用な評価基準になりますが、究極の目的は高品質なモデルを迅速に実現することであり、スケーリング効率を犠牲にしてもトレーニング時間の高速化やモデル収束の改善を求めることに意味がある場合があります。
収束スケーリング効率の必要性
ハードウェア使用率やスケーリング指標からシステムに関する重要な情報が得られるのに対し、収束スケーリング効率はトレーニングの根本的な目標、すなわち効率的なモデル収束に重点を置いています。収束とは、モデルの出力の改善が止まり、エラー率が定常的になる点を意味します。収束スケーリング効率は、コンピューティング リソースの追加によってトレーニング プロセスの完成が加速する程度を評価するものです。
収束スケーリング効率は 2 つの主要測定値で定義します。すなわち、N₀ 個のアクセラレータのクラスタが時間 T₀ で収束する場合の基本値と、N₁ 個のアクセラレータが時間 T₁ で収束する場合のスケーリング値です。クラスタサイズの増加に対して得られる収束時間の高速化の比率で、次のように得られます。
収束スケーリング効率が 1 であれば、クラスタサイズと同一の比率で解決時間が改善することを意味します。このため、収束スケーリング効率が可能な限り 1 に近くなることが望まれます。
このコンセプトを応用して、Trillium と Cloud TPU v5p を使用した GPT3-175b トレーニング タスクの MLPerf での評価結果を確認してみましょう。
Trillium のパフォーマンス
GPT3-175b トレーニングの結果を、Trillium の 4 つの異なる構成、Cloud TPU v5p の 3 つの異なる構成について計測しました。以下の分析では、比較のために、トータルピーク FLOPS が同じクラスタサイズに結果をグルーピングしています。たとえば、Cloud TPU v5p-4096 構成は 4xTrillium-256 と、Cloud TPU v5p-8192 は 8xTrillium-256 と、というようにそれぞれ比較します。
この分析で示すすべての結果は、Cloud TPU と GPU 向けの Google の高性能リファレンス実装である MaxText を使用したものです。
弱いスケーリングの効率
クラスタサイズをバッチサイズと比例的に拡大した場合、Trillium と TPU v5p は以下のようにどちらもほぼ線形なスケーリング効率を示します。
図 1: ソースデータ: Trillium と v5p の GPT3-175b トレーニング タスクでの MLPerf™ 4.1 Training Closed の結果、2024 年 11 月時点: Trillium と Cloud TPU v5p の弱いスケーリングの比較。v5p-4096 と 4xTrillium-256 をスケーリング ファクター比較のベースとし、「n x Trillium-256」は単一 ICI ドメイン中に 256 チップを搭載した Trillium Pod が n 個あることを示す。v5p-n は単一 ICI ドメインに v5p チップが n/2 個あることを示す。
図 1 はベース構成からクラスタサイズが増加したときの相対的なスループットのスケーリングを示しています。Trillium は、Cloud TPU マルチスライス テクノロジーを使用して複数のデータセンター ネットワーク間で実行した場合でもスケーリング効率が 99% に達しており、単一 ICI ドメイン内の Cloud TPU v5p クラスタが示す 94% のスケーリング効率を上回っています。この比較では、1,024 チップ（Trillium-256 Pod 4 基）のベース構成を使用し、最小の v5p 構成（v5p-4096、2,048 チップ）と対応するベースラインとしています。最小構成の 2x Trillium-256 Pod との比較で計測すると、Trillium では 97.6% ものスケーリング効率を維持しています。
収束スケーリング効率
先に示したように、弱いスケーリングは有用ですが、価値を示す尺度としては十分ではありません。一方、収束スケーリング効率では解決までの時間を考慮に入れます。
図 2: ソースデータ: Trillium と v5p の GPT3-175b トレーニング タスクでの MLPerf™ 4.1 Training Closed の結果、2024 年 11 月時点: Trillium と Cloud TPU v5p の収束スケーリング効率
最大のクラスタサイズに対して、Trillium と Cloud TPU v5p で収束スケーリング効率はほぼ同等となりました。この例では、最も右の構成で CSE が 0.8 になっているのは、クラスタサイズが（ベース）構成の 3 倍になっていて、収束までの時間がベース構成に対して 2.4 倍改善しているということを意味します（2.4/3 = 0.8）。
収束スケーリング効率は Trillium と TPU v5p でほぼ同等ですが、Trillium の真価はより低コストで収束に達するということですので、ここで最後の指標を見てみましょう。
トレーニング コスト
弱いスケーリング効率と収束スケーリング効率によってシステムのスケーリング特性が示されますが、最も重要な指標であるトレーニング コストは確認できていません。
図 3: ソースデータ: Trillium と v5p の GPT3-175b トレーニング タスクでの MLPerf™ 4.1 Training Closed の結果、2024 年 11 月時点: Cloud TPU v5p と Trillium の実経過時間とオンデマンドの正規料金に基づくトレーニング コストの比較。
Trillium では TPU v5p と比較してトレーニングのコストが最大 1.8 倍低減（45% 低下）しており、同一の検証精度までの収束が得られています。
情報に基づいたクラウド アクセラレータの選択
この記事では、アクセラレータ システムの比較の複雑さについて解説し、単純な指標の確認にとどまらず真のパフォーマンスと効率を評価することの重要性を強調しました。ピーク パフォーマンス指標は開始点としては機能しますが、実世界での有用性を予測するには多くの場合不足があります。実際、モデルの実効 FLOP 使用率（EMFU）やメモリ帯域幅の使用率（MBU）のような指標の方が、アクセラレータの対応能力について概してより意味のある情報を示します。
ワークロードとリソースの拡大に対応するシステムのパフォーマンス変化を評価するうえで、スケーリング特性（強いスケーリングと弱いスケーリング）が重要であるということを強調しました。しかし、最も客観的な尺度であると判断したのは収束スケーリング効率で、これによれば、単なる速度ではなく、同一の結果を得るのに要する能力に基づいてシステムを比較できます。
これらの指標を GPT3-175b トレーニングを使用したベンチマーク結果に適用することにより、Trillium が Cloud TPU v5p とほぼ同等の収束スケーリング効率を達成する一方で、1 ドルあたりのパフォーマンスが最大 1.8 倍優れていて、トレーニングのコストが低減されることを示しました。これらの結果によって、アクセラレータ システムをパフォーマンスと効率の複数の尺度で評価することの重要性が示されました。
ML アクセラレータの評価では、リソース使用率指標（EMFU、MBU）、スケーリング特性、収束スケーリング効率を組み合わせて総合的に分析することをおすすめします。この多角的なアプローチによって、対象のワークロード要件とスケールに対してデータ駆動の意思決定が可能になります。
Trillium についての詳細については、発表時のブログ記事またはドキュメントをご覧ください。
ー グループ プロダクト マネージャー Mohan Pichika
ー グループ プロダクト マネージャー Vaibhav Singh
