著者：Huang Jun、Alibaba Cloud シニアアルゴリズムエキスパート

この記事では、Alibaba Cloud AI Platform PAIチームとNVIDIA Megatron-Coreチームによる、MoE（Mixture of Experts）大規模言語モデル（LLM）の実装と学習最適化に関する革新的な取り組みを紹介します。プレゼンテーションは以下の流れで進めていきます。
1. まず、MoE技術の発展の歴史を簡単に振り返り、実用化において早急に取り組む必要がある中核概念と主要な課題を抽出します。
2. 次に、両者が共同で開発したMoEフレームワークの独自の機能と優れたパフォーマンスについて詳しく説明し、モデルのトレーニング効率、リソース利用率、モデルパフォーマンスの向上における検証結果を示します。
3. 最後に、今回のコラボレーションに基づいて Alibaba Cloud が構築したプラットフォーム ツールと推奨プラクティスを簡単に紹介し、開発者が MoE テクノロジーを効率的に活用し、大規模モデル トレーニングの詳細な調査と広範な適用を促進できるようにします。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

MoEは、密なモデル構造を複数のサブ構造に分解し、各サブ構造をエキスパートとするモデルアーキテクチャです。学習および推論中に計算対象となるエキスパートセットを動的に選択することで、モデルパラメータのスパースな更新を実現します。
簡単に言えば、MoEモデルはモデル全体を複数の専門サブモジュール（エキスパート）に分解し、入力に関連する少数のエキスパートサブ構造のみを一度にアクティブ化および更新します。どのエキスパートが計算に参加するかは、ルーティングメカニズムによって決定されます。GPTのような超大規模言語モデルでは、MoE技術を採用することで、すべてのパラメータを毎回更新する必要がないため、学習および推論中の計算負荷を大幅に軽減でき、学習効率が大幅に向上します。さらに、推論フェーズではアクティブなネットワークパラメータのサブセットのみが使用されるため、計算リソース要件が大幅に削減されます。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

2023年後半、Mistral AIはオープンソースのMoEアーキテクチャに基づく大規模言語モデル「Mixtral 8x7B」をリリースしました。わずか467億個のパラメータを持つこのモデルは、複数の下流タスクのリーダーボードにおいて、当時最高のオープンソース高密度モデルであったLlama-2 70Bを凌駕しました。右のグラフは、推論リソースを同一条件で使用した場合、MoEモデルが高密度モデルを大幅に上回るパフォーマンスを示していることを示しています。この成果はMoEモデルへの業界の強い関心を呼び起こし、Alibaba CloudとNVIDIA Megatron-Coreチームは、大規模モデル分野におけるMoE技術の応用についてさらに協力することになりました。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

MoEアーキテクチャを用いて産業グレードのアプリケーションを実装する際に解決すべき問題を説明するため、本稿ではまずMoEの開発の歴史を振り返る。1990年代初頭に提案されたマルチエキスパートモデルの概念は、複数のエキスパートモデルを組み合わせて協調的にタスクを解決し、ゲートルーティングを用いてエキスパートを選択するというMoEの理論的基礎を築いた。
2017年のディープラーニングの急速な進歩により、モデル性能におけるパラメータ数の重要性がますます明らかになり、従来の高密度ディープラーニングフレームワークは容量不足に直面しました。GoogleはMoEとディープラーニングの統合を先駆的に進め、RNNにおけるその実現可能性を検証しました。2020年までに、Transformerアーキテクチャは言語モデリングにおいて卓越した性能と優れたスケーラビリティを実証しました。その後、GoogleのGshardプロジェクトはMoEをTransformerに組み込み、綿密な設計と実験を通じて性能を大幅に向上させ、Switch Transformerのパラメータ数を兆レベルにまで押し上げ、大規模モデリングにおけるMoEアーキテクチャの基盤を築きました。
ST-MoE、Tutel MoE、FasterMOE、MegaBlocksなどのその後の研究では、速度と安定性が深く探求されてきました。2023年には、GPT-4の優れたパフォーマンスが、MoE構造の採用をめぐって業界内で激しい議論を巻き起こしました。同年、Mixtralの実装はこの仮説を検証したように思われ、大規模モデルアーキテクチャにおけるMoEの適用が徐々に成熟し、広く受け入れられつつあることを示しています。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

次に、MoEと深層学習を統合したSparsely-Gated MoEをはじめとするいくつかの画期的な研究を簡単に分析します。この研究では、RNNネットワークを複数のエキスパートサブネットワークに分割し、トークンルーティングを制御するためにSoftmax Gatingメカニズムを用いています。しかし、従来のSoftmax Gatingロジックは負荷の不均衡を引き起こす可能性があります。学習が進むにつれて、一部のエキスパートネットワークが過剰に活性化する一方で、他のエキスパートネットワークのパラメータが十分に学習されない可能性があります。
この問題に対処するため、研究者らは、ソフトマックス関数を計算する前に、ゲーティング行列パラメータと入力特徴量にノイズを注入することで、よりバランスの取れたエキスパート選択を実現しました。さらに、ゲーティングネットワークの損失関数に補助的な損失関数を追加することで、モデルがより良い負荷分散を行えるようにしました。同時に、各ルーティング戦略において上位2～4人のエキスパートを選択する実験も行いました。これらの最適化により、最終的にモデルのパラメータ容量は1000億という閾値を超えることができました。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

Google の Gshard プロジェクトは、MoE を Transformer アーキテクチャにさらに適用しており、その主な革新は次のとおりです。
1. MoE Transformerの基本構造と並列モードを設計しました。各交互フォワード層はMoE層に置き換えられました。エキスパートは異なるデバイスに分散され、独立して計算され、他の層のパラメータは共有されました。
2. エキスパートの負荷分散問題に対処するために、「エキスパート容量」という概念が提案されています。これは、各エキスパートによって処理されるトークンの数を制限し、残余接続戦略を使用して、すでに最大容量になっているエキスパートをバイパスします。
3. 補助損失とランダムトップkルーティング戦略を導入して、エキスパート選択プロセスを最適化します。
4. 通信効率に関しては、ゲートネットワークを通じてデータを事前に選別し、エキスパートIDごとにデータをグループで送信する革新的な「ローカルグループディスパッチ」方式を提案し、通信効率を効果的に向上させます。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

実験結果によると、Gshard は Transformer パラメータの数を 150B から 600B に増加させましたが、計算コストは​​ 50% しか増加しませんでした。これは、モデルのスケーラビリティにおける MoE の大きな利点を実証しています。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

Switch Transformerのその後の研究では、MoE Transformerの学習プロセスを徹底的に最適化し、モデルパラメータサイズの最大スケーリングを実現しました。まず、より積極的なトップ1ルーティング戦略を採用しました。この戦略では、パラメータ更新のために一度に1つのエキスパートのみが選択されます。これにより、ルーティングの計算負荷が大幅に軽減され、トップ2戦略と比較して計算コストがほぼ半減するだけでなく、各エキスパートが処理するバッチサイズも削減され、通信コストがさらに軽減されます。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM の実装とトレーニングの最適化」。

第二に、既存のエキスパート負荷分散問題に対処するため、本研究では補助的な損失関数を導入する。この関数において、Nはエキスパートの総数、fはi番目のエキスパートに割り当てられたトークンの割合、Pはルータがi番目のエキスパートにトークンを割り当てる確率を表す。最適化の目的は、トークンが各エキスパートに均等に分配されるようにすることである。損失関数を最小化することで、fとPは1/Nに近づく。この時点で、損失関数の調整パラメータαは通常0.01程度に設定され、比較的良好な負荷分散効果が得られる。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

分散トレーニングに関しては、前述のデータ並列処理とエキスパート並列処理に加えて、この研究では非エキスパート アーキテクチャ用のモデル並列処理戦略を実装し、シングル カードのメモリ要件をさらに削減して、モデル パラメータの数を大幅に増やしています。
ルーティング負荷分散によって引き起こされるモデル学習におけるランダム性と不安定性に対処するため、本研究では、安定した学習のための一連の戦略を提案する。これらの戦略には、ローカル計算時には高いFP32精度を使用し、通信フェーズでは低いBF16精度を使用する確率的精度戦略、正規分布と小さなスケーリング係数で重みを初期化する戦略、そして過剰適合を低減するためのエキスパートドロップアウト（エキスパート内部ドロップアウト）の導入が含まれる。
これらの最適化により、MoEトランスフォーマーの学習収束速度は、反復回数と学習時間の両方において、高密度モデルよりも数倍高速になりました。最大のモデルであるSwitch-Cは、驚異的な1.5兆個のパラメータを誇ります。その後のST-MoEに関する研究では、MoEモデルの学習安定性と性能を向上させる方法がさらに検討されました。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

事前学習の安定性に関しては、エキスパートルーティングのランダム性と数値精度によって引き起こされる丸め誤差（つまり、ルーティングネットワークによって計算されるロジット値が過度に高いことによる学習の非収束問題）に対処するため、本研究ではz-loss関数を提案する。この関数において、Bはバッチ内のトークン数、Nはエキスパート数、Xはルーティングネットワークに入力されるロジットの次元に対応する。z-lossを導入することで、精度などの要因によって引き起こされるロジット値が過度に高い現象を効果的に抑制し、モデル学習の安定性を向上させる。z-lossは、高密度LLM学習の安定性向上にも適用できる。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

さらに、本研究では、MoE トレーニングの安定性とパフォーマンスのバランスに影響を与えるさまざまな要因を詳しく調べ、詳細な実験を通じて一連のベスト プラクティスをまとめています。
例えば、ドロップアウトなどの適切な量のノイズをルーティングネットワークに導入すると、トレーニングの安定性は向上しますが、モデルのパフォーマンスに悪影響を与える可能性があります。ST-MoEの研究では、ファインチューニングフェーズにおいて、ドロップアウトパラメータの増加は、密なMoEモデルと非密なMoEモデルで大きく異なる影響を及ぼし、密なモデルではより大きな影響を与えることがさらに明らかになりました。さらに、密なモデルとは異なり、バッチサイズを小さくし、学習率を大きくすることが、MoEモデルのトレーニングパフォーマンスを向上させる上でより効果的です。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

MixtralモデルはMegablocksフレームワークから独自のアプローチを借用しており、その主な特徴はMoE層の計算をブロックスパース演算として表現することです。前述の動的ルーティング問題（異なるエキスパートが異なる数のトークンを処理する問題）に対処するため、ユーザーは計算中にトークンをドロップするかパディングするかを選択する必要があります。前者はモデルのパフォーマンスに影響を与え、後者は学習コストを増加させます。
Megablocks は、複数の行列乗算の演算を大きなブロックの疎行列として革新的に定義します。つまり、複数のエキスパートの計算は固定サイズの大きな行列とみなされ、計算タスクは複数の小さな行列ブロックの GEMM (General Matrix Multiply) 演算に細分化されます。
このようなブロックスパース行列の処理を最適化するために、Megablocksフレームワークはブロック圧縮スパース行（BCSR）データ構造を活用し、行と列へのアクセスおよび変換操作の速度を最適化します。このアプローチにより、ルーティング行列の効率的な処理が可能になります。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

前述のMoEの開発履歴には、開発中に解決する必要があった問題が詳述されています。以下では、これらの問題に対処するためにAlibaba CloudがNVIDIA Megatron-Coreチームと行った技術協力とその成果について簡単に紹介します。Megatron-Coreは、多様なモデル並列処理のサポート、演算子最適化、通信最適化、メモリ最適化、低精度トレーニング（例：FP8）などの高度な技術など、大規模LLMのトレーニングのためのコア技術を統合した、成熟した軽量の大規模LLMトレーニングフレームワークです。Megatron-CoreはMegatron-LMの優れた機能を継承し、コード品質、安定性、機能の完全性、テストカバレッジの点で包括的に改善されています。重要なのは、フレームワークの設計が分離とモジュール性を重視しており、二次開発を行う際や新しいモデルアーキテクチャを探索する際に開発者に高い柔軟性を提供していることです。そのため、私たちはMegatron-Coreチームとの協力を選択しました。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

以下のセクションでは、MoEアーキテクチャにおけるMegatron-Coreの主なサポート機能について説明します。並列化に関しては、Megatron-Core MoEはエキスパート並列処理だけでなく、データ並列処理、テンソル並列処理、パイプライン並列処理、シーケンス並列処理を含む3D並列処理もサポートしています。超大規模MoEモデルの場合、エキスパート並列処理と他の並列処理戦略を柔軟に組み合わせることができます。
トークン分配メカニズムに関して、Megatron-Core MoEはドロップレスMoE動作を採用しており、トークンが破棄されることはありません。ルーティングおよび負荷分散の最適化レベルでは、一般的なTop-kなどの複数のルーティング戦略と、Sinkhorn、Zロス、負荷分散ロスなどの複数の負荷分散アルゴリズムをサポートしています。
さらに、複数のエキスパートが可変長の入力を受け取る問題に対処するために、Megatron-Core MoE は GroupedGMM テクノロジを導入し、非効率的な操作を最適化された CUDA カーネルに置き換えることで最適化します。
一方、モデルの転送と適応の面では、Megatron-Core MoE は豊富なモデル チェックポイント変換機能を提供しており、ユーザーは HuggingFace モデルをインポートし、TP (テンソル並列処理)、PP (パイプライン並列処理)、EP (エキスパート並列処理) などの構造を自由に調整した後、Megatron-Core を使用してモデルのトレーニング タスクを効率的に開始できます。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

Alibaba Cloud AI Platform（PAI）チームは、Megatron-Coreをベースとした大規模モデルのトレーニングをシンプルかつ効率的に行うツールを開発しました。Alibaba CloudのLingjunコンピューティングサービスを活用し、チームはPAI-Megatron-Patchツールキットをリリースしました。このツールキットは、10種類を超える主流のオープンソース大規模モデル形式からMegatron-LMおよびMegatron-Coreへのシームレスな変換を可能にします。ユーザーは、PAI DSWやDLCなどの製品を使用することで、Megatron-LMおよびMegatron-Coreを簡単に起動し、大規模な高密度モデルやMoEモデルのトレーニングを行うことができます。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

多くの大規模オープンソースモデルがMegatron学習フレームワークとは異なるHuggingFace形式を使用しているという問題に対処するため、PAI-Megatron-patchの主な機能は、HuggingFaceモデルからMegatronフレームワークへの重み変換サービスを提供することです。この変換は基本的にモデルの名前空間をマッピングし、層ノルム層、アテンション層、MLP層などのコアコンポーネントの定義変換をカバーします。PAI-Megatron-patchには簡潔な重み変換スクリプトが含まれており、ユーザーはモデル形式の変換操作を簡単に実行でき、移行プロセスを大幅に簡素化します。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

MoEモデルの変換を行う際には、変換されたファイルが分散学習に直接使用できることを確認する必要があります。このため、広範な検証作業を実施しました。Mixtralモデルを例に挙げると、8人のエキスパートと16枚のカードを持つシナリオにおいて、TP=4、EP=4の分割戦略を採用しました。各TPパーティションには、フォルダ内に4つのファイルが含まれており、8人のエキスパート全員のFFN重みが格納されています。つまり、各ファイルには2人のエキスパートのFFN重み情報が含まれています。
変換後のモデルのゼロショット損失精度の性能は、以下の表に示されています。データから、変換前後の精度の差は非常に小さく、モデル変換の有効性と精度が確保されていることがわかります。

私たちの研究は、モデル変換プロセスに焦点を当てるだけでなく、学習プロセス全体の堅牢性を厳密に検証することにも取り組んでいます。Mixtral 8x7Bモデルを例に挙げ、ゼロからの事前学習、チェックポイントベースの再学習、そして微調整という3つの学習フェーズすべてにおいて詳細なテストを実施しました。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

de novo事前学習フェーズでは、以下のパラメータを設定しました：global_size=256、LR（学習率）=1e-4、seq_len=2048、TP=4。24時間、2.4Kステップの学習後、損失は約1.9に収束しました。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM の実装とトレーニングの最適化」。

継続的な事前トレーニング段階では、global_size は 256、LR は 5e-5、seq_len は 2048、TP は 4 に設定されました。18 時間と 2,000 回のトレーニング ステップの後、損失関数は通常の収束動作を示しました。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

finenetune フェーズでは、次の 3 セットのテストを行いました。
1. オレンジ色の曲線は、global_size=128、LR=1e-5、seq_len=2048、TP=4 で微調整されたベース モデルを表します。
2. 青い曲線は微調整指示モデルを表し、そのパラメータ構成はオレンジ色の曲線と同じです。
3. 赤い曲線は、global_size=32、LR は依然として 1e-5、seq_len は 128 に調整され、TP は依然として 4 のままで、回答損失のみを計算する、微調整された指示モデルを表しています。
すべてのモデルはトレーニング中に通常の収束を示し、微調整されたベース モデルでは損失が大幅に減少し、これは予想と一致しました。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

Finetuneの有効性をさらに検証するため、評価ベンチマークとしてコード生成タスクを選択しました。約8万件の公開コードサンプルを学習に使用し、前述のSFT設定に従ってHumanEvolプラットフォーム上でモデルを評価しました。2,500ステップの学習後、HumanEvolにおけるモデルのパフォーマンスは、学習開始時の45.73%から53.05%へと大幅に向上し、学習プロセスの合理性と有効性を強く示しました。さらに、速度面では、この手法は同様のリソース条件下で実行されたMegablocksを上回りました。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

Alibaba CloudのAIプラットフォームPAIを通じて、公式のベストプラクティスガイドを提供しています。このガイドは、データ処理、トレーニング、推論、評価、オンラインサービスの展開に至るまで、AI開発プロセス全体におけるテスト結果をまとめたものです。
このガイドでは、データ準備フェーズからプロセス全体を網羅し、OSS（オブジェクトストレージサービス）およびNAS（ネットワーク接続ストレージ）から生データを読み取り、PAI DSW環境内で効率的なデータ前処理を実行する方法を解説します。前処理済みのデータは、NASまたはOSSに簡単に保存し、後続のモデルトレーニングに備えることができます。
このモデル学習方法は、PAI DLC上での大規模分散学習をサポートすると同時に、DSWが提供する単一マシン学習環境との互換性も備えているため、ユーザーは実際のニーズに応じて柔軟に学習方法を選択できます。学習完了後、モデルのチェックポイントをOSSまたはNASストレージシステムに直接エクスポートできます。
トレーニング完了後、オフライン推論とモデル性能評価のためのチェックポイントを設定できます。推論と評価の正しさが検証されると、開発者はEASが提供するオンラインモデルサービスにワンクリックでモデルをデプロイできます。このサービスのインターフェースを使用することで、開発者は様々なアプリやビジネスシナリオを容易に構築できます。

画像ソース: GTC 2024 China AI Day オンライン セッションのプレゼンテーション「NVIDIA Megatron-Core に基づく MoE LLM 実装およびトレーニング最適化」。

このベスト プラクティス ガイドでは、各ステップの操作方法をユーザーが明確に理解できるように、詳細な手順を段階的に説明します。
今後、Alibaba Cloud AI Platform PAIチームは、NVIDIA Megatron-Coreチームとの連携をさらに深め、インテンシティモデルとスパースモデルの両方において、トレーニング性能と効率性において更なる飛躍を目指し、AGI（人工知能）技術の発展に貢献していきます。世界中の開発者の皆様がオープンソースコミュニティプロジェクトやAlibaba Cloudの構築にご参加いただき、共にインテリジェントな未来を創造していくことを心よりお待ちしております。

上記のコンテンツは、GTC 2024 China AI Dayのオンライン中国語プレゼンテーションセッションからの抜粋です。画像のQRコードをスキャンするか、カンファレンスのウェブサイトにアクセスして、プレゼンテーションビデオを視聴し、配布資料をダウンロードしてください。