GoogleがAI向けチップ「TPU v5p」を発表、前世代から最大2.8倍の性能向上で「Gemini」のトレーニングにも使用される

現地時間2023年12月6日に、Googleが独自開発を進める機械学習特化のプロセッサ「Tensor Processing Unit(TPU)」の新モデル「TPU v5p」が発表されました。GoogleはこのTPU v5pについて「TPU v5pと同時に発表されたマルチモーダルAIのGeminiのトレーニングにも使用されています」と述べています。

Introducing Cloud TPU v5p and AI Hypercomputer | Google Cloud Blog
https://cloud.google.com/blog/products/ai-machine-learning/introducing-cloud-tpu-v5p-and-ai-hypercomputer

Introducing AI Hypercomputer with Cloud TPU v5p - YouTube

TPU v5pは、2023年8月30日にリリースされた「TPU v5e」をベースに構築されています。しかし、TPU v5eは演算性能よりも電力効率やコストパフォーマンスが重視され、実際の性能は従来モデルである「TPU v4」に及ばないものとなっていました。

GoogleがAI特化プロセッサTPUの第5世代モデル「TPU v5e」を発表、前モデル比で1ドル当たりのトレーニングパフォーマンス最大2倍・推論パフォーマンス最大2.5倍 - GIGAZINE

一方で今回発表されたTPU v5pは性能特化のTPUで、TPU v4と比べて、TPU v5pは合計8960個ものチップで構成され、1チップ当たりのメモリは95GB、メモリ帯域も毎秒2765GBと大幅に強化されています。Googleは「この性能によって、より高いAI学習のための演算のニーズに応えることが可能です」と述べています。

Googleが公開したTPU v4とTPU v5e、TPU v5pの性能を比較した表が以下。ポッド当たりのチップ数や16ビット浮動小数点演算性能、1秒間に実行できる演算回数などにおいて、従来のTPU v4やTPU v5eの性能をTPU v5pが大幅に上回っていることが確認できます。

また、TPU v5pはint8を利用した学習にも対応しており、Bf16を利用したTPU v4の学習速度を1とすると、Bf16を利用したTPU v5pの学習速度はTPU v4の約1.9倍、int8を用いるとその差は約2.8倍まで開きます。

一方で、費用対効果はTPU v5eの方が優れていて、GPT3-175Bのトレーニングを行った場合、TPU v4の1ドル当たりの相対性能を1とすると、TPU v5pは約2.1倍になるのに対し、TPU v5eでは約2.3倍に向上します。

Google DeepMindとGoogle Researchのチーフサイエンティストであるジェフ・ダン氏は「私たちはTPU v5pを使用した大規模言語モデルのトレーニングにおいて、TPU v4世代のパフォーマンスと比較して約2倍の高速化を確認しました。JAXやPyTorch、TensorFlowなどの機械学習フレームワークとさまざまな自動化ツールを強固にサポートすることで、従来のモデルよりもさらに効率的にスケーリングすることが可能になっています」と述べています。

GoogleはTPU v5pについて「実際に、TPU v5pと同時に発表されたGoogleの最も高性能なAIモデルであるGeminiは、TPU v5pでトレーニングされました」と報告。

文字・音声・画像を同時に処理して人間以上に自然なやりとりができるGPT-4を超える性能のマルチモーダルAI「Gemini」がリリースされる - GIGAZINE

ダン氏は「TPU v5pは、マルチモーダルAIのGeminiのような最先端のAIモデルを使って、研究やエンジニアリングに取り組むために不可欠な存在です」と語っています。

なお、GoogleはTPU v5pの提供時期について「アクセスをご希望の方は、Google Cloudアカウントマネージャーまでご連絡ください」と述べています。