9月26日、AWSが「Scaling MoE reinforcement learning on Amazon EKS with EFA and DeepEP with 40% more throughput」と題した記事を公開した。MoEモデルに強化学習を適用しようとすると、スパースなトークンルーティングが生み出す不均一なall-to-all通信がボトルネックになり、スケールさせるほど通信コストが膨らむ——この"MoE強化学習の壁"を、Amazon EKS・EFA・DeepEPの組み合わせで突き破り、スループットを40%向上させたアーキテクチャが本記事では詳しく紹介されている。
なぜMoEの強化学習は難しいのか
Mixture-of-Experts（MoE）は、スパース性を活かして推論コストを抑えながら数千億〜数兆パラメータに拡張できるアーキテクチャとして、LLMの標準的な選択肢になりつつある。しかし「推論が軽い」という特性は、学習の手軽さを意味しない。
MoEモデルの後学習（post-training）において、RLHF（人間のフィードバックによる強化学習）やGRPO（Group Relative Policy Optimization）を適用する場合、次の3つの課題が同時に発生する。GRPOはDeepSeekが提案したRLHFの派生手法で、個別の報酬モデルを用意せず同一プロンプトから複数出力をサンプリングしてグループ相対スコアでポリシーを更新する点が特徴だ。
- ロールアウト生成とポリシー学習という異質なワークロードのバランス調整
- 数百台のアクセラレータにわたる高スループット通信の維持
- すべてのサブシステムを動的にオーケストレーションし、均衡を保つこと
なかでも通信のオーバーヘッドが核心的な問題だ。MoEが採用するExpert Parallelism（EP）は、トークンをデバイス間で動的にルーティングするall-to-all通信を生み出す。Tensor Parallelismなどの整然とした通信パターンと異なり、EPはスパースで細粒度・不均一なトラフィックを発生させる。EPの並列度が増すほど、この通信はノード間（inter-node）にまたがり、ボトルネックになる。
DeepEP over EFAで40%のスループット向上
この課題に対する最大の解決策が、**DeepEPをEFA（Elastic Fabric Adapter）上で動作させる構成**だ。DeepEPはDeepSeekが開発・公開したオープンソースライブラリであり、MoEのExpert Parallelism向けに特化した通信カーネルを提供する。DeepSeek自身もMoEアーキテクチャを採用しており、そのスパース通信の最適化知見がDeepEPに凝縮されている。タイトルで「DeepSeekも採用するスパースAIアーキテクチャ」と表現しているのはこの文脈に由来する。
標準的なNCCLのall-to-allコレクティブは、密で規則的な通信には効率的だが、MoEのスパースで不均一なトラフィックには過剰なオーバーヘッドを生む。DeepEPはこれを、2つの専用GPUカーネルで置き換える。
- dispatchカーネル: ローカルGPUからリモートのエキスパートへトークンをルーティング
- combineカーネル: 処理済みトークンを収集して戻す
ノード内の転送にはNVLink（NVSwitch経由）を使い、ノード間の転送にはlibfabric経由でEFAを使う。P5・P6などのインスタンスでは、NVIDIA GPUDirect RDMAによってGPUメモリバッファ間でCPUとOSをバイパスした直接転送が可能になる。
AWSはDeepEPのオープンソースへの貢献として、通信プリミティブをCUDA固有のRDMAバックエンドからlibfabricベースの実装に移植した。これによりDeepEP v2がネイティブのEFAサポートを獲得している。またNCCL 2.31にも最新のEFA最適化が組み込まれた。
48台のP5enインスタンス（16台がトレーニング、32台が推論に専用）でスーパースパースMoEモデルを動かした結果、DeepEP over EFAを有効化することでRLロールアウトスループットが40%増加した。
アーキテクチャ全体像：EKS中心の3層構成
このスケーリング構成の骨格は、Amazon EKS・EFA・Amazon S3の3層に分離されている。それぞれが独立してスケールできる設計だ。
EKSクラスタのノードグループ構成
元記事では、EKSクラスタを以下の3種のノードグループで構成する方法が整理されている。
| ノードグループ | 役割 |
|---|---|
| GPUインスタンス | ロールアウト生成、報酬モデル推論、ポリシー学習 |
| CPUインスタンス | 環境実行、前処理 |
| メモリ最適化インスタンス | 経験バッファ、チェックポイントキャッシュ |
ロールアウトワーカーはCPUベースの環境Podとやり取りしてサンプルを生成し、メモリ最適化バッファに書き込む。ポリシートレーニングワーカーはそのバッチを消費してモデルの重みを更新し、新しいチェックポイントを次のロールアウトラウンドに戻す。完成した学習成果物はAmazon S3に永続化される。
通信ドメインの分離
- ノード内: NVLink/NVSwitch → 高帯域幅
- ノード間: EFA + GPUDirect RDMA → DeepEPで最適化
Spotインスタンスによるコスト削減
もう一つの最適化として、ロールアウト生成にEC2 Spotインスタンスを活用できる点が紹介されている。
ロールアウト生成は独立したワーカーに分散できる推論タスクであり、Spot中断が発生しても未完了タスクをキューに戻して他のワーカーに再割り当てすればよい。ポリシートレーニングワーカーはSpot中断から切り離して安定した容量を確保しつつ、ロールアウト側はSpotで柔軟にスケールするという分離が可能だ。
必要な環境要件
記事で示されている動作環境は以下のとおりだ。
- Amazon EKS 1.31以降
- EFAインストーラ 1.49（AWS OFI NCCLプラグイン含む）
- DeepEP 2.0.0、NCCL 2.31.2、SGLang 0.5.17、PyTorch 2.12.1（CUDA 13.0）
- 対応GPUインスタンス:
p5.48xlarge、p5e.48xlarge、p6-b200.48xlarge
LLMのポストトレーニングコストが注目される中、MoEアーキテクチャの採用が広がるにつれてEP通信のボトルネックは普遍的な課題になりつつある。AWSがDeepEPのlibfabric移植にコントリビュートしてEFA上でのMoE学習を最適化した今回の取り組みは、この課題に対する実践的な回答だ。
詳細はScaling MoE reinforcement learning on Amazon EKS with EFA and DeepEP with 40% more throughputを参照していただきたい。
