Scaling and Optimizing Frontier Model Training (Fireworks AI)

TL;DR

Fireworks describes the trainer-side half of the RL infrastructure behind Cursor’s Composer 2 and now offered as a managed training shape for Kimi K2.5, Qwen3.5 397B, MiniMax M2.5, and Nemotron 3 Super — i.e. the same trillion-parameter MoE class. The recipe is a composition of four parallelism dimensions (FSDP / Pipeline / Context / Expert) chosen automatically per “training shape,” MXFP8 native grouped-GEMM expert kernels on Blackwell, FA4 with native MLA shapes (QK=192, V=128), per-RL-algorithm fused-loss kernels (GRPO/DRO/DAPO/GSPO/CISPO/SFT) replacing the two-forward-pass forward_backward_custom path, FP8 QAT that mirrors inference numerics, streaming pipeline schedules redesigned for RL’s asynchronous data arrival, and LoRA-side tricks (frozen-expert low-precision storage with on-the-fly bf16 dequant, optimizer-state CPU offload, multi-session adapter hot-swap). Validated training of a 1T+ MoE at 1M-token context on GB200 — the largest model × context combination reported in the published comparison table.

Key claims

Training shapes catalog covers both LoRA and full-parameter modes across Qwen3.5 397B-A17B (32×B200, 262K ctx), Qwen3 235B (16×B200, 128K), Kimi K2.5 (64×B200, 256K), MiniMax M2.5 (16×B200, 192K), Nemotron 3 Super 120B Hybrid Mamba-MoE (16×B200, 128K), Nemotron Nano 3 30B-A3B (8×B200, 262K), Llama 3.3 70B (8×B200, 128K), and Qwen3-VL 8B (4×H200, 65K) — “the broadest set of fine-tunable frontier MoE models available on any training platform” [§“Training Shapes catalog”].
Composer 2 was trained on this stack; numbers cited from Cursor’s release: 61.3 CursorBench, 73.7 SWE-bench Multilingual, 61.7 Terminal-Bench [§Intro].
LoRA on a 1T-class MoE is a memory problem, not a parameter-update problem: frozen experts are stored in a reduced-precision packed format (~4× memory reduction) and dequantized to bf16 on the fly during forward pass; because experts are frozen, there is no gradient-precision loss. For Kimi K2.5 (384 experts) this is the difference between multi-node and single 8-GPU node [§“Low-precision expert quantization”].
Optimizer-state CPU/GPU offload on Qwen3-30B MoE (128 experts, 8×H200) reduces peak GPU memory by >40% with bit-identical training results to the non-offloaded baseline [§“Optimizer state offloading”].
Multi-session LoRA serves multiple clients hot-swapping adapters on a shared frozen base, with verified cross-GPU parity and zero state leakage between rapid adapter switches [§“Multi-session LoRA”].
The four parallelism dimensions are composed from a single configuration: FSDP, Pipeline, Context, Expert — a dense 8B model may use only FSDP, while a 1T MoE at 256K context uses all four. The Training SDK selects the combination per shape rather than asking the user to specify [§“Composable 4D parallelism”].
MXFP8 native grouped GEMMs on Blackwell exploit block-scaled tensor-core MMAs: hardware dequantizes during the systolic-array multiply rather than in a separate kernel. On DeepSeek V3-class expert shapes (32 experts/rank, hidden 7168, intermediate 2048) this delivers “significant speedup over BF16” in both forward and backward; symmetric-KL divergence vs BF16 stays below 0.0063 across all tested configurations, well within Fireworks’ 0.01 acceptance threshold [§“MXFP8 expert kernels”].
Attention uses FA4 (CuTeDSL) kernels handling the native MLA shapes (QK=192, V=128) used by DeepSeek V3 and Kimi K2.5, without padding or reshaping, in both forward and backward; Fireworks collaborated on the public FA4 backward PR (Dao-AILab/flash-attention#2270) for these dimensions [§“FA4 / attention”].
FP8 Quantization-Aware Training is supported: fake-quantization ops in training exactly mirror the inference engine’s math, so QAT-trained models deploy at reduced precision with no post-training-quantization surprises [§“FP8 QAT”].
Custom RL loss via forward_backward_custom requires two forward passes (one for log-probs, one fused with backward through a CE surrogate); per-algorithm fused losses collapse this to one forward + one backward and yield ~2× speedup for GRPO, ~1.7× for DRO, ~1.4× for DAPO, ~1.3× for SFT on a Qwen3.5-35B MoE / 8×H200 setup. Numerically identical at step 1; within MoE-routing non-determinism thereafter [§“Fused RL losses”].
The HTTP-based training API redesigns the pipeline-parallelism schedule for streaming data arrival rather than batch-oriented loading — RL rollout data is asynchronous. Reports “up to an order-of-magnitude improvement in first-result latency” for RL workloads at low input QPS relative to batch size, with exact gradient parity to accumulated batch [§“Streaming pipeline schedule”].
Numerical-correctness validation is non-trivial: Qwen3.5-35B alone required solving 9 distinct gradient-correctness bugs across shared experts, router gates, GQA, and DeltaNet layers; every model ships with SFT memorization validation confirming end-to-end correctness [§“Numerical parity”].
Validated training of a trillion-parameter MoE at >1M-token context on GB200 clusters, claimed as the first published demonstration at this combination of scale × context. The comparison table places this above DeepSeek V3 (671B / 128K), Llama 3.1 (405B / 128K dense), Qwen3-235B (235B / 262K), and Nemotron 3 Super (120B / 1M Hybrid Mamba-MoE — smaller and architecturally easier at long context) [§“1M context”].

Method

Fireworks exposes per-model “training shapes” — JSON-resolvable profiles that fix the GPU layout, parallelism configuration, and deployment topology for a given (base model, mode, context length) tuple. resolve_training_profile() returns the shape; the user picks LoRA vs full-parameter and the SDK launches.

For LoRA the engineering problem is fitting the frozen base. Frozen-expert weights are quantized to a packed low-precision format (~4× compression vs bf16), dequantized on the fly during forward; gradients are clean because experts are frozen. CPU/GPU optimizer-state offloading reclaims another >40% on dense or smaller-MoE models with bit-identical training. Multi-session LoRA puts a shared frozen base behind hot-swappable adapter handles with verified state isolation across rapid switches.

For full-parameter training the engine composes four parallelism axes — FSDP, Pipeline, Context, Expert — from a single config string per shape. Expert-count-vs-GPU-count modular arithmetic and per-MoE-layer all-to-all dispatch are the two MoE-specific bottlenecks the schedule has to handle; dense models avoid routing but still hit memory walls. Numerical precision is asymmetric: MXFP8 native grouped GEMMs on Blackwell for MoE expert compute (with block-scale dequant inside the tensor-core MMA), FA4-CuTeDSL kernels on the native MLA QK=192/V=128 shape, FP8 QAT mirroring the inference engine for deployable low-precision models.

RL custom losses are decomposed into a generic two-forward-pass path (forward_backward_custom — flexible enough for any Python-defined objective) and a built-in fused path that drops the second forward for GRPO / DRO / DAPO / GSPO / CISPO / SFT. The pipeline schedule is redesigned to begin execution as data arrives over HTTP rather than wait for batch accumulation, so RL rollouts streaming in at low QPS don’t stall the trainer; gradients are identical to the equivalent batched run.

Results

Composer 2 (Cursor): 61.3 CursorBench / 73.7 SWE-bench Multilingual / 61.7 Terminal-Bench, trained on this infrastructure [§Intro].
Optimizer-state offloading: >40% peak-GPU-memory reduction on Qwen3-30B MoE / 8×H200, bit-identical to baseline [§“Optimizer state offloading”].
MXFP8 vs BF16 on DeepSeek V3-class expert shapes: “significant speedup” both directions; symmetric-KL ≤0.0063 across configurations [§“MXFP8 expert kernels”].
Fused RL losses vs forward_backward_custom (Qwen3.5-35B / 8×H200): GRPO ~2×, DRO ~1.7×, DAPO ~1.4×, SFT ~1.3× [§Table].
Streaming pipeline schedule: up to an order-of-magnitude improvement in first-result latency at low input-QPS-to-batch-size ratios, with exact gradient parity [§“Streaming pipeline schedule”].
1T+ MoE at >1M context validated on GB200 — no other published system has demonstrated this combination of total parameters × max train context [§Comparison table].

Why it’s interesting

This is the trainer-side companion to two pages already heavily populated. The fused-loss-kernel result connects directly to Composer 2 Technical Report — Cursor’s report describes the asymmetric NVFP4-forward / MXFP8-backward precision recipe, the Dr. GRPO objective, and the FA4 backward kernel collaboration for QK=192/V=128, but stops at a credit to the Colfax / FA-team collaboration. This Fireworks post is what the other side of that collaboration looks like as a productized stack, and it generalizes the recipe: same MXFP8 expert kernels, same FA4 MLA shapes, same streaming RL infrastructure are now also being applied to Kimi K2.5, Qwen3.5 397B, MiniMax M2.5, and Nemotron 3 Super. For the Reasoning RL page, this is the first filed datapoint that disentangles “fused on-device loss” (≈2× for GRPO) from the algorithmic choice of the loss itself — a knob orthogonal to the reward-shaping / KL-estimator / no-KL discussion already on that page.

For IO-Aware Kernel Design, this post is direct corroboration of the FA4-MLA-shape backward kernel (FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling) and a new instance of the template applied to RL-loss kernels — the two-forward “extract logprobs, then fuse loss through CE surrogate” pattern is exactly the produce/consume pair that an IO-aware fusion eliminates. For Training stability at scale, the FP8 QAT + MXFP8 grouped-GEMM + symmetric-KL ≤0.0063 acceptance criterion is an empirical floor for “how much numerical drift is tolerable for training-inference parity at MoE scale” — a calibration the wiki didn’t have. The published numerical-bug count (9 distinct gradient-correctness bugs for one model) is also a sobering datapoint for anyone porting an MoE recipe across architectures.