Nvidia 3090が180億パラメータの大規模モデルに単独で挑む。今度は国内オープンソースプロジェクトが大暴れ

この記事はAI新メディアQuantum Bit（公開アカウントID：QbitAI）より許可を得て転載しています。転載の際は出典元にご連絡ください。

何？単一の GPU で大規模なモデルをトレーニングできますか?

それとも20シリーズで入手できるものなのでしょうか？ ? ?

冗談ではなく、事実はすでに私たちの目の前にあります。

• RTX 2060 6GB の通常のゲーミング GPU は、15 億のパラメータ モデルをトレーニングできます。
• RTX 3090 24GB ホストは、180 億パラメータの大規模モデルに直接挑戦します。
• Tesla V100 32GB は240 億のパラメータも処理できます。

PyTorchや業界主流のDeepSpeed方式と比較すると、パラメータ容量を10倍以上増やすことができます。

さらに、この方法は完全にオープンソースであり、数行のコードだけで実行でき、必要な変更もごくわずかです。

この一連の作戦により、大規模モデルのトレーニングの閾値が実に半分にまで下がります。Lao Huang は大きな損失を被ることになるのではないでしょうか。

では、このような大物の名前を作った大ボスは誰でしょうか?

国産オープンソースプロジェクトColossal-AIです。

オープンソース化されて以来、GitHub のホットリストで何度もトップにランクされています。

△オープンソースアドレス: https://github.com/hpcaitech/ColossalAI

主なタスクは、GPT-2、GPT-3、ViT、BERTなどのさまざまな大規模モデルのトレーニングを高速化することです。

たとえば、ViT-Base/32 の事前トレーニングには約 30 分、15 億パラメータの GPT モデルのトレーニングには 2 日、83 億パラメータの GPT モデルのトレーニングには 5 日かかります。

GPUも節約できます。

たとえば、GPT-3 のトレーニング時に使用される GPU リソースは、NVIDIA の Megatron-LM の半分だけになります。

そこで今回は、どのようにして単一の GPU で数百億のパラメータを持つ大規模なモデルをトレーニングできるのでしょうか?

原理を深掘りしてみました〜

GPU+CPU異種メモリの効率的な使用

単一のコンシューマーグレードのグラフィック カードで大規模な AI モデルをトレーニングすることが難しいのはなぜですか?

ビデオメモリの制限が最大の難点です。

最近は大型モデルが人気で効果も抜群です。試してみたくない人はいないのではないでしょうか。

しかし、頻繁に発生する「CUDA メモリ不足」エラーは本当に耐え難いものです。

現在、業界で主流となっているのは、マイクロソフトDeepSpeedが提案するZeROである。   (ゼロ冗長オプティマイザー)。

その主な原則は、モデルを分割し、モデル メモリを単一の GPU に均等に分散することです。

データの並列性が高くなるほど、GPU のメモリ消費量は少なくなります。

このアプローチでは、CPU と GPU メモリ間のモデル データの静的パーティション分割のみが使用され、メモリ レイアウトはさまざまなトレーニング構成に対して一定です。

これにより、2つの問題が発生します。

まず、GPU または CPU のメモリが対応するモデル データ要件を満たすのに不十分な場合、他のデバイスに使用可能なメモリがあってもシステムがクラッシュします。

第二に、細粒度テンソルが異なるメモリ空間間で転送される場合、通信効率は非常に低くなります。モデルデータを事前にターゲットコンピューティングデバイスに配置できる場合は、CPU-GPU 通信は不要です。

現在、DeepSpeed の改良版が数多く登場しており、コンピューターのハードドライブを使用してモデルを動的に保存することを提案しています。ただし、ハードドライブの読み取りおよび書き込み速度はメモリやビデオメモリよりも大幅に低く、トレーニング速度は依然として低下します。

これらの問題に対処するために、Colossal-AI はGPU + CPU の異種メモリを効率的に活用するソリューションを提供します。

具体的には、ディープラーニング ネットワーク トレーニングの継続的な反復特性を活用し、反復回数に応じてトレーニング プロセス全体をウォームアップと正式なトレーニングの 2 つの段階に分割します。

ウォームアップフェーズでは、非モデルデータメモリ情報が監視および収集されます。

正式な段階では、収集された情報に基づいて、コンピューティング デバイス上の次の演算子に必要なピーク メモリが予約され、一部の GPU モデル テンソルが CPU メモリに移動されます。

ロジックは次のとおりです。

ここでもう少し詳しく説明すると、モデル データはパラメーター、勾配、オプティマイザーの状態で構成され、それらのフットプリントはモデル構造の定義に関連しています。

非モデル データは、オペレーターによって生成された中間テンソルで構成され、トレーニング タスクの構成 (バッチ サイズなど) に基づいて動的に変化します。

よく行われるのは、GPU ビデオメモリを取得することです。

したがって、CPU は GPU ビデオ メモリが不足しているときに支援すると同時に、他の状況でのメモリの浪費を回避できる必要があります。

これが、Colossal-AI が GPU + CPU の異種メモリを効率的に使用するロジックです。

上記のプロセスでは、実際には非モデルデータのメモリ使用量を取得することは非常に困難です。

非モデル データのライフ サイクルはユーザーによって管理されないため、既存のディープラーニング フレームワークでは、非モデル データ追跡インターフェイスがユーザーに公開されません。次に、CUDA コンテキストなどのフレームワーク以外のオーバーヘッドも考慮する必要があります。

ここで、Colossal-AI のソリューションは、ウォームアップ フェーズでサンプリングを使用して、非モデル データによる CPU と GPU メモリの使用状況を取得することです。

簡単に言えば、これは加算と減算の演算です。

非データモデル使用量 = 2 つの統計モーメント間の最大システムメモリ使用量 - モデルデータメモリ使用量

モデル データのメモリ使用量は、クエリ マネージャーを通じて取得できることがわかっています。

具体的には、次のようになります。

すべてのモデル データ テンソルはメモリ マネージャーによって管理され、各テンソルには HOLD、COMPUTE、FREE などのステータス情報がマークされます。

次に、動的に照会されたメモリ使用量に基づいて、テンソルの状態が継続的かつ動的に変換され、テンソルの位置が調整されて、GPU ビデオメモリと CPU メモリがより効率的に使用されます。

ハードウェアが非常に限られている場合に、モデル容量を最大化し、トレーニング速度のバランスをとります。これは、AI の普及と、大規模モデルの下流タスクの低コストの微調整にとって大きな意味を持ちます。

そして最も重要なことは、メモリを追加する方が、ハイエンドのグラフィック カードを購入するよりもはるかにコスト効率が高いということです。

少し前に、Colossal-AI は Google の最新の研究成果 PaLM (Pathways Language Model) の再現にも成功し、そのパフォーマンスも非常に良好でした。ただし、Microsoft DeepSpeed は現在 PaLM モデルをサポートしていません。

Colossal-AI は他に何ができるのでしょうか?

前述したように、Colossal-AI はトレーニングの高速化や GPU リソースの節約など、多くのタスクに取り組むことができます。

それで、これはどうやって実現されるのでしょうか?

簡単に言えば、Colossal-AI は複数の並列手法を統合したシステムであり、多次元並列処理、大規模オプティマイザー、適応型タスク スケジューリング、冗長メモリの排除などの機能を提供します。

現在、Colossal-AI をベースとした高速化ソリューションである FastFold は、タンパク質構造予測モデル AlphaFold のトレーニング時間を、当初の 11 日からわずか67 時間に短縮できます。

さらに、総コストが低くなり、長いシーケンスの推論タスクでは 9 ～ 11.6 倍の速度向上が達成されます。

このアプローチは、Google やコロンビア大学の方法を上回ることに成功しました。

さらに、Colossal-AI は半分の数の GPU のみを使用して GPT-3 をトレーニングできます。

NVIDIA のソリューションと比較すると、Colossal-AI はトレーニングを開始するために必要なコンピューティング リソースが半分しかなく、同じコンピューティング リソースを使用した場合、速度を 11% 向上させることができ、GPT-3 トレーニングのコストを 100 万ドル以上削減できます。

同時に、Colossal-AI はオープンソース コミュニティの構築にも大きな注意を払っており、中国語のチュートリアルを提供し、ユーザー コミュニティ フォーラムを開設し、全員のニーズとフィードバックに基づいて継続的に更新と反復を行っています。

たとえば、以前、ある読者が、Colossal-AI が一般的な消費者向けグラフィック カードで実行できれば素晴らしいだろうというメッセージを残しました。

数か月後、すでに手配が完了していました。

舞台裏：LAMBオプティマイザーの作者You Yangが主導

これを読んで、Colossal-AI は本当に注目に値すると思いますか?

実際、この国内プロジェクトを支える研究開発チームは非常に素晴らしいです。

リーダーは、LAMB オプティマイザーの考案者である You Yang 氏です。

彼は清華大学のコンピュータサイエンス学部に修士課程で首席で入学し、その後カリフォルニア大学バークレー校に進学してコンピュータサイエンスの博士号を取得しました。

彼は、IPDPS/ICPP 最優秀論文賞、ACM/IEEE ジョージ・マイケル HPC フェローシップ、Forbes 30 Under 30 (Asia 2021)、IEEE-CS スーパーコンピューティング優秀新人賞、UC Berkeley EECS Lotfi A. Zadeh 優秀大学院賞を受賞しています。

Google でのインターンシップ中に、You Yang 氏は LAMB メソッドによる BERT 事前トレーニングの世界記録を更新しました。

NVIDIA の公式 GitHub によると、LAMB は Adam オプティマイザーよりも 72 倍高速です。 Microsoft の DeepSpeed も LAMB アプローチを採用しています。

2021年、楊有氏は中国に戻り、分散型ソフトウェアシステム、大規模人工知能プラットフォーム、エンタープライズレベルのクラウドコンピューティングソリューションを主な事業とするAIスタートアップ企業、 Luchen Technologyを設立しました。

チームの中心メンバーは、カリフォルニア大学バークレー校、ハーバード大学、スタンフォード大学、シカゴ大学、清華大学、北京大学、シンガポール国立大学、シンガポール南洋理工大学など国内外の有名大学の出身者であり、Google Brain、IBM、Intel、Microsoft、NVIDIAなどの有名メーカーでの勤務経験を持っています。

同社は設立直後から、Innovation WorksやZhenFundなど複数のトップVC機関からシードラウンドの投資を受けた。

Luchen CSO の James Demmel 教授は、カリフォルニア大学バークレー校の著名な教授であり、ACM/IEEE フェローであり、米国科学アカデミー、米国工学アカデミー、米国芸術科学アカデミーの会員です。

ポータル:

https://github.com/hpcaitech/ColossalAI