NVIDIA Megatron Core 0.20.0：FSDP v2統合、MoE推論の改善

1. モデルを分担する学習構成を細かく設定可能にした

中心となるのは、既存の学習基盤に機能を接続する実験段階のモジュール、Megatron-FSDP v2アダプターだ。FSDPは学習に必要な状態を複数のデバイスに分散して保持する方式で、今回のアダプターは入力コンテキスト・エキスパートモジュール・モデルのステージごとの並列化と、ハイブリッドデータ並列シャーディングをサポートする。Megatron Bridgeとの統合も追加された。

ハイブリッドFSDPでは、デバイス間の接続網であるNVLinkで結ばれた領域の内側と外側に、異なる分割レベルを指定できる。モデル全体を一様に計算する密な部分と、一部のエキスパートモジュールを選んで計算する混合エキスパート部分にも、別々の設定を適用できる。ハードウェアの接続構造とモデルの部分に合わせて分割方式をそれぞれ選ぶ構成だ。

学習状態の保存と復元も統合の対象となる。PyTorchの分散チェックポイント方式で、モデルとオプティマイザーの状態を保存・読み込みする経路が追加された。オプティマイザーは学習中に重みを更新する構成要素であり、この変更には、モデルだけでなく学習を継続するための状態もまとめて扱えるという意味がある。

2. リクエストの組み合わせが変わってもMoEの計算を一定に保つ

混合エキスパートモデルであるMoEの推論に、バッチ不変の実行が導入された。同じリクエストでも、一緒に処理するリクエストによって計算結果が揺らぐことを防ぐ経路だ。対応精度は16ビット浮動小数点形式のBF16と、マイクロスケーリング8ビット数値形式のMXFP8だ。対象には、複数の演算をまとめて行うvLLMの融合エキスパート演算と、エキスパートの重みを行列として乗算するエキスパート行列積が含まれる。

リクエストごとの乱数シード処理も強化した。乱数生成の初期値を指定したリクエストが、バッチ構成やリクエストの移動、他のリクエストによる乱数消費の影響を受けないようにした。エキスパート計算の一貫性とともに、リクエストを混在させて処理するサービスで結果を再現するうえで重要な変更だ。

推論エンジンでは、エキスパートテンソル並列化もサポートする。エキスパート内部のテンソル計算を複数のデバイスに分割する方式だ。学習アダプターとは別に推論エンジンの並列実行の選択肢が加わったことで、学習構成とサービス構成をそれぞれ検討する根拠が生まれた。

3. 演算効率の利点には構成検討の負担も伴う

新しい学習アダプターは、MXFP8演算と量子化パラメーターの収集、混合精度、複数の数値精度を併用するオプティマイザーをサポートする。量子化は、数値をより少ないビットで表現する方式だ。メモリと演算の効率化を狙う選択肢が増えた分、導入チームは精度だけでなく、並列化と状態復元の経路も整合させる必要がある。

この構成検討には具体的な条件がある。アダプターのテンソル並列サイズとエキスパートテンソル並列サイズは、それぞれ1でなければならず、オプティマイザーのチェックポイント機能は現時点でエキスパート並列レイアウトを対象外としている。

対応する密なモデル構成には、学習イテレーション全体をCUDAグラフとしてキャプチャし、再生する機能も追加された。GPUの処理手順をまとめて再実行する経路だ。推論ではvLLMのエキスパート行列積バッファーを保持してグラフ実行に利用するようにし、同社発表による融合MoE経路の性能改善もリリースの変更一覧に含まれた。

4. 保存と接続処理による推論の停滞を減らした

運用面では、生成されたトークンを保存用に準備する処理を推論ループの外へ移した。低速なチェックポイント書き込みがデコードを停止させていた経路を見直したものだ。結果の保存処理中も生成を継続することに重点を置いた変更となる。

HTTPフロントエンドはデータ並列サイズに合わせて拡張され、テキストをモデルの入力単位に分割するトークン化処理も、リクエスト処理を調整するイベントループの外へ移した。複数のリクエストを受け付ける入口と、実際の推論実行の間の処理フローを改善する取り組みだ。

クライアントの接続が切れると、該当リクエストの生成も停止する。ユーザーが離れた後も残った処理が処理能力を占有する状況を防ぐ措置だ。バッチ不変の実行が出力の再現性を扱うのに対し、この変更はサービスが使う処理能力の管理に近い。

今回のリリースを導入する価値は、学習チームには分割構成の柔軟性に、推論チームにはリクエストの組み合わせや運用フローが変わってもサービスを一貫して実行できる点にある。