GoogleのAI用プロセッサ「TPU v4」はNVIDIAの「A100」より高速で効率的だとGoogleの研究者が主張

Googleは2021年、機械学習に特化したプロセッサ「Tensor Processing Unit(TPU)」の第4世代モデルである「TPU v4」を発表しました。新たにGoogleが、2023年4月に公開した公式ブログとプレプリントサーバーのArXivに投稿した論文で、「GoogleのTPU v4を組み合わせたシステムは同サイズのNVIDA『A100』を用いたシステムより高速かつ効率的」と主張しています。

[2304.01433] TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings
https://arxiv.org/abs/2304.01433

TPU v4 enables performance, energy and CO2e efficiency gains | Google Cloud Blog
https://cloud.google.com/blog/topics/systems/tpu-v4-enables-performance-energy-and-co2e-efficiency-gains

Google Claims Its TPU v4 Outperforms Nvidia A100
https://www.datanami.com/2023/04/05/google-claims-its-tpu-v4-outperforms-nvidia-a100/

Google reveals its newest A.I. supercomputer, says it beats Nvidia
https://www.cnbc.com/2023/04/05/google-reveals-its-newest-ai-supercomputer-claims-it-beats-nvidia-.html

機械学習モデルの開発では大量のデータを訓練することが必要なため、モデルがより大きく複雑になるにつれ、コンピューティングリソースの需要も増えています。Googleが開発するTPUは機械学習やディープニューラルネットワークに特化したプロセッサであり、TPU v4の性能は前世代のTPU v3を2.1倍上回っており、電力あたりの性能も2.7倍向上しているとGoogleは主張しています。

GoogleはクラウドコンピューティングサービスのGoogle Cloudを通じて、TPU v4を4096個組み合わせたスーパーコンピューターの「TPU v4 Pod」を提供しています。TPU v4 Podは独自のoptical circuit switches (OCS／光回路スイッチ)を介して4096個のTPU v4を相互接続したもので、パフォーマンス・スケーラビリティ・可用性などに優れており、LaMDA・MUM・PaLMといった大規模言語モデルにおいて主力製品となっているとのこと。

以下の写真は、TPU v4 Podの一部を撮影したもの。GoogleはTPU v4 Podで使用するOCSについて、競合する相互接続テクノロジーであるInfiniBandと比較してはるかに安価・低電力・高速であり、OCSがTPU v4 Podのシステムコストおよびシステムパワーに占める割合はわずか5％未満だと述べています。

機械学習用スーパーコンピューターにOCSを用いる利点としては、回路を切り替えて簡単に故障したコンポーネントを迂回(うかい)できるため、可用性に優れているという点も上げられます。GoogleのBardやOpenAIのChatGPTといった大規模なAIモデルや製品をトレーニングするには、大量のチップを数週間～数カ月にわたり稼働させる必要があるため、可用性は機械学習用スーパーコンピューターにとって重要です。

Googleは論文で、TPU v4 PodはNVIDIAのA100を用いた同規模のシステムよりも1.2～1.7倍も高速で、消費電力量も1.3～1.9倍少ないと主張しています。また、Google Cloudに最適化されたTPU v4 Podは、一般的なオンプレミスデータセンターにおける現代のドメイン特化アーキテクチャ(DSA)と比較して、エネルギー消費量は2～6倍、CO2排出量は20倍少なくなっているとのこと。

なお、NVIDIAの最新AI用チップである「H100」と比較していない理由についてGoogleは、H100がTPU v4よりも高度な4nmプロセスで作られたものであり、可用性も限られているためだとしています。

TPU v4 Podは画像生成AI「Midjourney」のバージョン4のトレーニングにも用いられました。Midjourneyの創業者兼CEOを務めるデビッド・ホルツ氏は、「Google Cloudと協力し、Googleのグローバルでスケーラブルなインフラストラクチャーを活用したシームレスな体験を、クリエイティブコミュニティに提供できることを誇りに思います」「最新のTPU v4によるGoogle JAXを使ったバージョン4のアルゴリズムトレーニングからGPUでの推論実行まで、ユーザーの活気に満ちたアイデアをTPU v4が実現するスピードに感動しています」と述べました。

なお、NVIDIAは2023年4月5日のブログで、サードパーティーのベンチマーク「MLPerf」が公開したベンチマークテストのデータを引用し、「H100」のパフォーマンスが前世代の「A100」を4倍上回ったと報告しています。

H100, L4 and Orin Raise the Bar for Inference in MLPerf | NVIDIA Blogs
https://blogs.nvidia.com/blog/2023/04/05/inference-mlperf-ai/