Google、大規模モデル トレーニングにおける Cloud TPU v4 Pod の実力を実証
Google Cloud Japan Team
※この投稿は米国時間 2021 年 12 月 2 日に、Google Cloud blog に投稿されたものの抄訳です。
最近、数十億または数兆ものパラメータを持つモデルが機械学習の機能と精度において目覚ましい進化を遂げています。たとえば、Google の LaMDA モデルは、多岐にわたるトピックについてユーザーと自由に流れるような会話を繰り広げることができます。機械学習の研究とプロダクト開発を行うコミュニティの間では、大規模モデルを活用した画期的な機能の提供に大きな関心が寄せられています。これら大規模モデルは高い演算能力を必要としており、それに伴ってモデル トレーニング プロセスの効率改善により焦点を当てる必要が生じています。そして、ベンチマーキングは効率向上のために ML システム コミュニティをまとめるための重要な手段です。
最近終了した MLPerf v1.1 Training ラウンド1 で、Google は Open 部門に 2 つの大規模言語モデル ベンチマークを提出しました。1 つは 4,800 億のパラメータ、2 つ目は 2,000 億のパラメータを持つものです。提出したベンチマークでは、一般公開されているインフラストラクチャを使用しました。たとえば、Cloud TPU v4 Pod スライスやオープンソースの Lingvo モデリング フレームワークです。
従来、この規模のトレーニング モデルは、何千万、もしくは何億ドルといったごく一部の企業しか払えないコストがかかるスーパーコンピュータを必要とすることが常でした。お客様はこれと同じ成果を、オンプレミス システムのインストールや維持にまつわる費用を負担することなく、エクサフロップ級の処理能力を持つ Cloud TPU v4 Pod を使って得ることができます。
大規模モデル ベンチマーク
Google が Open 部門に提出したのは、TensorFlow を使った、パラメータ 4,800 億の高密度 Transformer ベースのエンコーダのみのベンチマークと、パラメータ 2,000 億の JAX ベンチマークです。これらのモデルは、MLPerf の BERT モデルと類似したアーキテクチャを持ちますが、ディメンションがさらに大きく、レイヤ数も多くなっています。提出したベンチマークでは、2 つの異種フレームワークにまたがった大規模モデルのスケーラビリティと TPU の高パフォーマンスが実証されました。特筆すべきは、スタック Transformer アーキテクチャを持つこれらのベンチマークが、演算特性の面で他の大規模言語モデルにほぼ匹敵するものであったことです。
Google の 2 つの提出結果は、それぞれ 2,048 個のチップと 1,024 個のチップの TPU v4 Pod スライスでベンチマークを実施して得られたものです。エンドツーエンドのトレーニング時間は、4,800 億のパラメータ モデルで約 55 時間、2,000 億のパラメータ モデルで約 40 時間を達成できました。それぞれの実行での演算効率は 63% を達成しました（使用したシステムのピーク FLOP に対するモデルの浮動小数点演算とコンパイラ再実体化の割合として計算2）。
大規模モデル トレーニング向け次世代 ML インフラストラクチャ
このような素晴らしい結果を残すには、いくつかの最先端テクノロジーが必要でした。まず、各 TPU v4 チップの演算能力は TPU v3 チップの 2 倍超（最大 275 ピーク TFLOPS）です。2 つ目は、典型的な GPU ベースの大規模トレーニング システムと比べてチップあたり 10 倍の帯域幅を実現する超高速の相互接続によって、4,096 個の TPU v4 チップがまとめて Cloud TPU v4 Pod にネットワーク接続されます。大規模モデルでは通信が非常に頻繁に行われ、ローカルの演算がネットワーク全体で通信されるリモートの演算結果に依存することが多々あります。TPU v4 の超高速相互接続は、ネットワークのレイテンシと輻輳を取り除くことで、大規模モデルの演算効率に非常に大きな影響を与えます。
Google の提出結果で示されたパフォーマンス値には XLA 線形代数コンパイラが大きな役割を果たしています。また、Lingvo フレームワークも活用されています。XLA は最適化を透過的に実施します。たとえば、ML モデルの構成要素を形成する演算グラフの多くの GSPMD ベースの自動並列化があります。また、XLA では演算と通信をオーバーラップさせることによって、レイテンシの削減も可能です。Google の 2 つの提出結果は、TensorFlow と JAX の 2 つのフレームワークにまたがる Google のソフトウェア スタックの汎用性とパフォーマンスを実証しました。
MLPerf の大規模モデル
Google が提出したベンチマークは、ML の研究や本番環境でますます重要性を高めているものの、MLPerf の Closed 部門ベンチマーク スイートでは現在評価できないクラスのモデルが持つ可能性を表すものです。
これらのモデルをベンチマーク スイートに加えることが次の重要なステップであり、これによって、大規模モデルが提示するスケーラビリティの課題に焦点を当てるよう ML システム コミュニティに働きかけることができると Google は信じています。
Google の提出結果では、業界最高水準である 63% の演算効率が認められました。演算効率が高いとトレーニングの処理速度が上がり、実験の速度も向上します。すなわち、Google の Cloud TPU をお使いのお客様にとって、コスト削減に直結します。
TensorFlow、PyTorch、JAX を使った Cloud TPU の活用方法について詳しくは、Cloud TPU ホームページとドキュメントをご覧ください。
1. MLPerf の名前とロゴは、米国および他の国々における MLCommons Association の商標です。権利はすべて同組織が有しており、無断使用は固く禁じられています。詳しくは、www.mlcommons.org にアクセスしてください。
2. 演算効率とエンドツーエンドのトレーニング時間は公式の MLPerf 指標ではありません。
- TPU テクニカル プログラム マネージャー Pankaj Kanwar
