AWS、EKS・EFA・DeepEPでMoE強化学習のスループット40%向上

AWS は、Mixture-of-Experts（MoE）型のモデルを RLHF や GRPO といった強化学習で追加学習する際の構成を、公式ブログで紹介した。Kubernetes のマネージドサービスである Amazon EKS、ノード間の高速通信を担う Elastic Fabric Adapter（EFA）、エキスパート並列の通信を最適化する DeepEP を組み合わせ、スループットを 40% 向上させたとしている。

MoE は、数千億から兆単位のパラメーターを持つモデルを、推論時には一部のエキスパートだけを動かすことで効率よく扱える構造として標準的になった。しかし学習の負担は軽くならない。AWS によれば、大規模な強化学習には三つの課題が同時に生じる。回答を大量に生成する「ロールアウト」と方策の学習という異なる計算の調整、数百台規模のアクセラレーター間での高速な通信の維持、そして全体の処理の釣り合いをとり続ける制御だ。

ロールアウトは分散推論で、最初の応答の速さよりも全体の処理量を重視する。一方の学習は、各ワーカーが足並みをそろえて進む必要があり、1 台の遅れが全体を止めたり NCCL のタイムアウトを招いたりする。学習が遅ければ生成側が待たされ、生成が遅ければ学習側の計算資源が遊ぶ。加えて MoE では、トークンを各エキスパートへ振り分けるエキスパート並列が、全ノード間で動的な all-to-all 通信を生む。モデルが疎になるほど、学習は計算より通信に縛られるようになる。ジョブが 1 台に収まらなくなると、通信はノード内の高速な NVLink から、帯域の狭いノード間接続へ移る。DeepEP はこの通信を EFA 上で効率化する役割を担う。

入手した材料には、40% の向上を測った具体的なモデルや台数、比較の基準までは含まれていない。それでも、MoE の強化学習では GPU の性能よりネットワークと全体制御が成否を分けるという整理は、自社でモデルを追加学習する企業にとって基盤選びの判断軸になる。