GPT-NeoX 20Bの実戦投入とDeepspeedへのこだわり

GPT-NeoX 20B

GPT-NeoX 20B」は、オープンソースの自然言語処理モデルとしては、現時点で最大のものです。

2週間前にEleutherAIによってリリースされました。このモデルは200億個のパラメータを使って学習されており、1年間の苦労の末に完成したものです。彼らの主な課題の一つは、モデルを複数のGPUで効率的に学習できるように、適切な並列アーキテクチャを設定することだったようです。

彼らの前モデルであるGPT-Jは、GoogleのTPUで学習させたもので、TPUで使うことを前提に作られています。GPT-NeoX 20Bの内部アーキテクチャは根本的に異なり、GPUで使用することを前提に作られています。このモデルは、さらに大きなモデルのシリーズの最初のものになるはずです（ここでの目標は、GPT-3に追いつくことです）。

GPT-NeoX 20はテキスト生成モデルです。つまり、あなたに代わってテキストを書くことができ、ブログ記事生成、チャットボット、テキスト分類、感情分析、キーワード抽出、コード生成、エンティティ抽出、意図分類、質問応答など、ほとんどすべての自然言語処理のユースケースを高い精度で実現します...。

多くの研究者や企業がこの新しいモデルを活用し、大きな成果を得ることは間違いありません。しかし、GPT-NeoX 20Bをどのように導入するかという課題があります。

必要なハードウェア

GPT-NeoX 20Bを実運用で展開しようとすると、まず直面するのが、高度なハードウェアを必要とすることです。

GPT-Jは、たとえ非常に高価であっても、民生用ハードウェアに導入することが可能でした。例えば、非常に優れたCPU（結果的に痛みを伴うほど遅くなったとしても）や、NVIDIA RTX 3090のような高度なゲーミングGPU上に展開することができました。しかし、GPT-NeoX 20Bは非常に大きいので、それはもうできません。

基本的にGPT-NeoXは、最低でも42GBのVRAMと40GBのディスクスペースを必要とします（ここでは、スリムなFP16バージョンの話をしています）。この条件を満たすGPUはほとんどありません。主なものとしては、NVIDIA A100、A40、RTX A6000などがあります。

これらのGPUは非常に高価であるだけでなく、世界的な半導体不足のため、最近では手に入れることが難しくなっています。

ここでの最善の解決策は、マルチGPUアーキテクチャを採用することです。

DeepspeedによるマルチGPUの実現

モデルが1つのGPUに収まらない場合は、複数のGPUに分割してみましょう。

これは2つの理由から興味深い選択肢です。まず、インフラを水平方向に拡張する方が、垂直方向に拡張するよりもコストが低くなることがよくあります。これはサーバーでも同じですが、GPUでも同じことが言えます。第二に、NVIDIA A40のような大きなGPUよりも、NVIDIA Tesla T4のような小さなGPUをいくつか手に入れる方が簡単です。

残念ながら、機械学習モデルを複数のGPUに分割することは技術的に難しいことです。Hugging Faceがモデルの並列化とその選択肢について素晴らしい記事を書いています。 しかし幸いなことに、EleutherAIはGPT-NeoX 20Bを、学習用であれ推論用であれ、複数のGPUで簡単に並列化できるように設計しました。

GPT-NeoXのアーキテクチャはDeepspeedをベースにしています。Deepspeedはマイクロソフトが提供するフレームワークで、もともとは複数のGPU間で学習を並列化するために作られたものですが、推論にも使われるようになってきました。EleutherAIがすべての作業を行ったので、GPT-NeoXの並列化は設定ファイルの数字を変更するだけで簡単に行うことができます。

Extract from the GPT-NeoX docs about parallelism

しかし、APIを介したオンライン推論のためにDeepspeedを本番環境に導入するのは難しい...。DeepspeedはWebサーバに簡単に統合できるようには設計されていません。NLPクラウドでは、それを苦労して学びましたが、苦労の末、DeepspeedとGPT-NeoXを私たちのAPIの後ろで動作させる効率的な方法をようやく考え出すことができました。