MoEトレーニングの効率とパフォーマンスのボトルネックを打破し、Huawei Panguの新しいスパース大規模モデルアーキテクチャLocMoEがリリースされました

2023年12月、初のオープンソースMoE大規模モデルMixtral 8×7Bがリリースされました。複数のベンチマークテストでは、そのパフォーマンスはGPT-3.5とLLaMA 2 70Bをほぼ上回りましたが、推論オーバーヘッドは約12Bの密なモデルと同等でした。モデルのパフォーマンスをさらに向上させるために、密な LLM では、パラメータ スケールの急速な拡大により、厳しいトレーニング コストに直面することがよくあります。 MoE を広く適用することで、計算コストを比較的変えずにモデル容量を大幅に拡張することができます。この特徴により、MoE は間違いなく LLM の発展を促進するための重要なテクノロジーとなります。

MoE 設計の本来の意図は、モデルの学習をより「専門的」にすることであり、その有効性は業界で認められています。しかし、既存の MoE アーキテクチャ トレーニングの欠点が徐々に明らかになりつつあり、主な欠点としては、専門家の負荷の不均衡、専門家内のサンプルの混合と専門家間の深刻な均質性、追加の通信オーバーヘッドなどがあります。

既存のMoEに共通するトレーニング効率とパフォーマンスのボトルネックを軽減するために、高性能コンピューティングとLLMトレーニング加速を専門とするHuawei GTS AIコンピューティングラボの研究チームは、LocMoEと呼ばれる新しいMoEアーキテクチャを提案しました。ルーティングメカニズムの観点から、スパースLLMトレーニングのコストを削減しながらパフォーマンスを向上させることを目指しています。

論文リンク: https://arxiv.org/abs/2401.13920

論文紹介

この研究で、著者らは、従来の MoE ルーティング メカニズムでは、同じエキスパート ネットワークにルーティングされるトークンに大きな違いが生じることが多く、エキスパート ネットワークの収束が妨げられることを発見しました。一方、異なるエキスパートにルーティングされるトークンは非常に類似しているため、エキスパート間で深刻な均質性が生じ、最終的にはモデルのセマンティクスの理解と生成の能力に影響を及ぼします。著者らは、エキスパートルーティングと入力データ特性の関係を理論的に説明し、エキスパートネットワークには下限容量があることを NLP 分野で初めて証明しました。この理論的根拠に基づいて、エキスパート ルーティングのゲーティング ウェイトは直交化され、エキスパート ネットワーク間の識別が大幅に強化され、元のスケールよりもはるかに小さいトークンを処理することでドメイン評価で同様の結果を達成できます。同時に、本研究では、MoE アーキテクチャに固有の All-To-All 通信のボトルネックをターゲットとし、負荷/通信の最適化を組み合わせて、効率的で高エネルギーの MoE アーキテクチャを提案します。

具体的には、著者らは、Pangu モデルのバックボーンに組み込まれて機能を強化する、LocMoE という新しい MoE アーキテクチャを提案しました。 LocMoE は、追加の通信と計算のオーバーヘッドを削減しながら、ルーティング メカニズムの説明可能性を高めることを目指しています。まず、著者らは、トークンは常に、トークンとの角度が最も小さいエキスパートにルーティングされる傾向があることを発見しました。エキスパート間のゲーティング重みベクトルが直交している場合、エキスパート ネットワークによって処理されるトークンは、可能な限り均一性を回避することができます。

そのため、本稿では、GAP レイヤーを使用してトークンの特徴を抽出し、それをルーティングの基礎として使用します。 GAP レイヤーは、直交ゲーティング重みの条件を満たすという特徴があり、Dense レイヤーに比べて計算量が大幅に削減されます。上記の構造に基づいて、著者はモデル損失に影響を与えずにエキスパートが処理するトークンスケールの下限を理論的に解決し、エキスパートネットワークの負荷を軽減します。さらに、著者らは補助損失と提案された局所性損失を組み合わせて、ルーティングに対するソフト制約を実行し、マシン間の All-To-All 通信のオーバーヘッドを削減しました。最後に、通信マスキングなどのエンジニアリング最適化を使用して、スパースな大規模モデルの全体的なトレーニング パフォーマンスをさらに向上させます。

著者らは、LocMoE アーキテクチャを Pangu-Σ 38B モデルに組み込み、意味的類似性の高い ICT ドメイン データをトレーニングに使用して、ドメイン知識の学習能力をテストしました。 10 個のダウンストリーム タスクでは、LocMoE の精度はネイティブ Pangu-Σ よりも一般的に高く、トレーニング パフォーマンスはステップごとに 10% ～ 20% 向上します。 MoE アーキテクチャは汎用性と移植性も高く、他のハードウェア仕様や他の MoE アーキテクチャの LLM バックボーンに低コストで組み込むことができます。

現在、LocMoEは、ICTサービス分野におけるファーウェイの専門知識の質疑応答アプリケーションであるAskO3に導入されています。AskO3はファーウェイのO3ナレッジコミュニティで開始され、数万人のエンジニアリングユーザーから高い評価を受けています。

イノベーション分析

ルーティングとデータ特性の関係

既存の MoE ルーティング メカニズムの一般的な説明可能性の欠如に対応して、著者らはトークン ルーティングの性質を分析し、トークンを効果的に区別できる構造を設計しました。特定のトークンの場合、学習したルーティング戦略では通常、割り当てるトークンに乗じたゲーティング ウェイトが大きいエキスパートが選択されます。

そして、エキスパートのゲーティング重みが直交性を満たす場合、エキスパートはより高い識別可能性を持つことができます。同時に、トークンは、より小さな角度を持つ専門家にルーティングされる傾向があると結論付けることができます。

著者は最終的にトークン特徴を抽出するための構造として GAP を選択し、そのゲーティング重みは直交条件を満たすことができます。

上記のルーティングメカニズムの本質は、図に示すように、ルーティング識別能力とエキスパートトークン間の最小角度の関係を説明しています。

図: LocMoEルーティングメカニズムの概略図

専門家の能力の下限

上記の理論に基づいて、著者らはエキスパートの能力には下限があることを発見しました。つまり、入力データが確実な場合、エキスパートは経験値のスケールよりもはるかに小さいトークンを処理することで同等のパフォーマンスを達成できます。この問題は、少なくとも 1 つのクラス識別トークンがエキスパートにルーティングされることを保証する最小トークン サイズを見つけることに変換できます。これは、すべてのエキスパートの容量が揃っている場合の下限として機能します。同時に、合理的なエキスパート容量は、トークンとゲーティングウェイト間の最小角度と負の相関関係にあり、角度が減少するにつれて指数関数的に増加すると結論付けることができます。実験により、エキスパート容量をこの下限に設定すると、モデルの収束と損失曲線に影響がないことが確認されました。

地域制約

MoE レイヤーでの LocMoE の損失は、補助損失と局所性損失の 2 つの部分で構成されます。補助損失は、スパースゲート MoE で最初に提案され、エキスパート負荷分散を改善するために SwitchTransformer にも適用されています。

ただし、マシン間の All-To-All によって発生する追加の通信オーバーヘッドは回避できません。そのため、著者は、エキスパートの負荷分散を前提として、トークンがローカルデバイス上のエキスパートに割り当てられる可能性が高くなり、最終的に制約バランスが達成されるように、局所性制約を追加しました。局所性損失は、現在のトークン分布と完全に局所化された分布の差、つまり KL ダイバージェンスによって定量化され、これにより一部のマシン間通信がマシン内通信に変換され、マシン内相互接続の高帯域幅が最大限に活用されます。

実験結果

著者らは、64、128、256 個の Ascend 910A NPU を含むクラスターで実験を行い、主に 2 つの従来の MoE 構造、Hash (Facebook 製) と Switch (Google 製) と比較しました。

トレーニング効率

著者らは、各実験グループのトレーニングプロセス中に計算、通信、マスキング、アイドルに費やされた時間を記録しました。その中でも、64Nと128Nの場合、LocMoEの計算オーバーヘッドと通信オーバーヘッドが最も低くなります。 LocMoE の計算オーバーヘッドは 256N と依然として最も低いですが、一部のデバイスにはエキスパートが含まれていないため、ローカル通信変換が無効になります。これは、LocMoE が計算と通信の両方で大幅な利益を得るための適用条件は、エキスパートの数が少なくともノードの数に等しいことであることを示しています。

図: 異なるクラスター構成におけるさまざまな MoE 構造のトレーニング効率

エキスパートロード

制約がエキスパートの負荷に与える影響を検証するために、著者らは各エキスパートにルーティングされるトークンの分布を分析しました。負荷分散を実現するために、RRE モジュールによって実装されたハッシュ ルーティングは、静的ルーティング テーブルのハード制約を使用して分散のバランスを確保し、LocMoE とスイッチはトークンの特定の特性を考慮してルーティングを実行します。学習ルーターとして、補助制約と局所制約の影響下で、LocMoE エキスパートのバランスはスイッチよりも大幅に優れており、安定した高いリソース使用率を示しています。

図: 文部科学省のさまざまな組織における専門家の負担

専門家に割り当てられたサンプルの類似性

著者らは、LocMoE を支持する主要な理論を検証するために実験を使用しました。左の図は、同じエキスパートにルーティングされたトークンの類似度が高く、1 に近いことを示しています。右の図は、トークンとそれがルーティングされるエキスパート間のゲーティング重み類似性分布が、他のエキスパートよりも右側にあることを示しています。これにより、トークンは角度が最も小さいエキスパートにルーティングされる傾向があるという理論的前提が確認され、エキスパート容量の下限を解決するための重要なパラメータ δ が示されます。

図: 同じ/異なるエキスパートにルーティングされたトークンの類似性 (左)、トークンとルーティング先エキスパートの類似性 (右)

NLP分野の下流タスク

Pangu-Σは、金融や健康など40以上の分野のコーパスを使用して事前トレーニングされており、マルチドメインのテキストデータから知識を学習する能力を実証しています。本研究では、無線ネットワーク、光、オペレータ IT などの製品ラインの技術レポートやツールマニュアルを含む ICT サービスのドメインデータを使用して、専門分野の知識における LocMoE の学習パフォーマンスを評価します。概念間の論理的複雑さに応じてL1からL3に分け、10種類のNLPタスクの評価データセットを整理します。ネイティブ Pangu-Σ と比較すると、LocMoE はモデルの意味理解と表現能力をある程度向上させます。

図: NLP 分野の下流タスクのパフォーマンスをネイティブ Pangu-Σ と比較