フロンティアモデル学習のスケーリングと最適化

Fireworks は先日、フロンティアモデル、特に数兆パラメータの MoE モデルの学習をスケーリングおよび最適化する方法について詳細な技術ブログ記事を公開しました。記事では、Cursor がリリースした Composer 2 を例に挙げています。このモデルは複数のベンチマークでトップスコアを達成しており、その強化学習推論インフラは Fireworks によって支えられています。

数兆パラメータの MoE モデルの学習は、メモリの壁と複雑なクラスタオーケストレーションによって長らく制約されてきました。Fireworks のトレーニング SDK は、モデルカタログ、並列化戦略、精度最適化、メモリ最適化を含む完全なソリューションを提供します。現在、このプラットフォームは Qwen3.5、Kimi K2.5、Llama 3.3 などを含む多様なモデルをサポートし、LoRA と全パラメータ学習の両方に対応しています。

LoRA モードでは、学習可能なパラメータはわずかですが、凍結されたベースモデルは依然として GPU メモリに常駐する必要があります。Fireworks は、低精度エキスパート量子化技術を使用して、凍結されたエキスパートの重みを圧縮形式で保存し（約 4 倍のメモリ節約）、フォワードパス中に動的に逆量子化します。同時に、オプティマイザ状態のオフロードにより、CPU と GPU 間で状態を移動し、さらにメモリを解放します。例えば、Qwen3-30B MoE モデルでは、ピーク GPU メモリが 40% 以上削減され、スループットは低下しません。さらに、マルチセッション LoRA により、複数のクライアントが独立して異なる LoRA アダプタをロードおよびホットスワップでき、RL ワークフローに適しています。

全パラメータ学習はさらに大きな課題を伴います。すべての重みに勾配、オプティマイザ状態、アクティベーションが必要です。Fireworks のエンジンは、単一の設定から 4 次元の並列化（FSDP、パイプラインパラレリズム、コンテキストパラレリズム、エキスパートパラレリズム）を構成します。各戦略は異なるボトルネックに対処します。たとえば、8B の dense モデルは FSDP のみで十分ですが、256K コンテキストの 1T MoE モデルは 4 つすべてを必要とします。コンテキストパラレリズムは長コンテキストスケーリングの主要な軸として、シーケンスを複数のノードに分割します。エキスパートパラレリズムは DeepEP を使用して効率的なトークンルーティングを実現します。混合アテンションアーキテクチャの場合、エンジンは異種レイヤータイプのコンテキストパラレリズムもサポートしています。

計算精度に関しては、Fireworks は Blackwell アーキテクチャの MXFP8 ネイティブグループ化 GEMM を利用し、エキスパート計算で大幅な高速化を実現しています。BF16 と比較して、DeepSeek V3 クラスのエキスパート形状では、フォワードおよびバックワードパスの両方で明らかな速度向上が見られ、対称 KL ダイバージェンスは 0.0063 未満に保たれています。アテンション計算には、Blackwell SM100 向けに設計された FA4 カーネルが統合されており、DeepSeek V3 や Kimi K2.5 で使用されるネイティブのマルチヘッド潜在アテンション（MLA）形状をサポートしています。さらに、FP8 量子化対応学習（QAT）により、トレーニング中の擬似量子化演算が推論エンジンの計算と完全に一致し、トレーニング後の量子化による誤差を排除します。

強化学習の損失計算に関しては、Fireworks は融合損失計算パスを提供し、損失計算をフォワードパスに直接組み込むことで、余分なフォワードパスを排除します。GRPO、DRO などの一般的なアルゴリズムでは、1.3 倍から 2 倍の高速化が達成され、数値精度は 2 パスリファレンスと一致しています。ストリーミングパイプラインパラレリズムはスケジューリング戦略を再設計し、データが到着次第実行を開始することで、バッチ蓄積のボトルネックを排除し、RL シナリオでの最初の結果レイテンシを桁違いに改善します。

記事の最後では、Fireworks が GB200 クラスタ上で 100 万トークンを超えるコンテキストでの数兆パラメータ MoE モデルの学習を検証したと述べられています。これは業界初の成果とされています。現在利用可能なトレーニングシェイプは最大 256K トークンのコンテキストをサポートしており、チームはさらにこの限界を押し広げています。