TransformersコミュニティがMoEのコンパイル最適化を提案、3つのテストで起動後の学習時間を10〜24％短縮

Transformersコミュニティは、混合エキスパートモデル（MoE）の学習を最適化する提案をレビューしている。コントリビューターのTitus-von-Koeller氏は10月28日にPR #50102を提出し、29日にレビュー待ちとなった。作者によると、地域コンパイルを有効にした3つのテストでは、起動後の学習時間が約10〜24％短縮した。確認時点で、提案はまだマージされていない。提案とテストの概要

地域コンパイルでは、モデル内で繰り返し使われるブロックをコンパイルし、ほかの同一ブロックで計算グラフを再利用する。PyTorchの公式ドキュメントによると、これによりコンパイル範囲を小さくし、初回コンパイル時のコストを抑えられる。一方、一部の演算はコンパイル範囲外に残り、モデル全体をまとめてコンパイルする場合とは最適化の余地が異なる。PyTorch技術ドキュメント

作者の調査では、各デコーダーレイヤーがルーター出力を共有リストに追加していた。レイヤーが増えるにつれてリストが長くなり、コンパイラーが新しいグラフを繰り返し生成して、再コンパイルの上限に達するとeager実行に切り替わっていた。提案では、各レイヤーの出力収集が、それ以前に蓄積された出力数に左右されないようにし、コンパイル領域の外で結果をまとめる。作者によれば、この設計は出力順序と勾配を維持し、補助損失に必要な勾配も保つ。技術説明

比較対象となったのは、単一のB300上でのQwen3-30B-A3B、単一のRTX 4090上でのGranite、そして8基のH100にFSDP2とcheckpointingを組み合わせたQwen3の3構成だ。起動後の学習時間は、それぞれ約24％、10％、21％短縮した。起動時間も含めた全体の短縮率は約11％、13％、20％だった。8基構成の結果には以前のプロトタイプが使われており、現在の修正を完全に検証した結果とは見なせない。作者のテスト表

学習エンジニアにとって、この提案は、コンパイルを有効にした後も、繰り返し使われるブロックで計算グラフが継続して再利用されているか確認する価値があることを示している。今後はメンテナーのレビュー、第三者による再現、自身の構成での再コンパイル回数、ジョブ全体の所要時間、損失曲線を確認したい。現時点のデータから、すべてのMoEモデルで同様の効果が得られるとは結論できない。