TRIAGE: Direction-Aware Mismatch Stabilization of Native NVFP4 Reinforcement Learning

1 Introduction

• Longer generation trajectories make rollout a dominant wall-clock bottleneck in RL training for large language models.Native Blackwell W4A4 operations offer up to 4× peak matrix-compute throughput of BF16, motivating NVFP4 for long-sequence rollout.Source: §1 · p. 1
• Magnitude-only corrections cannot determine whether a policy-gradient update will reduce or amplify learner–sampler mismatch.The paper asks what determines whether mismatch contracts or amplifies under policy optimization.Source: §1 · p. 1
• In native NVFP4 runs, mismatch-amplifying updates become imbalanced before mismatch is broadly severe, while large-gap tokens cluster in short response portions.TRIAGE responds with segment-level diagnosis, selective gating of negative-advantage, negative-gap tokens, update rebalancing, and bounded repair of residual severe mismatch.Source: §1 · p. 2
• Up to 2.3× rollout throughput over BF16, with stable optimization and near-BF16 average performance across five mathematical reasoning benchmarks, is reported for TRIAGE.The method retains native NVFP4 W4A4 forward execution on both sampler and learner.Source: §1 · p. 2
• TRIAGE characterizes mismatch by policy-gradient direction and identifies a directional, locally concentrated precursor to native NVFP4 instability.Its selective stabilization combines segment-level risk diagnosis, direction-selective token gating, response-level rebalancing, and targeted repair while preserving contracting updates.Source: §1 · p. 2
• Up to 1.3× end-to-end RL speedup with BF16-level performance is reported for native W4A4 execution on sampler and learner forward paths.The experiments total more than 34,500 GPU-hours on B300 GPUs across dense and MoE models.Source: §1 · p. 2

2.1 Group-Relative Policy Optimization

• GRPO assigns each response a group-relative advantage from its reward and optimizes token-level importance ratios with a clipped policy objective.The advantage is normalized by the group reward mean and standard deviation and assigned to valid tokens of that response.Source: §2.1 · pp. 2–3

2.2 Decoupled Rollout and Training

• In decoupled RL, the sampler generates tokens while the learner computes gradients, and their same-prefix token probabilities can differ.Under synchronous weight transfer with one update per rollout, exp(δ_i,t) carries the update’s off-policy error; magnitude-only corrections do not indicate how the gradient changes the gap.Source: §2.2 · p. 3

2.3 Low-Precision Rollout with NVFP4

• NVFP4 reduces matrix-computation cost and model-memory traffic for rollout, but quantized weights and activations perturb hidden states and logits.During decoding, these local perturbations can alter sampled tokens and subsequent prefixes, accumulating into trajectory-level divergence.Source: §2.3 · p. 3

3.1 Directional Dynamics of Policy Mismatch

• The local product δ_i,tg_i,t amplifies mismatch when positive and contracts it when negative; this is a directional risk indicator, not an unconditional predictor of realized token motion.Cross-token coupling, optimizer-history terms, and higher-order effects can alter an individual token’s realized motion.Source: §3.1 · p. 4
• The two amplifying regions are negative-advantage, negative-gap tokens and positive-advantage, positive-gap tokens.For vanilla GRPO, the directional criterion reduces to the sign interaction A_iδ_i,t.Source: §3.1 · p. 4

3.2 Directional Asymmetry Precedes Negative-Tail Explosion

• By pre-terminal, ρ_asym rises from 1.02 to 2.32 on 4B and from 1.34 to 1.85 on 30B, while at least 83.0% and 91.9% of tokens remain within |δ| < 0.05, respectively.The two amplifying regions become imbalanced before broad magnitude growth is visible.Source: §3.2 · p. 5, Figure 2 · p. 5
• At terminal, the 5% gap quantile q0.05(δ) reaches −42.75 on 4B and −10.94 on 30B, after the earlier directional-asymmetry warning interval.The negative tail expands abruptly in the terminal stage.Source: §3.2 · p. 5, Figure 2 · p. 5

3.3 Localized Accumulation before Global Failure

• From healthy to pre-terminal, amplifying-tail tokens fall from 4.8–4.9% to 2.8–3.6% globally, yet the top 5% of segments contain about one quarter of them—roughly five times uniform placement.Adjacent tail-heavy persistence increases from about 2.6× to 3.8–4.0× the chance rate, indicating local reorganization rather than an expanding global tail.Source: §3.3 · p. 5, Figure 3 · p. 5
• At pre-terminal, 97.2% of 4B and 79.9% of 30B responses containing tail-heavy segments still have |δ̄_i| < 0.05; by terminal, roughly half of eligible segments are tail-heavy.These localized patterns motivate segment-level diagnosis and restricting intervention to 𝒜− tokens rather than suppressing every token in an affected response or segment.Source: §3.3 · p. 5, Figure 3 · p. 5

4 Method

• TRIAGE combines a reallocation term for negative-advantage responses with a bounded corrective term for positive-advantage responses.The modifications act on the effective update coefficient of the underlying importance-corrected objective.Source: §4 · p. 6

4.1 Segment Diagnosis with Direction-Selective Gating

• Segment statistics identify where intervention is warranted, but only negative-advantage, negative-gap tokens are gated; naturally contracting tokens in the same segment are left ungated.The method uses short contiguous segments and maps persistent displacement and sparse deep-tail mismatch to a segment gate; the exact operating rule and thresholds are deferred to Appendix C.Source: §4.1 · p. 6
• Normalization reallocates relative update mass away from the targeted amplifying subset rather than uniformly weakening the negative-advantage response.Gating can attenuate an amplifying update toward zero, but cannot reverse its sign or directly reduce the remaining learner–sampler gap.Source: §4.1 · p. 6, §4.2 · p. 7

4.2 Repair on Positive-Advantage Responses

• Repair adds a bounded corrective term on positive-advantage responses, where raising sampled-token log-probabilities on negative-gap tokens reinforces mismatch-contracting motion aligned with reward-driven learning.Applying the same correction to negative-advantage responses would oppose their policy gradients, so that side is handled by gating alone.Source: §4.2 · p. 7
• Gating operates throughout training, while repair starts at a fixed update and uses thresholds and a schedule fixed before each run.All valid tokens within an active segment share the same corrective coefficient; the repair objective targets negative segment-level displacement rather than enforcing contraction for every token.Source: §4.2 · p. 7

5.1 Experimental Setup

• The experiments use Qwen3-30B-A3B-Base as the primary model and Qwen3-4B-Base for replication and ablations, with GRPO on DAPO.Evaluation covers five mathematical reasoning benchmarks and tracks reward, policy entropy, and the mean absolute train–rollout log-probability gap.Source: §5.1 · p. 8

5.2 Experimental Results

• TRIAGE remains stable throughout the evaluated 600- and 1,700-step schedules, approximately doubling the observed stable horizon on 4B versus naive NVFP4.Naive NVFP4 becomes unstable after about 300 steps on 4B and shows rapid mismatch growth after roughly 700 steps on 30B; TIS delays the 4B failure but does not prevent it.Source: §5.2 · p. 8
• 58.49% versus 58.26% on 4B and 70.96% versus 72.41% on 30B are TRIAGE and BF16 average benchmark scores, respectively.At the shared step-300 checkpoint on 4B, TRIAGE exceeds naive NVFP4 and NVFP4+TIS by 5.93 and 2.27 percentage points, respectively.Source: §5.2 · p. 8

5.3 Rollout Throughput and Learner Overhead

• 2.27× and 2.30× BF16 rollout-throughput speedups are observed with TRIAGE at global batch sizes 256 and 512, respectively.End-to-end speedups are 1.21× and 1.30× at those batch sizes; median incremental learner-update overhead is 0.84%.Source: Table 2 · p. 9, §5.3 · p. 9

5.4 Ablation Studies

• Full TRIAGE achieves the highest reward in the 4B ablations; removing repair lowers reward despite keeping the mean log-probability gap low.TIS alone remains stable when continued from the TRIAGE checkpoint, while naive NVFP4 allows mismatch to grow sharply before late-stage collapse.Source: Figure 5 · p. 9, §5.4 · p. 9

Low-precision training

• Prior low-precision training methods address quantization error through quantization design, improved gradient estimation, LoRA, adaptive quantization noise, or importance-based corrections.The paper focuses on residual learner–sampler mismatch dynamics under native NVFP4 W4A4 execution.Source: §6 · p. 9

Training–inference alignment

• QaRL aligns the learner with quantized rollout, while QUADS uses learner-side weight-only fake quantization and rollout-side activation compensation; TRIAGE modifies the objective and retains native W4A4 on both paths.The cited approaches are presented as responses to training–inference alignment differences that can destabilize RL even with synchronized weights.Source: §6 · p. 10

Directional control

• ACRL uses advantage-guided four-quadrant probability adjustment, whereas TRIAGE analyzes local mismatch energy and identifies directional asymmetry and localized negative tails.Source: §6 · p. 10

• TRIAGE combines selective gating, within-response rebalancing, and bounded repair while preserving native W4A4 forward execution on learner and sampler.Source: §7 · p. 10
• On Qwen3-4B-Base and Qwen3-30B-A3B-Base, TRIAGE sustains stable optimization over evaluated horizons and achieves near-BF16 average performance across five mathematical reasoning benchmarks.Source: §7 · p. 10
• Native NVFP4 with TRIAGE achieves up to 2.30× rollout throughput and 1.30× end-to-end speedups over BF16 in the measured configurations.Source: §7 · p. 10

• The study does not evaluate models at 100B+ scale or substantially longer training horizons, so persistence of the directional dynamics in those settings remains open.Source: §8 · p. 10
• Experiments focus on synchronous RL; TRIAGE has not been experimentally validated with asynchronous policy staleness.The paper identifies joint analysis of quantization mismatch and asynchronous staleness as future work.Source: §8 · p. 10
• End-to-end speedup is approximately 1.3× despite faster NVFP4 rollout; synchronization and quantization-related reductions and transformations remain runtime costs.The paper proposes tighter kernel–framework integration and finer-grained repair as future directions.Source: §8 · p. 10

A.1 Fixed-Sampler Setup and Update Coefficients

• During one learner update, trajectories and recorded sampler log-probabilities are fixed, and the mismatch gradient equals the learner log-probability gradient.For valid tokens, the effective coefficient sign matches the advantage; masked tokens have zero coefficient.Source: §A.1 · p. 13

A.2 Preconditioned Gap and Energy Dynamics

• The local optimizer update separates a preconditioned current-gradient component from residuals such as history-dependent momentum and decoupled weight decay.The matrix K(P) is positive semidefinite, though its off-diagonal entries can have either sign.Source: §A.2 · p. 13

A.3 Diagonal Contributions and Their Scope

• The sign of the diagonal contribution is sign(δ_p g_p): positive values amplify the existing gap, while negative values contract it.Cross-token coupling, optimizer residuals, and higher-order effects can change the realized sign of token motion.Source: §A.3 · p. 14

A.4 Amplifying and Contracting Masses

• Advantage-weighted token masses summarize amplifying regions but do not estimate B_amp or B_con, which also depend on gradient geometry and gap magnitude.The two quantities retain optimizer and cross-token effects.Source: §A.4 · p. 14

B Characterization Protocol and Additional Analyses

• The characterization appendix specifies runs, mismatch estimands, and analysis windows, and reports indicator-agreement analyses, localization statistics, and segment-length sensitivity.Table 3 summarizes model-specific settings and the characterization updates analyzed, which are distinct from configured maximum run lengths.Source: B Characterization Protocol and… · p. 14, Table 3 · p. 15

B.1 Models, Training Setup, and Evaluation

• Characterization uses Qwen3-4B-Base and Qwen3-30B-A3B-Base on DAPO-Math-17K, with 16 prompts and 16 responses per prompt per rollout and one optimizer update.Native NVFP4 W4A4 is applied to MLP projections while specified attention, normalization, and embedding components remain BF16; training-time AIME-2024 tracking is distinct from independent final checkpoint evaluation on five benchmarks.Source: §B.1 · p. 15

B.2 Mismatch Estimands and Analysis Windows

• The 4B estimand is the ordinary learner–sampler gap, while 30B measures residual mismatch after replaying sampler MoE routing; their absolute gap magnitudes are not directly comparable.Analysis anchors are at 25%, 50%, 75%, and 100% of each trajectory with inclusive 25-update trailing windows; collapse onset is run-specific, at update 304 for 4B and update 724 for 30B, without asserting a shared failure mechanism.Source: §B.2 · p. 16, Table 4 · p. 16

B.3.2 Agreement as Mismatch Becomes Severe

• For 4B, weighted agreement rose from 47.73% across eligible late-stage A− tokens to 63.05% at |δ| > 3, whereas 30B agreement stayed below 50% and declined with severity.The 4B |δ| > 3 group contained 1,419 tokens across 316 responses, and its weighted agreement exceeded 50% at each of five saved anchors; token-wise predictiveness claims are restricted to the 4B vanilla characterization.Source: §B.3.2 · p. 17, Table 5 · p. 18

B.4.1 Segment Construction and Threshold Calibration

• Localization partitions valid completion tokens into within-response segments of W = 64 and calibrates a model-specific severe-gap threshold on a healthy window before drift.It analyzes negative-advantage responses; the fixed threshold separates localization from later global tail growth, and top-5% segment containment is not a random-position null because segments are selected by observed tail fractions.Source: §B.4.1 · p. 18, §B.4.2 · p. 19

B.4.3 Stage-Wise Localization with 𝑊= 64

• Before terminal failure, r_amp fell from 4.90% to 2.75% on 4B and 4.79% to 3.59% on 30B, while Ctok rose to 25.68% and 25.01%; terminal Hseg reached 48.64% and 57.02%.At the pre-terminal stage, 97.22% of affected 4B responses and 79.94% of affected 30B responses had |δ̄_i| < 0.05; signed cancellation contributed to dilution, while local tails could also coexist with a small mean absolute gap.Source: §B.4.3 · p. 19, Table 6 · p. 19, Figure 6 · p. 20

B.4.4 Sensitivity to Segment Length

• With W = 128, the pre-terminal small-mean-gap pattern persisted, while terminal top-5% token containment fell to 14.83% on 4B and 17.63% on 30B.At the pre-terminal stage, 96.11% of affected 4B responses and 75.99% of affected 30B responses had |δ̄_i| < 0.05; the fixed δ < −6 criterion was too sparse for stable early adjacency ratios.Source: §B.4.4 · pp. 19–20, Table 7 · p. 20

B.4.5 Response-Conditioned Localization Controls

• At pre-terminal, observed Ctok exceeded within-response shuffle medians by 7.51 percentage points on 4B and 6.51 on 30B, but P_adj did not exceed its shuffle median.The exploratory segment-order test found 2,892 versus 1,613.85 expected heavy–heavy pairs on 4B and 1,584 versus 1,055.58 on 30B; this does not establish causal propagation or exclude position trends.Source: §B.4.5 · pp. 20, 22, Figure 7 · p. 21

B.5 Full-Resolution Characterization Trajectories

• Full-resolution trajectories track marginal gap statistics and directional-region weights across updates, retaining the ordinary 4B gap and routing-conditioned 30B residual as distinct estimands.Marginal statistics include zero-advantage tokens, while directional comparisons use strict A− and A+ regions and compare raw |A_i| with response-length-normalized |A_i|/n_i weights.Source: §B.5 · p. 22
• Full-resolution trajectories show a non-monotonic directional ratio that broad windows conceal, while length normalization reduces raw-weighted imbalance, particularly on 30B.The shaded windows are descriptive summaries of one characterization run per model.Source: §B.5 · p. 22, Figure 8 · p. 23

B.6 Sensitivity to Aggregation Choices

• Widening the pre-terminal window from 25 to 50 updates changes the raw-weighted ratio from 2.319 to 2.454 on 4B and from 1.851 to 1.735 on 30B; both windows retain the imbalance.The compared intervals are 219–243 versus 194–243 for 4B and 519–543 versus 494–543 for 30B.Source: §B.6 · p. 22, Table 8 · p. 24
• Length normalization lowers pre-terminal ratios to 1.434 on 4B and 1.049 on 30B, and can change whether ratios exceed one in earlier stages.Raw advantage-weighted token mass and length-normalized coefficient mass answer different questions; the ratio describes allocation between two amplifying regions, not the balance of all amplifying and contracting contributions.Source: §B.6 · p. 24, Table 8 · p. 24

C.1 Segment Construction and the Gate Operating Rule

• Negative-advantage responses are eligible for gating; segment risk combines negative mean gaps with deep-tail concentration, and the resulting weight applies only to negative-gap tokens.Masked tokens are excluded, and the final shorter nonempty segment is retained.Source: §C.1 · pp. 24–25

C.2 From the Weighted Loss to the Effective Update

• Response normalization converts raw token weights into effective update multipliers, reallocating gradient mass within a response rather than conserving parameter-gradient norms.If all valid tokens share one raw weight, their effective multipliers are one; an attenuated raw weight need not yield a multiplier below one.Source: §C.2 · p. 25

C.3 Repair Objective and Additive Coefficient

• Repair adds a positive-advantage segment penalty with bounded coefficients; it can reinforce negative-gap tokens while also assigning coefficients to positive-gap tokens in the same segment.The coefficient bound applies to log-probability gradients, not parameter-gradient norms, which also depend on gradient norms.Source: §C.3 · p. 26

C.4 Implementation and Configuration

• The implementation computes detached gate diagnostics and live repair gaps, applies the scheduled repair coefficient, and combines the gate-weighted policy loss with repair for one optimizer step.The repair denominator counts every nonempty segment in the update batch; activation begins inclusively at the configured rollout-update index.Source: §C.4 · pp. 26–27

D.1 Main Training Comparisons

• The main comparison evaluates BF16, native NVFP4, NVFP4 with TIS, and NVFP4 with TRIAGE using native W4A4 forward execution for NVFP4 on both sampler and learner.Qwen3-30B-A3B-Base is the primary model, with Qwen3-4B-Base used for replication and ablation.Source: §D.1 · p. 27

D.2 Downstream Evaluation

• Five mathematical reasoning benchmarks are evaluated, with collapsed runs represented by their last saved checkpoint before collapse.The benchmarks are AIME 2024, AIME 2025, MATH500, AMC23, and GSM8K.Source: §D.2 · p. 27

E.1 Natural-Generation RL Measurements

• Natural-generation rollout throughput averages per-iteration token rates, while iteration wall time includes rollout, learner updates, weight synchronization, and orchestration.Ten measured observations follow two warm-up iterations; initialization, evaluation, and checkpoint saving are excluded.Source: §E.1 · p. 28

E.2 Paired Learner-Overhead Measurement

• 0.84% median learner-update overhead, with a 0.225 s median absolute increase, was measured for TRIAGE on paired frozen batches.The comparison used five paired repetitions on eight B300 GPUs.Source: §E.2 · p. 28

E.3.1 Independent Fixed-Output Configuration Selection

• 1.175× is the NVFP4/BF16 throughput ratio under independently selected configurations for the fixed 4,096-token, 512-response workload.Both selected configurations use the FlashInfer TRT-LLM MoE backend; Table 10 reports medians and observed ranges.Source: §E.3.1 · p. 28, Table 10 · p. 29

E.3.2 Concurrency and output length

• Throughput rises with concurrency in the measured 4K sweep, falls as output length grows at fixed concurrency, and the relative gain varies by operating point.Figure 9 shows the concurrency and length sweeps; its common-configuration C64/4K point is omitted from the selected-configuration length curve.Source: §E.3.2 · p. 29, Figure 9 · p. 29

E.4 Controlled Long-Tail Workload

• 1.292× and 1.278× are the NVFP4/BF16 median-time ratios for mixed-length batches of 256 and 512 responses, respectively.The corresponding equal-token uniform controls yield ratios of 1.538× and 1.513×; Figure 10 reports complete-batch wall time and late completion quantiles.Source: §E.4 · p. 29, Figure 10 · p. 30