miniTE: Diagnosing and Bridging the Execution Gap in Native NVFP4 Reinforcement Learning

AI summary

Overview

Research area: Systems and numerical methods for reinforcement learning (RL) post-training of large language models — specifically native 4-bit (NVFP4) computation and cross-engine numerical consistency between a sampler (inference engine) and a learner (training engine).

Technical level: Advanced. The paper assumes familiarity with RL fine-tuning (PPO-style clipped objectives, importance sampling), Mixture-of-Experts models, and low-precision number formats, though its core argument is stated plainly.

Scope: A diagnostic-plus-systems study showing that matching bit-widths and weight versions across two engines is not enough for stable NVFP4 RL, and presenting a lightweight framework (miniTE) that closes the gap while speeding up training.

What This Paper Is About

In RL post-training, one engine generates text (the sampler) and another scores and updates the model (the learner). If those two engines compute slightly different probabilities for the same token, those differences leak into the importance ratios and clipping decisions that drive learning. The paper shows that simply giving both engines NVFP4 weights and identical parameters does not make them numerically agree — the residual gap lives in unstated implementation details like scale domains, rounding, accumulation order, and operator composition. The goal is to diagnose that gap precisely and then build a practical native W4A4 (4-bit weights and activations) expert-computation path that learns stably and runs faster than both BF16 and naive Transformer Engine NVFP4.

Key Contributions

1. Execution-gap diagnosis. Frozen-policy scoring, routing-conditioned comparison, and scaling controls show that numerical differences persist after nominal recipe agreement — even with matched expert assignments at the observed layer — and a "strict execution" reference supplies an actionable target for alignment.

2. A practical native NVFP4 RL path. An integrated W4A4 expert-forward path (per-token activation outer scales, FP32 epilogue rescaling, BF16 backward on saved high-precision operands) sustains long-horizon learning and beats naive NVFP4 on all four mathematics benchmarks without changing the RL objective.

3. A lightweight, numerically faithful framework. miniTE exposes configurable numerical recipes, replaceable expert operators, and reusable diagnostics, validated by operator-level numerical checks (28 bitwise checks against the upstream reference) plus complete learner benchmarks.

4. Joint quantization/grouped-GEMM co-design. A shared producer–consumer layout (packed FP4 codes, block scales, and outer-scale metadata written directly into contiguous GEMM-consumed buffers) plus amax computed during permutation rather than in a separate scan.

Main Findings

• The gap exists before any optimizer step. At a shared frozen parameter version, the P99 of the absolute learner–sampler log-probability difference is 0.1271 nats under BF16 versus 0.2911 under NVFP4; the paired increase is 0.1641 with a 95% confidence interval of [0.1311, 0.2005]. Statistics come from 32 prompts and 69,088 matched tokens with 2,000 bootstrap resamples.

• Clipping is reached more often under NVFP4. Numerical-only ratios fall outside the analysis clipping bounds [0.80, 1.28] for 1.203% of NVFP4 tokens versus 0.627% under BF16.

• Matched routing does not fix it. With expert assignments fixed at each target layer (set-match 1.0, no invalid experts), FC1 symmetric relative-L2 discrepancies still range 0.0381 to 0.6124 and FC2 0.0810 to 0.7624, peaking at layer 32.

• Finer scales alone are not the answer. Changing only the sampler activation outer scale from per-tensor to per-token leaves FC1/FC2 discrepancies unchanged, and an early per-token candidate actually shows a larger tail than standard NVFP4 (P99 0.3179, clipping potential 1.282%).

• Strict execution can align the engines. A configuration matching learner kernels, formulas, and reduction orders to the sampler (retaining native NVFP4 expert forward and BF16 gradients, run without TIS) records a log-probability difference of about 1e-8 and maintains reward through a region where the unaligned run deteriorates.

• Training quality beats naive NVFP4, trails BF16. At step 999: AIME 2024 36.67 (vs 26.67 NVFP4, 41.11 BF16), AIME 2025 24.44 (vs 20.56, 31.11), MATH-500 87.45 (vs 86.90, 90.05), GSM8K 94.79 (vs 94.71, 93.57), IFEval 38.40 (vs 43.99, 38.36). The five-task average is 56.35%, a gain of +1.78 points over NVFP4 and 2.49 points below BF16. IFEval is the one task where miniTE ranks below NVFP4.

• Complete learner throughput. On eight B300 GPUs at 524,288 tokens per update, the optimized path reaches 33,540 tokens/s = 1.39× BF16 (24,134) and 1.47× TE NVFP4 (22,801). The repeated ten-trial core study reaches 32,517 tokens/s and a trial-ID-paired trimmed speedup of 1.334× over BF16; dividing aggregate throughputs instead gives 1.344×.

• Large module-level gain under a common recipe. Expert forward is 2.46× faster (geometric mean) than the pinned upstream per-token implementation TE U3045 across FC1/FC2 and four workload sizes, ranging 1.587× to 3.065×.

• BF16 backward beats MXFP8 backward. Holding expert forward fixed, BF16 backward lowers forward-plus-backward latency in all eight module cells with a geometric-mean speedup of about 1.21×, and its largest gradient relative-L2 difference from an FP32 reference is 0.001663 versus 0.037653 for MXFP8.

• Balanced module timings mislead. miniTE with BF16 backward is slower than BF16 in seven of eight module forward-plus-backward cells. A routing control on GB200 shows random weights with mock data activate 50.8 experts per layer on average versus 126.2 for pretrained weights on real DAPO text, out of 128 — so complete-learner measurements need real weights and real text.

• Integration is the new bottleneck. Standalone NVFP4 serving reaches 80,676 output tokens/s versus 49,119 for BF16 (1.64×). But in a colocated RL window, learner time falls only from 47.04 s to 42.56 s while weight conversion/synchronization rises from 9.94 s to 55.26 s, summing to 56.98 s for BF16 versus 97.82 s for miniTE.

Methodology in Plain English

The researchers first held everything else constant to isolate the question. They took a frozen Qwen3-30B-A3B-Base checkpoint and, using teacher forcing, had both engines score the exact same token sequences with the same token IDs and the same parameter version. Any difference in the resulting log-probabilities must then come from implementation, not from weight staleness or sampling randomness. They compared BF16 against NVFP4 to measure the size of that difference, and used bootstrap resampling over prompt clusters to attach confidence intervals.

Next they eliminated the obvious suspect — expert routing. In a Mixture-of-Experts model, different experts may be selected on the two engines. They froze the expert assignments at each target layer and re-measured; discrepancies remained, which localizes the problem to the operators themselves rather than to which experts got picked. They also tried changing just the scale granularity on the sampler side to see whether that alone would restore agreement. It did not.

With the problem pinned down, they built miniTE around an explicit "execution specification": a record of scale domains, rounding, layouts, and where casts and router coefficients sit. Their chosen recipe gives each dispatched activation row its own outer scale, applies the outer factors in FP32 in the epilogue after the native block-scaled matrix multiply, and casts to BF16 only at the end — so the scaling connects to native Tensor Core arithmetic without ever materializing decoded BF16 operands. Hadamard transforms and stochastic rounding are switched off, so a token's FP4 codes do not depend on which other tokens happen to be grouped with it (verified by a context-swap test: target-row codes and outputs stayed bitwise identical, while an RHT-enabled per-tensor reference changed in five of six swaps). Backward is deliberately decoupled: a straight-through estimator using saved original BF16 activations and weights, avoiding backward quantization and its preparation cost.

Finally, they treated quantization and grouped GEMM as one dataflow. The quantizer writes packed codes, block scales, and outer-scale metadata straight into contiguous buffers the GEMM consumes, with device-side offsets marking expert boundaries; activation maxima are gathered during permutation and padded alongside the activations, eliminating a separate amax scan. Efficiency was then measured two ways — isolated module cells and complete learner updates on fixed text — because the paper is explicit that these have different references and timing boundaries. Numerical fidelity of the optimizations was checked with bitwise comparisons against the upstream reference.

Why This Matters

Impact on research. The paper reframes low-precision RL as a cross-engine numerical execution problem rather than a quantization-error problem. Its finding that matching weights, formats, and even expert assignments is insufficient — and that a sampler-only scaling change cannot buy agreement — suggests that papers reporting low-bit RL results should specify the full operator path, not just a format name. The distinction between three separate comparisons (quantization error within a stack, cross-stack execution difference, and implementation difference for one operator) is a reusable measurement framework.

Real-world applications.

• Cost reduction for RL post-training: complete learner updates at 33,540 tokens/s on a single B300 node rather than a BF16 baseline is a direct reduction in the wall-clock and GPU-hour cost of math-reasoning post-training runs.
• Faster inference serving: the 80,676 versus 49,119 output tokens/s serving measurement matters for any deployment that uses NVFP4 sampling.
• Reasoning-model development pipelines: teams training mathematical reasoning with DAPO-style data can adopt the recipe while monitoring the diagnostics for silent quality regressions.
• Reproducible low-bit experimentation: released Dockerfiles, reference configurations, and single-command launch scripts lower the barrier to testing custom numerical recipes.

Industry relevance. The paper is explicit that the naive path can lose throughput — TE NVFP4 reaches only 0.945× BF16 in the optimized study — so "use FP4" is not automatically a win. Its identification of weight conversion and synchronization (9.94 s to 55.26 s) as the dominant integration cost gives systems teams a concrete target, and the 1.64× serving gain versus an end-to-end RL iteration that sums to 97.82 s versus 56.98 s is a warning that component speedups do not compose for free.

Future Directions

• Close the remaining quality gap. miniTE still trails BF16 by 2.49 points on the five-task average and trails strict execution in cross-engine agreement. The authors note the gap cannot currently be attributed to quantization error versus residual execution difference, and propose combining tighter execution alignment with format-aware error compensation.

• Cut the learner–sampler boundary cost. Weight conversion and synchronization dominate the colocated measurement. Reducing repeated conversion and packing at that boundary is named as a priority, with the caveat that current measurements cover selected stages rather than complete RL iterations.

• Fuse amax, quantization, and GEMM online. The authors flag this as a complementary direction whose numerical behavior and performance need joint validation, especially when scales evolve during accumulation, and note the SM resource constraints involved in coordinating scale updates, accumulator rescaling, and data movement.

• Broaden the evidence base. The learning results cover one MoE model with one run per configuration, and the implementation gains were measured on Blackwell hardware. Validation across more models, tasks, independent training seeds, and hardware platforms is left as future work.

Target Audience

Researchers and engineers working on low-precision training or inference for large language models, particularly those running RL post-training pipelines with separate sampler and learner engines. It is also valuable for systems engineers who own MoE grouped-GEMM kernels and quantization dataflows, and for practitioners who need to decide whether native FP4 is worth adopting — the paper's honest accounting of where FP4 loses (module-level forward-plus-backward in seven of eight cells, IFEval at 38.40 versus NVFP4's 43.99, and the synchronization bottleneck) makes it useful as a decision document, not just a results report. Readers without background in RL objectives or block-scaled number formats will find the middle sections dense, but the diagnostic logic in the first half is accessible.

Authors’ abstract

Native NVFP4 promises to accelerate reinforcement learning (RL) for large language models (LLMs), but achieving stable learning across separate training and sampling engines remains challenging. We show that matching model weights and numerical formats is insufficient to ensure consistent policy evaluation, highlighting the importance of cross-engine numerical execution. We introduce miniTE, a lightweight framework that combines execution diagnostics, cross-engine numerical alignment, and efficient native W4A4 (4-bit weights and activations) GEMM computation for LLM RL. It improves numerical execution consistency without modifying the RL objective, while reducing implementation overhead to realize the benefits of native low-bit computation. Long-horizon RL experiments on Qwen3-30B-A3B demonstrate stable training and benchmark performance close to BF16, outperforming naive Transformer Engine (TE) NVFP4 in evaluation quality. On single NVIDIA B300 GPU node, the optimized learner reaches 33,540 tokens/s, yielding 1.39x and 1.47x the throughput of BF16 and TE NVFP4, respectively, for complete learner updates. We released the framework, kernels, diagnostics, and training configurations to make native low-bit RL easier.