Huawei GTS LocMoE+: 高いスケーラビリティとアフィニティ MoE アーキテクチャ、アクティブ ルーティングを実現する低オーバーヘッド

AIxivコラムは、当サイトが学術的・技術的な内容を掲載するコラムです。過去数年間で、このサイトの AIxiv コラムには 2,000 件を超えるレポートが寄せられ、世界中の主要な大学や企業のトップ研究室がカバーされ、学術交流と普及を効果的に促進しています。共有したい優れた作品がある場合は、お気軽に寄稿するか、報告のために当社までご連絡ください。提出電子メール: liyazhou@jiqizhixin.com; zhaoyunfeng@jiqizhixin.com

この論文の共著者は、GTS AI Computing Lab の Li Jing 博士、Sun Zhijie 博士、Lin Dachao 博士です。主な研究および実装分野には、LLM トレーニングと促進、AI トレーニング保証、グラフ コンピューティングが含まれます。

MoE は、トレーニングと昇進プロセスにおける低コストと高効率により、過去 2 年間、大規模言語モデルの分野で頭角を現してきました。 MoE の魂として、専門家が学習の可能性を最大限に高める方法について、関連する研究と議論が無限に行われています。以前、Huawei GTS AI Computing Labの研究チームは、新しいルーティングネットワーク構造、通信オーバーヘッドの削減を支援するローカルロスなどを含むLocMoEを提案し、広く注目を集めました。

LocMoE の上記の設計は、次のようなトレーニングにおける一部の MoE の古典的な構造のボトルネックを効果的に軽減します。専門家のルーティング アルゴリズムではトークンを効果的に区別できない可能性があり、通信同期効率は内部の送信帯域幅の違いによって制限されます。ノード間など。さらに、LocMoE は、識別トークンを正常に処理できる専門家の能力の下限を証明し、解決します。この下限は、トークンが専門家に受動的に配布されるシナリオにおいて、トークン バッチ内に存在する識別トークンの確率分布に基づいて導出されます。そして、専門家にも最適なトークンを選択する能力があれば、識別トークンが処理される確率は大幅に高まり、専門家の能力の下限はさらに圧縮されることになる。

上記のアイデアに基づいて、チームはさらに、低オーバーヘッドのアクティブ ルーティングに基づく MoE アーキテクチャを提案し、LocMoE+ と名付けました。 LocMoE+ は、LocMoE の高識別エキスパートとローカル通信の利点を継承し、ルーティング戦略をさらに変革し、トークンとエキスパート間の親和性インデックスを定義し、このインデックスから開始してトークン配布をより効率的に完了することで、トレーニング効率を向上させます。

• 論文リンク: https://arxiv.org/pdf/2406.00023

論文の紹介

この論文の核となるアイデアは組み合わせる専門家による従来のパッシブルーティング アクティブルーティングは、特定の容量の下で識別トークンを処理する確率を向上させるため、サンプルノイズが減少し、トレーニング効率が向上します。この論文は、トークンとそれに割り当てられたエキスパートとの関係から始まり、計算オーバーヘッドの低いスキームでエキスパートとトークンの間の親和性を定量化して定義します。したがって、この論文では、グローバル適応ルーティング戦略を実装し、親和性スコアに基づいてエキスパート次元でトークンを再配置および選択します。同時に、トークン特徴量の分布が安定するにつれてエキスパート能力の下限が徐々に減少することが証明されており、トレーニングのオーバーヘッドを削減できます。

この論文は、学習ルーティング戦略においてトークンがより小さな角度のエキスパートにルーティングされる傾向があるという発見に基づいて、2 つのルーティング メカニズムを組み合わせた最初の論文であり、影響を与える既存のアクティブ ルーティング ソリューションの過剰なオーバーヘッドの障害を打ち破ります。トレーニングの効率も向上し、パッシブ ルーティングの性質との一貫性が保たれます。

この一連の作業の高い効率性を証明するために、著者が LocMoE とはまったく異なるハードウェア環境 (サーバー モデル、NPU カード モデル、クラスター ネットワーキング スキーム)、トレーニング フレームワーク、およびバックボーン モデルを選択したことは言及する価値があります。そして持ち運びのしやすさ。

適応型双方向ルートディスパッチメカニズム

背景の紹介

従来の MoE には 2 つのルートディスパッチメカニズムがあります:

(1)ルーター、トークン全体を直接入れます特徴が割り当てられます。

(2) ソフトルーターは、トークン特徴の重み付けされた組み合わせを割り当てます。

この記事では、計算コストが低い (1) について引き続き検討します。ハードルーターのシナリオでは、1) 各トークンが上位 K のエキスパートを選択できるトークン選択ルーター (TCR)、2) 各エキスパートが適切な上位 -C を選択できるエキスパート選択ルーター (ECR) に分けることができます。トークン。容量の制限により、各エキスパートが受け取るトークンの数には上限 C があるため、シナリオ 1) では、各エキスパートが受け取るトークンは切り捨てられます:

以前の研究では、MoE トレーニングは 2 つの段階に分かれていると指摘しました。フェーズ 1。ルーティング トレーニングでは、ルーティングがトークンを合理的に割り当てることができるようにします。つまり、異なる分野のトークンや大きな違いがあるトークンを区別して、異なる専門家に割り当てることができます。フェーズ 2。トークン ルーティングにより、各エキスパートの役割は、同じ分野または同様の特性を持つトークンを受け取ることです。各エキスパートは、特定のトレーニングを受けた後、関連する分野および特性の知識を取得できます。要約すると、MoE トレーニングの各ステップの「成功」の鍵は、トークン配布の正確さと合理性にあります。

この記事の貢献

(1) ソフトマックス活性化関数演繹を通じて、エキスパートとトークン間のコサイン類似度はより正確にアフィニティを測定できます:

(2) の観点から理論モデリングでは、TCR と ECR の単一トレーニングの成功率が 2 つの一般的なシナリオで分析されます。

理論に基づいて、著者は次のことを指摘しました。ルーティング トークンの機能が不十分です。TCR はトレーニングされるたびに、ECR よりもトレーニングが成功する確率が高く、適切なトークンが確実に選択されるようにするために、より多くの専門家の能力が必要です。

モデルトレーニングの後半段階では、ルーターがトークンを正しく割り当てる一定の能力を備えている場合、ECR がトレーニングされるたびに、TCR よりもトレーニングが成功する確率が高くなります。この時点では、より小さい容量しかありません。適切なトークンを選択するために必要です。

• この理論も非常に直感的です。ルーターにディスパッチ能力がない場合、つまり、エキスパートがディスパッチ能力を持っている場合には、トークンにランダムにエキスパートを選択させる方がよいでしょう。適切なトークンの場合は、 ECR を使用する方が適切です。したがって、著者は TCR から ECR​​ への移行を推奨し、同時に、エキスパート キャパシティの需要推定に基づいて、トレーニングの後期段階ではより小さなエキスパート キャパシティを使用するグローバル レベルの適応型ルーティング スイッチング戦略を提案します。

実験結果

この論文の実験は、Huawei 独自の Cache Coherence Protocol High-Performance Computing System (HCCS) のおかげで、Ascend 910B3 NPU の自己構築クラスター上で実施されました。デバイス間のデータ通信はマルチカードシナリオで実現され、Ascend プロセッサー専用に設計された Huawei Collective Communication Library (HCCL) により、HCCS などの高速リンクでの高性能分散トレーニングが可能になります。実験では、Ascend NPU と互換性のある PyTorch for Ascend フレームワークと、加速ライブラリ AscendSpeed および Ascend デバイス用に特別にカスタマイズされたトレーニング フレームワーク ModelLink を使用し、LLM 並列戦略と通信マスキングの最適化に焦点を当てています。

トレーニング効率

実験結果は、モデルトレーニングの収束や有効性に影響を与えることなく、各エキスパートが処理する必要があるトークンの数をベースラインと比較して 60% 以上削減できることを示しています。通信の最適化と組み合わせることで、クラスタ サイズが 32 カード、64 カード、256 カードの場合、トレーニング効率が平均 5.4% ～ 46.6% 向上します。

ビデオ メモリの使用量

LocMoE+ では、特にクラスター サイズが小さく、コンピューティングが集中しているシナリオでは、ビデオ メモリの使用量がある程度増加します。 Ascend Insight ツールを使用してメモリ監視サンプルを分析すると、LocMoE+ メモリ使用量はベースラインと比較して 4.57% から 16.27% 低下し、LocMoE と比較して 2.86% から 10.5% 低下したことがわかります。クラスターのサイズが大きくなるにつれて、メモリ使用量の差は縮小します。

有効性評価

オープンソースの評価セット C-Eval と TeleQnA、および独自に構築された ICT ドメイン評価セット GDAD を使用して、一般知識とドメイン知識における LocMoE+ の機能を評価しました。このうち、GDAD は、ドメイン タスク、ドメイン コンピテンシー認定試験、一般能力の 3 つの主要な評価システムにおけるモデルのパフォーマンスを検査するために、18,060 サンプルを含む合計 47 の下位項目をカバーしています。

十分な SFT の後、LocMoE+ は、ドメイン タスク機能の 16 のサブ能力において、ベースラインと比較して平均約 20.1% 向上し、LocMoE と比較して約 3.5% 向上しました。ドメインコンピテンシー認定試験はそれぞれ 16% と 4.8% 増加しました。一般機能の 18 のサブ機能のうち、LocMoE+ はそれぞれ約 13.9% と 4.8% 向上しました。全体として、LocMoE+ は、GDAD、C-Eval、TeleQnA でそれぞれ 9.7% ～ 14.1% のパフォーマンス向上を示しています。

以上がHuawei GTS LocMoE+: 高いスケーラビリティとアフィニティ MoE アーキテクチャ、アクティブ ルーティングを実現する低オーバーヘッドの詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。