Cohere が Google Cloud TPU で言語モデルのトレーニングを加速している方法

※この投稿は米国時間 2022 年 7 月 27 日に、Google Cloud blog に投稿されたものの抄訳です。

この数年間で、大規模言語モデル（LLM）のトレーニングにおける進歩によって、自然言語処理（NLP）は一部の会社しかアクセスできない最先端のテクノロジーから、多くの一般的なアプリケーションの強力なコンポーネントになりました。chatbot からコンテンツの管理や分類まで、モデルが大規模であるほど、より高精度に言語を理解して生成できるというのが NLP の原則です。

しかし、より大規模で強力な言語モデルの作成を追求するなかで、スケールが主要な課題になっています。モデルが大規模になりすぎて 1 つのデバイスに適合できない場合、分散トレーニング戦略が必要です。その戦略には、膨大なメモリ容量と高速な相互接続を備えた多くのコンピューティング リソースが必要となります。また、ハードウェアと時間のリソースを最適化するためには、専門的なアルゴリズムも必要です。

Cohere のエンジニアは、このスケーリングの課題に対するソリューションに取り組んでおり、すでに成果を挙げています。Cohere では、この種のプロジェクトで一般的に必要となるインフラストラクチャや ML の深い専門知識がなくても、強力な大規模言語モデルを扱えるプラットフォームをデベロッパーに提供しています。新しい技術論文である「Scalable Training of Language Models using JAX pjit and TPUv4（JAX pjit と TPUv4 を用いた言語モデルのスケーラブルなトレーニング）」では、Google Cloud が最近発表した Cloud TPU v4 Pod にデプロイされた新しい FAX フレームワークによって、大規模言語モデルを数千億のパラメータにスケーリングする際の課題に対処する方法を、Cohere のエンジニアが説明しています。具体的には、Cohere によりテンソル並列処理とデータ並列処理を通して達成されたトレーニング効率の飛躍的進歩が、そのレポートに記載されています。

このフレームワークは、スケーラビリティとプロトタイピングの 2 つを大きく改善して、大規模言語モデルの研究と開発、本番環境を加速させることを目的としています。Cohere は、より大規模なモデルをさらに迅速にトレーニングすることでモデルを改善し、顧客にすばやく最適なモデルを提供できるようになります。また、そのフレームワークでは、新しいアイデアの試行とテストを行うことで、カスタマー サービス chatbot を動かす生成モデルを作成するなど、特定の目標に取り組むモデルの迅速なプロトタイピングもサポートします。モデルタイプ間の切り替えと異なる目標に対する最適化が可能であることで、最終的に Cohere は特定のユースケースに最適化されたモデルを提供できます。

FAX フレームワークは、JAX のパーティション化ジャスト イン タイム コンパイル（pjit）機能（デバイスとワークロードの関係を抽象化）に大きく依存しています。これによって、Cohere のエンジニアは、対象のタスクにとって理想的な構成でデバイスとプロセスを合わせることで、効率とパフォーマンスを最適化できます。pjit では、任意の関数を 1 つのプログラムにコンパイルします（XLA コンピューティング）。そのプログラムは、異なるホストに存在するデバイスを含む複数のデバイスで実行できます。

また、Cohere の新しいソリューションでは、Google Cloud の新しい TPU v4 Pod を利用して、テンソル並列処理を行っています。この処理は、これまでのパイプラインの並列処理を実装するよりも効率的です。その名のとおり、パイプライン並列アプローチは、1 つの長い組み立てラインのように、直線的にアクセラレータを使ってワークロードをスケールします。アクセラレータはデータの各マイクロバッチを次に渡す前に処理し、その後、バックワード パスを逆の順番で実行しなければいけません。

テンソル並列処理により、パイプラインの並列処理で発生するアクセラレータのアイドル時間（いわゆるパイプライン バブル）がなくなります。テンソル並列処理では、アクセラレータ全体で大規模テンソル（段落内の単語といった複数のオブジェクト間の関係を定義する数学的配列）をパーティショニングすることで、複数のデバイスで同時にコンピューティングを行います。パイプラインの並列処理が終わることのない組み立てラインであれば、テンソル並列処理は一連の並列な組み立てラインです。たとえば、一方はエンジンを作り、もう一方はボディを作るなどして、同時に組み合わせることでひとつの完全な車を短時間で作ることができます。

大規模言語モデルに Google Cloud を選ぶ理由

長年の技術パートナーシップの一部として、Cohere は Google Cloud の高度な AI と ML インフラストラクチャを活用して、プラットフォームを動かしています。Cohere は、大規模 ML に合わせて最適化された Google Cloud のカスタム設計の機械学習チップである、Cloud TPU でプロダクトを開発およびデプロイしています。Cohere は最近、Google Cloud TPU で FAX を使って大規模言語モデルをトレーニングすることによる新しいモデルの改善とスケーラビリティを発表しました。このモデルでは、TPU v3 から TPU v4 へ移行することにより、これまでで合計 1.7 倍の高速化を達成しました。パフォーマンスの大幅な向上に加えて、新しい TPU VM アーキテクチャにより、優れたユーザー エクスペリエンスが実現しています。重要な点として、Google Cloud を使い、世界最大の一般利用可能な ML ハブにおける 90% カーボンフリー エネルギーを活用することで、Cohere の最先端の ML トレーニングが高い水準のサステナビリティで達成されています。

Cloud TPU を導入することで、Cohere は大規模言語モデルのトレーニングをより迅速かつ経済的で、アジャイルにしています。このことが、より大規模で高精度な大規模言語モデルをデベロッパーに提供し、あらゆる規模のデベロッパーと企業が NLP テクノロジーを利用できるようにするために役立っています。

- Cohere、機械学習エンジニア Joanna Yoo 氏
- Google Cloud、プロダクト マネージャー Vaibhav Singh