Transformers、MoEエキスパートの勾配がゼロになる問題を報告　特定の分散構成でBF16を使うと更新が停止する可能性

TransformersのMixture of Experts（MoE）モデルの学習で、損失曲線からは見えにくい問題が報告されました。10月17日、プロジェクトのメンバーは、token dispatch経路でBF16を使用し、Expert Parallelismの規模が特定の範囲にある場合、すべてのエキスパートの勾配がゼロになる可能性があると報告しました。再現環境はTransformers 5.20.1、PyTorch 2.13.0、NVIDIA B300 4基で、報告には小規模なQwen3 MoEの再現コードも添付されています。問題報告

公式ドキュメントによると、token dispatchは各プロセスのトークンを対応するエキスパートに送り、処理結果を交換します。また、エキスパートの重みはFSDP2でさらにシャーディングできます。報告によると、tp_size < ep_size < fsdp_size * tp_sizeの場合、エキスパート勾配の除数と、そのFSDPグループのサイズが異なるため、NCCLのPREMUL_SUMリダクション経路が使われます。報告者の分析では、関連するBF16スケーリング係数の型処理に誤りがあり、特定の係数がゼロとして読み取られます。並列化の仕組みに関するドキュメント、原因と再現手順

この問題の技術的なリスクは、ほかのパラメーターは学習できるため、全体のlossが下がり続ける可能性があり、学習曲線だけではエキスパートが更新されているか判断できない点です。報告者はQwen3-30B-A3Bのフルファインチューニングで、1回の逆伝播後に96個すべてのエキスパート勾配テンソルがゼロだったと観察しました。これは特定環境での観察であり、すべてのMoE学習に影響すると一般化することはできません。報告では、FP32、エキスパートを凍結した場合、および一部の並列化構成の境界条件は例外として挙げられています。テスト結果と影響範囲

修正案では、エキスパート勾配を先に合算し、その後ローカルで除算することで、問題の経路を回避します。提案者によると、再現ケースではBF16の勾配ノルムがFP32に近い値まで回復しました。10月18日に追加されたコミットでは、lossと各勾配を比較するテストがさらに加わりました。ただし、確認時点でPRはまだオープンです。元のissueは完了としてマークされていますが、正式リリースで修正済みとは判断できません。エンジニアは今後、マージとリリースの状況を確認し、同じ構成でエキスパートの勾配と重みの更新を抜き取り確認して、lossの低下だけで学習の正常性を判断しないようにする必要があります。修正提案