DUICUO

1 台の Nvidia 3090 で 180 億のパラメータを持つ大規模なモデルを処理できます。今回は中国のオープンソース プロジェクトが大活躍しています。

この記事は、AI関連ニュースメディアQuantumBit(WeChat ID: QbitAI)の許可を得て転載しています。転載の許可については、元の情報源にお問い合わせください。

何ですか? 1 つの GPU で大規模なモデルをトレーニングできるようになりましたか?

それとも20シリーズくらいで手に入るものなのでしょうか?

冗談ではなく、事実は明らかです。

  • RTX 2060 6GB グラフィック カードは、一般的なゲーム向けに15 億のパラメータを持つモデルをトレーニングできます。
  • RTX 3090 24GB PC は、 180 億のパラメータを持つ大規模なモデルを単独で処理できます。
  • Tesla V100 32GB は240 億のパラメータも処理できます。

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

さらに、この方法は完全にオープンソースであり、必要なコードは数行だけで、変更もほとんど必要ありません。

この動きにより、大規模モデルのトレーニング閾値は大幅に半分に削減され、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 メモリ不足」エラーが頻繁に発生するのは本当に耐えられません。

現在、業界で主流となっているのは、Microsoft DeepSpeed が提案する ZeRO という方式です。   (ゼロ削減オプティマイザー)。

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

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

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

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

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

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

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

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

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

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

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

一般的なロジックは次のとおりです。

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

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

彼らが頻繁に行うのは、GPU メモリをめぐる競争です。

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

Colossal-AI はこのロジックに基づいて、GPU と CPU の異種メモリを効率的に利用します。

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

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

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

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

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

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

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

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

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

非常に限られたハードウェア条件下でモデル容量を最大化し、トレーニング速度のバランスをとることは、AI の普及と下流のタスクのための大規模モデルの低コストでの微調整にとって非常に重要です。

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

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 Best Paper Award、ACM/IEEE George Michael HPC Fellowship、Forbes 30 Under 30 (Asia 2021)、IEEE-CS Supercomputing Outstanding Newcomer Award、UC Berkeley EECS Lotfi A. Zadeh Outstanding Graduate Award など、数々の賞を受賞しています。

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

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

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

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

同社は設立後すぐに、Innovation WorksやZhenFundを含む複数のトップVC企業からシードラウンドの投資を受けた。

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

ポータル:

https://github.com/hpcaitech/ColossalAI