当涉及训练大型 AI 模型时，人们会想到使用数千个 GPU，昂贵的训练成本，只有少数几家科技巨头能够承受。而 AI ユーザー、例えばスタートアップ企業や大学の研究者は、大型モデルのトレーニングをただ眺めるしかありません。
現在、たった 1 つの GPU を搭載したパーソナルコンピュータでも、180 億のパラメータを持つ GPT をトレーニングすることができます。ノートパソコンでも、10 億以上のパラメータを持つモデルをトレーニングすることができます。既存の主流のソリューションと比較して、パラメータ容量は 10 倍以上増加することができます！
このような大幅な改善は、Colossal-AI によるものであり、これは汎用の大型 AI モデルの効率的なトレーニングシステムです。最も重要なことは、完全にオープンソースであり、既存のディープラーニングプロジェクトを最小限の変更で、個人用の消費者向けグラフィックスカード上でより大きなモデルをトレーニングできるようにするだけで実現できます。つまり、誰もが自宅で大型 AI モデルをトレーニングできるようになります！特に、大型 AI モデルのファインチューニングや推論などの下流タスクやアプリケーションのデプロイが容易になります！
Colossal-AI は、さまざまな人気のある効率的な並列処理を提供することで、ユーザーが既存のプロジェクトを簡単に大規模な計算クラスタにデプロイできるよう支援します。
2018 年に Google が提案した BERT モデルの 300 万のパラメータ以来、わずか数年で大型モデルのパラメータ数は何度も更新されています。例えば、OpenAI が提案した 1750 億のパラメータを持つ GPT-3 や、Microsoft と NVIDIA が共同で提供する 5300 億のパラメータを持つ MT-NLG などです...
密なモデルは、1000 億を超えるパラメータのスケールを実現していますが、スパースなエキスパートミックスモデル（MoE）は、2021 年に Google が発表した Switch Transformer など、パラメータ数を兆にまで引き上げています。
しかし、このような巨大なモデルをゼロからトレーニングすることは非常に高価になる可能性があります。通常、数百または数千の専用の高性能 GPU（例：NVIDIA A100）を同時に使用する必要があります。専用の InfiniBand 高速ネットワークを使用してスーパーコンピュータクラスタを構築する場合、トレーニングコストは 1000 万ドルに達する可能性もあります。
1 つの消費者向け GPU で大規模なモデルをトレーニング#
明らかなように、大学生や個人開発者などの AI ユーザーは、このような高額なコストを負担することはできませんが、AI コミュニティで最も人気のある計算リソースは NVIDIA RTX GPU です。
AI の生産性を向上させ、より多くの開発者が大型モデルの恩恵を受け、私たちのビジョンを実現し、大型 AI モデルの使用を「高速かつ低コスト」にするために、Colossal-AI はわずか数行のコードでモデルトレーニングの容量を 10 倍に増やすことができます。これにより、誰もが単一の普通の GPU を使用して大型 AI モデルをトレーニングできるようになります。
すべてのハードウェア上で、Colossal-AI は通常の PyTorch や主流の分散ソリューション（例：Microsoft の DeepSpeed）よりも優れたパフォーマンスを発揮します。
大型モデルの代表である GPT に対して、Colossal-AI は RTX 2060 6GB を搭載したゲームノートパソコンで最大 1.5 億のパラメータを持つ GPT をトレーニングすることができます。RTX3090 24GB を搭載したコンピュータでは、Colossal-AI は 180 億のパラメータを持つ GPT のトレーニングをサポートします。Tesla V100 などの高性能グラフィックスカードでも、Colossal-AI は大幅な改善をもたらすことができます。
Colossal-AI は、Google が最近発表した PaLM（Pathways Language Model）も成功裏に実現しました。さまざまなハードウェア上で、Colossal-AI は優れたパフォーマンス向上を示し、Microsoft の DeepSpeed の公式の PaLM 実装はまだリリースされていません。
キーテクノロジー：強化ヘテロジニアストレーニング#
消費者向けの単一の GPU を使用して大型 AI モデルをトレーニングする際の最大の問題は、GPU メモリの容量が非常に限られていることであり、モデルパラメータを収容できる量が制限されてしまうことです。Microsoft の DeepSpeed が提案した ZeRO-offload メソッドは、モデルを分割し、CPU メモリを利用することで、より大きな容量と低コストを実現しようとしています。現在、DeepSpeed をベースにしたいくつかの異種トレーニングの変更版が存在しています。しかし、下図の左側に示すように、GPU メモリが対応するモデルの要件を満たすことができない場合、システムはクラッシュしてしまいます。
Colossal-AI は、DeepSpeed 派生の方法とは異なり、自身のコアテクノロジーである ZeRO を構築し、DeepSpeed の問題を解決しました。DeepSpeed は、CPU と GPU メモリの間でモデルデータを静的に分割し、異なるトレーニング設定に対して固定のメモリレイアウトのみを使用します。Colossal-AI は、GPU と CPU メモリの使用効率を向上させるために多くの改善を行いました。結局のところ、CPU メモリは大容量の高性能グラフィックスカードよりもはるかに安価です。
Colossal-AI が設計した Gemini メカニズムは、GPU と CPU のヘテロジニアスメモリを効果的に管理し活用することで、GPU のメモリ制限を克服します。
私たちは、ディープラーニングネットワークのトレーニングプロセスの反復的な性質を利用し、トレーニングをプレヒートフェーズと非プレヒートフェーズの 2 つのステージに分割しています。反復回数に基づいて分割します。初期のプレヒートフェーズでは、メモリ情報を監視します。非プレヒートフェーズでは、収集した情報を使用してテンソルを効率的に移動し、CPU-GPU データ移動を最小限に抑えます。
これは簡単に聞こえますが、実装は非常に困難です。なぜなら、ユーザーが管理しない非モデルデータのライフサイクルや、既存のディープラーニングフレームワークが非モデルデータのトラッキングインターフェースをユーザーに公開していないためです。さらに、CUDA コンテキストなどのフレームワーク外のオーバーヘッドも考慮する必要があります。
Colossal-AI は、プレヒートフェーズでサンプリングを行うことで、CPU と GPU メモリの使用状況を取得します。また、非モデルデータの使用状況は、2 つの時点のシステムの最大メモリ使用量とモデルのメモリ使用量を比較することで取得できます。メモリマネージャーにクエリを実行することで、モデルのメモリ使用状況を知ることができます。以下の図に示すように、黒い実線で示されています。
これは簡単に聞こえますが、実装は非常に困難です。なぜなら、ユーザーが管理しない非モデルデータのライフサイクルや、既存のディープラーニングフレームワークが非モデルデータのトラッキングインターフェースをユーザーに公開していないためです。さらに、CUDA コンテキストなどのフレームワーク外のオーバーヘッドも考慮する必要があります。
Colossal-AI は、プレヒートフェーズでサンプリングを行うことで、CPU と GPU メモリの使用状況を取得します。また、非モデルデータの使用状況は、2 つの時点のシステムの最大メモリ使用量とモデルのメモリ使用量を比較することで取得できます。メモリマネージャーにクエリを実行することで、モデルのメモリ使用状況を知ることができます。以下の図に示すように、黒い実線で示されています。
すべてのモデルテンソルは、メモリマネージャーによって管理され、各テンソルには HOLD、COMPUTE、FREE などの状態情報が付いています。動的なメモリ使用状況に基づいて、Colossal-AI はテンソルの状態を変更し、テンソルの位置を調整し続けます。結果として、GPU と CPU メモリの効率的な使用が実現され、モデルの容量が最大化され、非常に限られたハードウェア環境でトレーニング速度がバランス良くなります。これは、AI の民主化と大型モデルの低コストなファインチューニングの下流タスクにとって重要です。
さらに、分散並列技術は、モデルトレーニングをさらに高速化するための重要な手法です。Colossal-AI は、並列次元の制約、効率の低さ、汎用性の欠如、デプロイの難しさ、メンテナンスの欠如など、既存のソリューションの痛点に対して、多次元並列やヘテロジニアス並列などの技術を使用して、ユーザーがわずかな修正で大型 AI モデルを効率的かつ迅速にデプロイできるようにしています。
例えば、GPT-3 のような超大型 AI モデルの場合、NVIDIA のソリューションと比較して、Colossal-AI は半分の計算リソースでトレーニングを開始することができます。同じ計算リソースを使用する場合、速度をさらに 11％向上させることができ、GPT-3 のトレーニングコストを 100 万ドル以上削減することができます。
タンパク質構造予測に使用される AlphaFold については、私たちのチームは Colossal-AI ベースの高速 Fold という加速ソリューションを導入しました。FastFold は、Google やコロンビア大学が提案した他のソリューションを上回り、AlphaFold のトレーニング時間を 11 日から 67 時間に短縮し、総コストも削減しました。さらに、長いシーケンスの推論において、私たちは 9.3〜11.6 倍の速度向上を実現しました。
さらに、Colossal-AI はオープンソースコミュニティの構築に非常に重要な価値を置いており、英語と中国語のチュートリアルを提供し、PaLM や AlphaFold などの最新の先端アプリケーションをサポートしています。また、定期的に新しい革新的な機能をリリースしています。私たちは常にコミュニティの提案と議論を歓迎しており、何か問題があればお手伝いします。質問をするか、フォーラムでディスカッションを開始することができます。提案に感謝します。最近、Colossal-AI は Github の人気プロジェクトランキングで 1 位になりました。これは、10K のスターを持つプロジェクトが多数存在する中でのことです。
原作者について
原作者はすべて HPC-AI Tech の中核メンバーであり、カリフォルニア大学バークレー校、スタンフォード大学、清華大学、北京大学、シンガポール国立大学、南洋理工大学などの有名な大学出身であり、Google Brain、IBM、Intel、Microsoft、NVIDIA などのテクノロジージャイアントでの経験もあります。同社は、Innovation Works、真格基金などのトップベンチャーキャピタルからのシードラウンドの資金調達も受けています。
HPC-AI Tech の創設者である楊優教授
博士、カリフォルニア大学バークレー校
IPDPS/ICPP ベストペーパーオーサー
ACM/IEEE CS ジョージ・マイケル記念 HPC フェローシップ
フォーブス 30 アンダー 30（アジア 2021）
IEEE-CS スーパーコンピューティングの優れた新人賞
UC バークレー EECS Lotfi A. ザデ賞
HPC-AI Tech の CSO であるジェームズ・デンメル教授
カリフォルニア大学バークレー校の著名な教授
ACM/IEEE フェロー
アメリカ科学アカデミー、エンジニアリングアカデミー、芸術科学アカデミーのメンバー
プロジェクトの詳細はこちらからご覧ください：https://github.com/hpcaitech/ColossalAI
元の記事（この記事は元の記事を翻訳したものです）：https://medium.com/@hpcaitech/train-18-billion-parameter-gpt-models-with-a-single-gpu-on-your-personal-computer-8793d08332dc