1兆パラメータM6モデルの事前トレーニングの背後にある分散フレームワークWhaleの解読

[[418050]]

最近、アリババクラウドPAIチームとDAMOアカデミーインテリジェントコンピューティングラボは、巨大モデルM6の「低炭素バージョン」を共同でリリースしました。これにより、数兆個のパラメータを持つ超大規模モデルのトレーニングのエネルギー消費が大幅に削減されます。当社が独自に開発した Whale フレームワークの助けを借りて、わずか 480 枚の GPU カードを使用して、人間のニューロンの 10 倍の大きさである 1 兆パラメータのマルチモーダル大規模モデル M6 をトレーニングしました。従来の海外企業が達成した兆パラメータ規模と比較すると、エネルギー消費量は80％以上削減され、効率は11倍近く向上しました。

M6は中国で商品化された初の大型マルチモーダルモデルです。 M6 は従来の AI を超える認知能力と創造力を備えています。絵を描くこと、書くこと、質疑応答が得意です。電子商取引、製造、文学、芸術など、多くの分野で幅広い応用が期待されています。

ここでは、1兆パラメータのモデルトレーニングをサポートするWhaleフレームワークの設計を紹介します。

1. モデル開発の動向と課題

1. モデル開発動向

ディープラーニングの普及に伴い、モデルのパラメータサイズも急速に増加しました。 OpenAI のデータは次のことを示しています:

• 2012 年以前は、ムーアの法則に従って、モデル計算に必要な時間は 2 年ごとに 2 倍になっていました。
• 2012年以降、モデル計算に必要な時間は3.4か月ごとに倍増し、ハードウェア開発のスピードをはるかに上回りました。

過去 1 年間で、モデル パラメータの規模は急速に拡大しました。 Google、NVIDIA、Alibaba、ARIはいずれも兆単位のパラメータを持つモデルをリリースしており、大手企業の中には数百億、数千億のパラメータを持つモデルをリリースしているところもあります。同時に、モデルパラメータの規模が大きくなるにつれて、モデル効果も徐々に向上しています。 Nvidia は Bert モデルのさまざまなパラメータ スケールをテストし、モデル パラメータのスケールが増加するにつれてモデルの困惑度が減少することを発見しました。

Google は GShard の論文で、MoETransformer モデルのパラメータが大きいほど翻訳の品質が高くなることを発見しました。

2. 大規模モデルトレーニングの課題

大規模なモデルはモデルのパフォーマンスを向上させますが、トレーニング フレームワークに大きな課題をもたらします。たとえば、1兆規模のモデルをトレーニングする場合、次のような課題に直面します。

• トレーニングの難しさ:
  • GPU メモリはモデルのコピーを保存するのに十分ではなくなり、データの並列処理では需要を満たすことができなくなりました。
  • フレームワークは、モデルを保存およびトレーニングするためのマルチ GPU 機能を調整するための新しい並列戦略を提供する必要があります。
  • ユーザーがモデルの分散バージョンを簡単に実装できるように、シンプルで使いやすいインターフェースをユーザーに提供する方法。
  • 超大規模モデルは、計算効率と通信効率に大きな課題をもたらします。コンピューティングと通信の効率を向上させるにはどうすればよいでしょうか?
  • 下流のタスクと接続し、バッチ予測とオンライン推論の要件をサポートする方法。
• 高コスト:
  • テラバイトモデルを例にとると、モデルパラメータのサイズは 4 TB、勾配も 4 TB で、さらにオプティマイザーの状態とアクティブテンソルが加わり、ビデオメモリの要件は膨大になります。
  • 業界で同規模のモデルをトレーニングするために必要なリソース：Nvidia 3072 A100、Google 2048 TPU v3、実装するにはコストが高すぎます。
  • コストを削減し、効率を高め、使用するリソースを減らし、トレーニングの収束を高速化する方法。

現在、Horovod、Tensorflow Estimator、PyTorch DDP など、データ並列処理をサポートする分散トレーニング フレームワークがいくつかあります。パイプラインの並列処理をサポートする Gpipe、PipeDream、PipeMare など。演算子分割をサポートする Mesh Tensorflow、FlexFlow、OneFlow、MindSpore など。ただし、これらのフレームワークにはまだいくつかの欠点があります。

• シングル モード: 多くのフレームワークは一部の並列戦略のみをサポートしており、さまざまなハイブリッド並列戦略を完全にサポートすることはできません。
• 高い参入障壁: ユーザーがモデルの分散バージョンを実装するのは困難でコストがかかり、効率的な分散並列戦略を実装するにはドメイン専門家の経験が必要です。
• 移行コストが高い: 異なる分散フレームワークには、個別の並列実装があります。フレームワークごとに独自の DSL があります。ユーザーが並列戦略を切り替える場合、さまざまなインターフェースを学習し、モデルを書き直す必要があります。
• 不十分なパフォーマンス: 一部のフレームワーク実装では、クラスターの物理環境が考慮されていません。

分散トレーニングの現在の課題に対処するために、私たちは以下の主な目標を掲げて分散トレーニング フレームワーク Whale を開発しました。

• 複数の並列戦略を統合: 1 つのフレームワークでさまざまな並列戦略とこれらの戦略のさまざまな組み合わせをサポートします。
• シンプルで使いやすいインターフェース: 並列戦略の構成を完了するには、ユーザーは数行の注釈を追加するだけでよく、モデル コードを変更する必要はありません。
• 効率的なトレーニング フレームワーク: ハードウェア リソース、ネットワーク トポロジ、モデルを組み合わせて共同最適化を行い、効率的な分散トレーニング フレームワークを作成します。

2. PAIが独自に開発したWhaleフレームワーク

1. クジラの建築

私たちは、複数の並列戦略を統合し、次のような観点から分散トレーニングの課題に対処する高性能な分散トレーニング フレームワークである Whale をリリースしました。

• さまざまな並列化戦略を統合および抽象化し、分散トレーニング フレームワークで複数の並列化戦略をサポートします。
• 分散並列インターフェースのセットは Tensorflow に基づいて設計されており、Tensorflow と完全に互換性があります。ユーザーは、数行の注釈を追加するだけで、さまざまな分散並列戦略を実装できます。
• モデル構造とネットワーク トポロジを組み合わせてスケジューリングと通信を最適化し、効率的な分散トレーニング機能を提供します。

Whale フレームワークは下図のとおりで、主に 4 つのモジュールに分かれています。

• API: シンプルで使いやすいインターフェースを提供し、ユーザーがさまざまなハイブリッド並列戦略を組み合わせて使用​​できるようにします。
• Whale IR: 並列戦略を内部表現に変換し、TaskGraph、Multi-Dimension、VirtualDevices の抽象化を通じてさまざまな並列戦略を表現します。
• Whale Engine: WhaleIR に基づいて、グラフ編集ツールを使用して分散実行グラフを構築します。
• ランタイム: 分散実行グラフを TFGraph に変換し、TF ランタイムを呼び出して実行します。

2. クジラの紹介 使いやすいインターフェース

Whale は、さまざまな並列戦略を記述するための簡潔で使いやすいインターフェースを提供します。主なプリミティブは次のとおりです。

• クラスター: 仮想デバイスの分割を構成する方法
• レプリカ: データの並列処理
• ステージ: タスクグラフを分割する
• パイプライン: パイプラインの並列処理
• 分割: 演算子分割

これらのインターフェースを使用すると、次のようなさまざまな並列戦略を組み合わせることができます。

• データ並列処理:

• パイプラインの並列処理:

パイプライン並列処理 + データ並列処理:

その他の並列戦略の例:

3. クジラの訓練プロセス

Whale を使用した分散トレーニング プロセス:

• 並列戦略構成:
  • Whale API を使用して、モデルの並列戦略を構成します。モデル コードを変更せずに、数行の注釈を追加するだけです。この方法はセクション2.2に示されています。
  • モデルは複数の TaskGraph に分割できます。 TaskGraph は複数の並列戦略の構成をサポートしており、各 TaskGraph は異なる並列戦略で構成できます。
• 仮想リソース部門:
  • 並列戦略に従って仮想デバイスを分割し、各 TaskGraph は仮想デバイスに対応します。
  • GPU リソースとネットワーク トポロジに応じて、仮想デバイスの物理デバイスを選択します。
• 分散実行グラフ:
  • 並列戦略とリソース割り当て情報に基づいて、グラフ編集ツールを使用して実行グラフを編集（グラフのコピー、分割、通信ノードの挿入など）し、最終的な分散実行グラフを生成します。
  • 分散グラフを実行するには TF ランタイムを呼び出します。

3. 兆M6モデルの事前トレーニング

1兆ドルモデルには膨大な計算能力が必要です。計算能力の要件を削減するために、Whale は MoE (Mixture-of-Experts) 構造を実装します。 MoE の主な特徴はスパース活性化です。ゲーティング (ルーター) は、入力計算の上位 k 人のエキスパート (k は通常 1 または 2) を選択するために使用され、計算能力の要件が大幅に削減されます。

Whale は MoE (Mixture-of-Experts) レイヤーを実装し、エキスパートの並列処理をサポートし、エキスパートを複数のデバイスに分割して、単一デバイスのメモリと計算能力の要件を削減します。同時に、データの並列処理はトレーニングの同時実行性の向上に役立ちます。そのため、M6 モデルのトレーニングには、データ並列処理とエキスパート並列処理を組み合わせたハイブリッド並列戦略が採用されています。MoElayer はエキスパート並列処理を採用し、他のレイヤーはデータ並列処理を採用しています。

Whale は、モデルのハイブリッド並列トレーニングのためのシンプルで使いやすいインターフェースを提供します。並列戦略を構成するには、数行の注釈を追加するだけでよく、モデル自体を変更する必要はありません。 M6モデルは、データ並列処理+エキスパート並列処理の戦略を採用しています。次の注釈を追加するだけです。

同時に、トレーニング リソースを節約し、トレーニング効率を向上させるために、Whale はさまざまな最適化テクノロジを提供します。

ビデオメモリの最適化:

• 自動勾配チェックポイントは、アクティベーション メモリを節約するために最適なチェックポイント ノードを自動的に選択します。
• グループ単位で適用し、オプティマイザー適用ステージでビデオ メモリを最適化します。
• CPU オフロード テクノロジーは、ビデオ メモリのオプティマイザー ステータスと重みを最適化します。
• 通信プーリングは、データ ブロック サイズと通信の同時実行性を制御し、通信用のビデオ メモリを節約します。

コンピューティングと通信の高速化:

• DP + EP ハイブリッド並列戦略を採用して計算能力要件を削減します。
• パケット融合通信、半精度通信、トポロジーを考慮したAll2All通信オペレータなどの技術を採用し、通信効率を向上します。
• 混合精度、コンパイル最適化などのテクノロジーを組み合わせてトレーニング効率を向上させます。
• Whale フレームワークの助けを借りて、兆 M6 モデルの事前トレーニングが 480 V100 上で初めて 3 日間で完了しました。 Nvidia が 3072 個の A100 GPU を使用して 1 兆個のパラメータを実現し、Google が 2048 個の TPU を使用して 1.6 兆個のパラメータの大規模モデルを実現したのに比べ、DAMO Academy は今回、480 個の V100 32G GPU のみを使用して 1 兆個のモデル M6 を実現し、コンピューティング パワー リソースを 80% 以上節約し、トレーニング効率を 11 倍近く向上させました。

IV.結論

モデルパラメータの規模はますます大きくなり、大規模モデルが開発トレンドとなってきました。超大規模モデルのトレーニングの課題に対処するために、私たちは独自の Whale フレームワークを開発しました。このフレームワークは、さまざまな並列化戦略を統一された方法で抽象化およびカプセル化し、分散トレーニング フレームワークで複数の並列化戦略をサポートします。 Whale はシンプルで使いやすいインターフェースを提供します。ユーザーは、モデル自体を変更することなく、数行の注釈を追加するだけでさまざまな並列戦略を実装できます。同時に、ハードウェア リソース、ネットワーク トポロジ、モデルを組み合わせてハードウェアとソフトウェアを最適化し、効率的な分散トレーニング フレームワークを提供します。

Whaleフレームワークを通じて、480枚のV100 GPUカードを使用して1兆規模のモデルをトレーニングし、3日以内にモデルトレーニングの収束を完了し、超大規模モデルトレーニングの実装を可能にしました。今後はWhaleフレームワークをさらに改良し、大規模化、高速化、コスト効率の向上という3つの側面から機能を拡張していきます。同時に、より多くのビジネス シナリオで Whale 機能の実装を促進し、技術機能から製品機能への変革を可能にします。