Qwen グローバルバッチ ロードバランス MoE LLM トレーニング
グローバルバッチ ロードバランスが MoE の性能とエキスパートの専門化を向上
Qwen は、Mixture-of-Experts (MoE) 大規模言語モデル (LLM) のトレーニングにおいて、グローバルバッチ ロードバランシング損失メソッドを導入しました。マイクロバッチレベルからグローバルバッチレベルへロードバランシング損失の計算をシフトすることにより、モデルはさまざまなタスクでの性能を向上させ、エキスパート間で重要なドメイン専門化を促進します。これは、マイクロバッチ バランスがしばしばエキスパートの特定データドメインへの専門化を妨げる制限に対処します。
マイクロバッチ ロードバランシングの制限
ほとんどの既存の MoE トレーニングフレームワーク（例えば、Megatron-core）は、マイクロバッチレベルでロードバランシング損失 ($L_{ ext{balance}}$) を実装しています。これは、損失が各マイクロバッチ内で計算され、その後グローバルバッチ全体で平均化されることを意味します。
このアプローチは、マイクロバッチに多様なデータが欠けている場合に問題となります。たとえば、マイクロバッチがコードデータのみを含む場合、マイクロバッチレベルの損失はルーターにそれらのコードトークンをすべてのエキスパートに均一に分配させます。これにより、エキスパートがコードなどの特定ドメインに専門化することを妨げ、全体のモデル性能を低下させる可能性があります。LLM トレーニングでは、単一のマイクロバッチのデータはしばしば同じドメインから来るため、これが多くのオープンソース MoE モデルが顕著なエキスパート専門化を欠く理由を説明しています。
グローバルバッチ ロードバランスの実装
この問題を解決するために、Qwen はグローバルバッチレベルでロードバランシング損失 ($L_{ ext{global}}$) を計算することを提案します。これは以下の3つのステップで達成されます：
- Synchronizing 並列グループ全体でのエキスパート選択頻度 ($f_{i}$)。
- Calculating 各並列グループ（たとえば、単一のGPU上）でのロードバランシング損失。
- Aggregating すべてのマイクロバッチ全体での損失。
エキスパート選択頻度はエキスパートの数に等しい次元を持つ小さなベクトルであるため、マイクロバッチ間でこのデータを同期することは計算コストが低く、実装は「ほぼ無料」と言えます。
性能向上とエキスパート専門化
Qwen は、3つの MoE 構成（3.4B で 0.6B アクティベート、15B で 2.5B アクティベート、および 43B で 6.6B アクティベート）および2つのデータ構成（120B トークンおよび 400B トークン）でこの方法をテストしました。
モデル効果の向上
マイクロバッチレベルの損失と比較して、グローバルバッチ アプローチはすべてのテスト設定（すべてのモデルサイズ、データ量、タスク）でより良い性能を達成しました。
ドメイン専門化の向上
グローバルバッチ バランスはエキスパートの専門化を可能にします。マイクロバッチ バランスはドメインに関わらずエキスパートが均一にアクティベートされる結果となりますが、グローバルバッチ バランスは特定のドメインによって特定のエキスパートが頻繁にアクティベートされることを許し、明確な専門化を示します。
バランスバッチサイズの影響
3.4B モデル（0.6B アクティベート）での実験では、バランスバッチサイズが 2 から 128 に増加するに従い、ptr-training PPL (perplexity) が急速に減少し、128 で飽和することが示されました。これは、主流の MoE フレームワークが通常、より大きなモデルに対してバランスバッチサイズを 8 から 16 の間で使用していることを考えると、グローバルアプローチの重要性を強調しています。
効率と効果のバランス
グローバルバッチ バランスはモデル品質を向上させる可能性がありますが、マイクロバッチ バランスの劣化を引き起こす可能性があり、これが計算効率に悪影響を及ぼすことがあります。
これを緩和するために、Qwen はグローバルバッチ バランス損失の上にマイクロバッチ バランス損失を追加する実験を行いました（グローバルバッチ損失の定数重み 0.01 を使用）。このハイブリッドアプローチは、更新ステップの速度を改善しました（1.64秒から1.59秒/更新ステップ）一方で、モデルの効果はほとんど影響を受けませんでした。
結論
グローバルバッチ ロードバランスを実装することにより、Qwen は MoE トレーニングにおける重要な制限に対処し、より高性能で専門化されたエキスパートを可能にします。この方法は、特にさまざまなドメインにおいてより大規模で専門化された MoE モデルへのスケーリング能力において、MoE モデルの最適化に新たな視点を提供します。