15億パラメータのモデルを2日間でトレーニングし、国内オープンソースプロジェクトがNvidiaのMegatron-LMを上回った

AIの現在の動向において、その徹底的な発展に影響を与える矛盾は何でしょうか?

一方では、大型モデルが大流行しており、その結果は驚くべきものであり、誰もがそれを試してみたいと思っています。しかし一方で、ハードウェア基盤は昼夜を問わず稼働する数万の GPU からなる大規模クラスターであることが多く、人々がお金を支払うのを思いとどまらせるには十分ではありません。

では、同じ GPT-3 トレーニングを GPU の数を半分に減らして完了できるとしたらどうでしょうか?

鍵は何だと思いますか?

あなたを不安にさせておくつもりはありません。この改善は、Colossal-AIと呼ばれる GitHub オープンソースプロジェクトを通じて実現されました。

さらに、プロジェクトがオープンソース化されてから間もなく、Python の世界のホットリストですぐにトップになりました。

↑GitHubアドレス: https://github.com/hpcaitech/ColossalAI

Colossal-AI は GPT-3 を高速化するだけでなく、GPT-2、ViT、BERT などの他のモデルでも非常に優れたパフォーマンスを発揮します。

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

業界で主流の AI 並列システムである NVIDIA Megatron-LM と比較すると、同じ 512 個の GPU を使用して GPT-2 モデルをトレーニングする場合、Colossal-AI の加速率は NVIDIA Megatron-LM の2 倍になります。 GPT-3 をトレーニングする場合、トレーニングコストを約 1,000 万元節約できます。

さらに、GPT-2 をトレーニングする場合、ビデオメモリの消費量を Megatron-LM の10 分の 1以下に抑えることもできます。

Colossal-AI はどうやってそれを実現したのでしょうか?

いつものように、紙から始めます。

効率的な6次元並列方式

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

まず多次元並列処理について見てみましょう。

いわゆる「多次元」とは、現在主流の分散並列ソリューションでは複数の並列方式が使用されることが多いことを意味します。

たとえば、NVIDIA の Megatron-LM では、データ並列処理、パイプライン並列処理、テンソル並列処理という3 つの方法が使用されます。そのため、このモードは3次元並列とも呼ばれます。 Microsoft の DeepSpeed は、Megatron-LM を並列基盤と呼んでいます。

Colossal-AI は、システムの並列次元を一度に6 次元まで増やすことができます。

データ並列性とパイプライン並列性の互換性を前提とし、プロジェクトチームが開発した2D/2.5D/3D テンソル並列方式とシーケンス並列実装をベースとしています。

その中で、高次元テンソル並列処理は、Colossal-AI が大規模モデルに対するメモリ使用率と通信効率を向上させる鍵となります。

実際、テンソル並列処理は新しいものではありませんが、過去によく見られるテンソル並列処理は主に1 次元に基づいています。

その原理は、モデルレイヤー内の重みパラメーターを行または列に分割して異なるプロセッサに割り当て、ブロックマトリックス乗算を使用して操作を複数のプロセッサに分散して同時に実行することです。

たとえば、Nvidia の Megatron-LM が典型的な例です。

ただし、この並列アプローチにはいくつかの欠点があります。

たとえば、各プロセッサは依然として中間アクティベーション全体を保存する必要があり、大規模なモデルを処理するときに大量のビデオメモリ領域が浪費されます。

一方、この単線方式では、各プロセッサが他のすべてのプロセッサと通信する必要が生じます。

つまり、GPU が 100 個ある場合、各 GPU は他の 99 個の GPU と通信する必要があり、計算ごとに最大9,900 回の通信が必要になります。

しかし、テンソル並列処理の次元を2 次元に拡張すると、単一の計算量はすぐに 1 桁減少します。

各 GPU は、自身の行または列内の GPU とのみ通信する必要があるためです。

同じく GPU が 100 個のケースでは、各 GPU が通信する必要がある GPU の数は9 個に削減でき、1 回の計算には900 回しかかかりません。

実際、これに基づいて、Colossal-AI には 2.5 次元および 3 次元テンソル並列方式も組み込まれており、これにより伝送コストをさらに削減できます。

2D並列方式と比較すると、2.5D並列方式では効率が1.45倍向上し、3D方式では効率が1.57倍向上します。

大きな画像、ビデオ、長いテキスト、長期の医療モニタリングなどのデータに対しても、Colossal-AI はシーケンス並列方式を採用しており、元のマシン機能の限界を打ち破り、長いシーケンスデータを直接処理できます。

Colossal-AI の API インターフェースはカスタマイズ可能であり、新しい並列ディメンションを簡単に追加できることは特筆に値します。

第二に、大規模オプティマイザーも Colossal-AI のハイライトです。

前述したように、分散並列システムではさまざまな並列方式が使用されており、データ並列化も一般的な方式の 1 つです。

この方法の原理は理解するのが難しくありません。トレーニングデータをいくつかの部分に分割し、異なるマシンに異なるデータを計算させ、パラメータサーバー（Paremeter Server）を介してターゲットデータを収集するというものです。

これにより、AI モデルのトレーニング中にバッチサイズが大幅に増加し、トレーニングプロセスが高速化されます。

しかし、大規模なトレーニングには、一般化エラー（一般化ギャップ）が発生し、ネットワークの一般化能力が低下し、AI モデルの精度が低下するという共通の問題があります。

そのため、Colossal-AI では、LAMB や LARSなどの自社開発の大規模最適化ツールをシステム内で使用しています。トレーニング精度を確保しながら、バッチサイズを512から65536に拡張します。

その中で、LARS オプティマイザーは、学習率を層ごとに調整し、学習率によって発生する非収束を軽減します。

LAMB オプティマイザーは LARS に基づいており、学習率を層ごとに調整するという考え方を適応勾配に適用します。

その結果、LAMBはBERTトレーニングにおけるLARSの違いの問題を解決でき、最大バッチサイズは64Kに達します。

これまで、LAMB オプティマイザーは、BERT の事前トレーニングにかかる時間を 3 日 3 晩から 1 時間以上に短縮することに成功していました。

3 番目に、Colossal-AI は適応型でスケーラブルなスケジューラを使用して、タスクを効率的に処理します。

Colossal-AI は、既存の一般的なタスクスケジューラとは異なり、GPU の数に基づいてタスクの規模を静的に決定するのではなく、バッチサイズに基づいて各タスクを動的かつ自動的に管理します。

タスクスケジューラは進化的アルゴリズムを通じて、スケジューリングの決定を継続的に最適化し、GPU の使用率を大幅に向上させることもできます。

評価結果によると、現在の最先端の方法と比較して、この方法は平均JCT（ジョブ完了時間）を45.6％短縮することができ、これは既存のディープラーニングタスクスケジューリングアルゴリズムよりも優れています。

さらに、この適応型でスケーラブルなスケジューラは、NCCL ネットワーク通信を通じて効率的なタスク移行も実現できます。

最後に、冗長メモリを排除することも、AI トレーニングを高速化するソリューションです。

この点に関して、Colossal-AI はゼロ冗長性オプティマイザーテクノロジー (略して ZeRO) を使用します。

この方法は、主に、オプティマイザーの状態、勾配、およびモデルパラメーターを分割して、GPU が現在の計算に必要な部分のみを保存するようにすることで、データ並列処理とモデル並列処理におけるメモリの冗長性を排除します。

特にモデル推論を展開する場合、ゼロオフロードを使用してモデルを CPU メモリまたはハードディスクにオフロードし、少量の GPU リソースのみを使用して最先端の AI 大規模モデルを低コストで展開できます。

以上のことから、Colossal-AI の加速効果が技術的なレベルで非常に明白であることは容易に理解できます。

アプリケーションレベルでは、Colossal-AI の設計ではエネルギー消費と使いやすさも考慮されています。

データの移動がエネルギー消費の主な原因であることを考慮して、Colossal-AI は計算量を増やすことなくデータ移動を最小限に抑えることでエネルギー消費を削減します。

一方、Colossal-AIは誰でも使えるオープンソースシステムなので、利用の敷居が低く、分散システムを勉強したことがない人でも始めることができます。

同時に、Colossal-AI は最小限のコード変更だけで、既存の単一マシンコードを並列コンピューティングクラスターに迅速に拡張できます。

GPT-3のトレーニング速度が10.7%向上

Nvidia の Megatron-LM では、GPT-3 のトレーニングを高速化するために、少なくとも 128 個の GPU が必要です。下の表に示すように、同じコンピューティングリソースを使用すると、Colossal-AI は各反復に要する時間を 43.1 秒から 38.5 秒に短縮できます。

これは、Colossal-AI が GPT-3 のトレーニング速度をさらに 10.7% 向上できることを意味します。

エンジニアリングの観点から見ると、このような大規模なモデルのトレーニングには数百万ドルかかることが多いことを考えると、この増加のメリットは明らかです。

さらに、Colossal-AI はシステム最適化により、トレーニング速度をあまり落とさずに GPU の数を 128 から 96 に削減 (43.1 → 48.5) し、トレーニングコストを大幅に削減できます。

さらに ZeRO3 (Zero Redundancy Optimizer) を有効にすると、必要な GPU の数を半分の 64 に減らすことができます。

GPT-2のトレーニングは2日間で完了します

GPT-2 の高速トレーニング結果では、4、16、64 GPU のいずれの場合でも、Colossal-AI が占有するビデオメモリが Megatron-LM と比較して大幅に削減されていることがわかります。

つまり、Colossal-AI を使用すると、エンジニアは同じ数の GPU を使用しながら、より大きなモデルをトレーニングしたり、より大きなバッチサイズを設定してトレーニングを高速化したりすることができます。

下の表の結果からも、バッチサイズが大きくなるにつれて、Colossal-AI のリソース使用率がさらに向上し、Megatron-LM の2 倍の速度に達することがわかります。

研究開発チームは 256 個の GPU で実験を行い、最終的に 82.8 時間で GPT-2 の 15 億パラメータバージョンのトレーニングを完了しました。

この推定に基づくと、Colossal-AI は 512 個の GPU で GPT-2 事前トレーニングを実行することで、トレーニング時間を45 時間に短縮できます。

複数の並列モードに完全対応

BERT で実施された実験は、世界で最も並列次元を備えたAI トレーニングシステムとしての Colossal-AI の利点を実証しています。

Megatron-LM と比較すると、Colossal-AI の順次並列アプローチでは必要なビデオメモリが少なくなり、より大きなバッチサイズを使用してトレーニングを高速化できます。同時に、開発者はより長いシーケンスデータを使用できるようになります。

Colossal-AI の順次並列アプローチは、パイプライン並列アプローチとも互換性があります。開発者が順次並列処理とパイプライン並列処理の両方を使用すると、大規模なモデルのトレーニングにかかる時間をさらに節約できます。

さらに、Colossal-AI は、最近の学術的にホットなViT モデルにおいて、高次元テンソル並列法の利点も実証しました。

Colossal-AI は、64 個の GPU を使用する場合、2/2.5 次元テンソル並列処理を採用し、より大きなバッチサイズを最大限に活用して、より高速な処理速度を実現します。

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

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

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

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

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

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

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

2020年に博士号を取得後、シンガポール国立大学のコンピューターサイエンス学部に入学し、2021年1月に大統領若手教授に就任しました。

また、2021年にはIEEE-CSスーパーコンピューティング優秀新人賞を受賞しました。この賞は、毎年世界中で最大 3 名を表彰するもので、博士号取得後 5 年以内に高性能コンピューティングの分野で影響力のある顕著な貢献を果たし、高性能コンピューティングの発展に長期的に貢献できる優秀な若手学者にのみ授与されます。

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

コアチームのメンバーは、カリフォルニア大学バークレー校、スタンフォード大学、清華大学、北京大学、シンガポール国立大学、シンガポール南洋理工大学など、国内外の著名な大学の出身者です。彼らは、高性能コンピューティング、人工知能、分散システムにおいて 10 年以上の技術蓄積を持ち、トップクラスの国際学術誌や会議で 30 本以上の論文を発表しています。

現在、Luchen Technologyは、Innovation WorksとZhen Fundの共同投資による1,000万人民元を超えるシードラウンドの資金調達を獲得しています。