# Aligning Training and Rollout Policies for End-to-End FP4 RL on Mixture-of-Experts LLMs

> Symmetric residual correction on both training and inference sides, unlike one-sided rollout correction, restores near-FP16 accuracy in FP4 RL for Mixture-of-Experts models.

- **Source:** [arXiv](https://arxiv.org/abs/2611.04376)
- **Published:** 2026-11-07
- **Permalink:** https://picx.dev/p/kiCFGu
- **Whiteboard:** https://picx.dev/p/kiCFGu/image

## Summary

## Summary (Overview)

- **Key finding**: RolloutResQ, a quantization recovery recipe that works for dense models in FP4 reinforcement learning (RL), **fails to transfer** to Mixture-of-Experts (MoE) LLMs, often performing *worse* than the naive FP4 baseline.
- **Root cause identified**: The failure stems from the **asymmetric application** of residual correction (RC) between training and inference sides. The MoE router's discrete expert selection amplifies this mismatch, causing training and inference policies to route tokens to different experts.
- **Proposed solution**: **MoEResQ**, a symmetric quantization-recovery recipe that applies residual correction on *both* training and inference sides, restoring training–inference alignment and closing the accuracy gap to full precision (BF16).
- **Results**: MoEResQ recovers near-BF16 accuracy on all four benchmarks (AMC23, AIME24, AIME25, MATH500). With experts-only quantization, the gap to BF16 shrinks to as little as **0.49%** (AMC23) and **0.54%** (MATH500).
- **Core principle established**: **Training–inference symmetry** is the key to accuracy recovery in FP4 RL for MoE models.

---

## Introduction and Theoretical Foundation

Large Language Models (LLMs) have achieved strong performance across mathematical reasoning, code generation, and agentic workflows [1, 2]. Reinforcement Learning with Verifiable Rewards, using policy-gradient algorithms such as GRPO [3], has become a successful paradigm for LLM post-training [4]. However, RL post-training is highly resource-intensive in both memory and computation.

**Quantization** has emerged as an effective tool to reduce memory footprint and computational cost. Prior work includes:
- **QeRL** [7]: weight-only quantization on the rollout policy with LoRA [8]
- **RolloutResQ** [9]: residual correction on rollout activations, quantizing both training and rollout sides to 4-bit

A central challenge is the **mismatch between training and rollout sides**, previously studied only for dense models [10].

**Mixture-of-Experts (MoE) LLMs** combine dynamic routing with sparse activation, achieving higher performance at reduced computational cost [11, 12]. The router performs a **discrete, input-dependent top-k expert selection**, making MoE training prone to instabilities such as routing collapse [16, 17, 21] and unusually sensitive to small perturbations of the router's input activations.

This sensitivity is amplified in low-precision RL: when rollout and training run under different numerical conditions, the router can select different experts for the same token, creating a **rollout–training mismatch** that widens the off-policy gap and can severely degrade or collapse training [22].

---

## Methodology

### 2.1 Residual Correction and RolloutResQ

Activation outliers are a primary obstacle to low-precision quantization, widening the dynamic range and collapsing smaller values to zero. **Residual correction (RC)** addresses this. For each linear layer $\ell$, the activation residual is:

$$\Delta X_{\ell} = X_{\ell} - Q(X_{\ell})$$

where $Q$ is the quantization operator. The corrected output is computed via two FP4 GEMMs:

$$\hat{Y}_{\ell} = Q(X_{\ell}) Q(W_{\ell})^{\top} + Q(\Delta X_{\ell}) Q(W_{\ell})^{\top}. \tag{1}$$

**RolloutResQ** applies this correction on rollout activations only (asymmetric). The paper uses **HiF4** [15] as the FP4 quantization format and keeps routers in BF16 (they constitute only 0.38% of active parameters).

### 2.2 Experimental Setup

- **Model**: Qwen3-30B-A3B [19]
- **Training**: GRPO on DAPO-Math-17K [30]
- **Infrastructure**: VeRL [27], vLLM [28] for rollout, FSDP [29] for training
- **Evaluation**: AMC-2023 [23], AIME-2024 [24], AIME-2025 [25], MATH-500 [26]
- **Scopes**: Full quantization (experts + attention projections) and experts-only quantization
- **Simulation**: All experiments under simulated quantization due to hardware limitations

### 2.3 MoEResQ

MoEResQ applies RC to the **same layers on both training and inference**:

- **Rollout**: Applies activation correction of Equation 1
- **Training forward pass**: Uses the same residual correction as inference
- **Training backward pass**: With input $X_{\ell}$, weight $W_{\ell}$, and output gradient $dY_{\ell}$, the gradients are $dX_{\ell} = dY_{\ell} W_{\ell}$ and $dW_{\ell} = (d\bar{Y}_{\ell})^{\top} \bar{X}_{\ell}$. Residual correction is applied to both $dY_{\ell}$ and $X_{\ell}$, retaining $dY_{\ell}, \Delta d\bar{Y}_{\ell}, X_{\ell}, \Delta X_{\ell}$ where $\Delta X_{\ell} = X_{\ell} - Q(X_{\ell})$ and $\Delta dY_{\ell} = dY_{\ell} - Q(dY_{\ell})$.

The **experts-only scope** quantizes only expert linear layers (up to 94.66% of parameters in Qwen3-30B-A3B), keeping attention layers in BF16 to limit perturbation to the router's input.

---

## Empirical Validation / Results

### Table 1: Main Evaluation Results

| Dataset | BF16 | HiF4 | RolloutResQ | MoEResQ | RolloutResQ (experts only) | MoEResQ (experts only) |
|---------|------|------|-------------|---------|---------------------------|------------------------|
| AMC23 | 85.10 | 65.23 (-19.87) | 62.34 (-22.76) | 84.45 (-0.65) | 78.75 (-6.35) | **84.61 (-0.49)** |
| AIME24 | 48.23 | 23.96 (-24.27) | 25.00 (-23.23) | 44.48 (-3.75) | 37.39 (-10.84) | **46.98 (-1.25)** |
| AIME25 | 34.38 | 20.83 (-13.55) | 19.06 (-15.32) | 29.69 (-4.69) | 29.06 (-5.32) | **32.92 (-1.46)** |
| MATH500 | 86.71 | 79.21 (-7.50) | 74.00 (-12.71) | 85.42 (-1.29) | 82.46 (-4.25) | **86.17 (-0.54)** |

*Table 1: Evaluation results for Qwen3-30B-A3B trained on DAPO-Math-17K with GRPO. Routing layers (0.041% of parameters) kept in BF16.*

**Key observations**:
- RolloutResQ performs **worse than naive HiF4** on 3 of 4 benchmarks (reversal of the dense-model result)
- MoEResQ narrows the gap to BF16 from ~20% to <5% in all cases
- Experts-only MoEResQ lands within ~0.5% of BF16 on AMC23 and MATH500

### Table 2: Training–Inference Alignment Metrics

| Metric | BF16 | HiF4 | RC Inference | RC Training | RC Both | RC (Exp. Only) Inference | RC (Exp. Only) Training | RC (Exp. Only) Both |
|--------|------|------|--------------|-------------|---------|--------------------------|--------------------------|---------------------|
| **PCC** | 0.9951 | 0.9410 | 0.9444 | 0.9580 | **0.9940** | 0.9702 | 0.9796 | **0.9952** |
| **top_k_match_frac** | | | | | | | | |
| Layer 0 | 0.9869 | 0.9801 | 0.8100 | 0.8195 | 0.9709 | 0.9867 | 0.9875 | 0.9867 |
| Layer 12 | 0.8799 | 0.6172 | 0.6279 | 0.6590 | 0.8682 | 0.7373 | 0.7879 | 0.8730 |
| Layer 23 | 0.8562 | 0.5883 | 0.5990 | 0.6402 | 0.8418 | 0.7209 | 0.7674 | 0.8506 |
| Layer 35 | 0.8459 | 0.6256 | 0.6514 | 0.6836 | 0.8359 | 0.7354 | 0.7947 | 0.8455 |
| Layer 47 | 0.8680 | 0.6926 | 0.7092 | 0.7297 | 0.8846 | 0.7748 | 0.8059 | 0.8687 |
| **topk_overlap_mean** | 0.9850 | 0.9580 | 0.9550 | 0.9600 | **0.9840** | 0.9710 | 0.9770 | **0.9840** |

*Table 2: Training-inference alignment metrics: Pearson correlation coefficient (PCC) between training and inference probabilities, fraction of tokens with exactly matching top-k expert sets (top_k_match_frac), and mean pairwise overlap (topk_overlap_mean).*

**Key observations**:
- One-sided RC (Inference or Training only) leaves alignment near HiF4 level
- **Symmetric RC (Both)** restores all metrics to near-BF16 levels
- The pattern holds for both full-scope and experts-only quantization

**Figure 1** visualizes training vs. inference probability alignment: RC Both clusters tightly along the diagonal (perfect alignment), closely mirroring BF16, while HiF4 and asymmetric recipes show significant scatter.

---

## Theoretical and Practical Implications

### Theoretical Implications

1. **Routing amplifies quantization error**: The discrete argmax over gate logits means small perturbations can flip expert selection entirely—a failure mode absent in dense models. This explains why dense-model recipes don't transfer trivially.

2. **Symmetry as a design principle**: The paper establishes that **symmetric application of recovery techniques** on training and inference sides is the critical factor for MoE models, not merely the presence of correction. Asymmetric recipes (even those that improve dense models) can actively *harm* MoE performance.

3. **Quantifying alignment**: The PCC, top_k_match_frac, and topk_overlap_mean metrics provide actionable tools for diagnosing training–inference mismatch before full evaluation.

### Practical Implications

1. **Efficiency**: The experts-only scope captures most of FP4's efficiency benefits (experts are up to 94.66% of parameters) while achieving near-BF16 accuracy—a practical recipe for deployment.

2. **Implementation guidance**: MoEResQ provides a concrete, implementable recipe for backward-pass residual correction, extending RC beyond the forward pass.

3. **Hardware constraints**: Keeping routers in BF16 (0.041% of parameters) is a low-cost safeguard that preserves routing quality.

---

## Conclusion

This paper studied end-to-end FP4 RL post-training for MoE LLMs and found that **RolloutResQ does not transfer** from dense to MoE models. The root cause is the router: since expert selection is a discrete argmax over gate logits, the mismatch from an asymmetric recovery recipe is amplified into divergent routing that degrades training–inference alignment.

**MoEResQ** applies residual correction symmetrically on both training and inference sides, restoring alignment and recovering near-BF16 accuracy. The paper establishes **training–inference symmetry** as the key principle for accuracy recovery in FP4 RL for MoE models.

**Future directions**:
- Pushing the efficiency of the recipe further while sustaining near-BF16 accuracy
- Exploring additional quantization formats and granularities beyond HiF4
- Investigating whether symmetric principles extend to other recovery techniques beyond residual correction

The work highlights that as low-precision RL extends to more complex architectures, the interaction between quantization and architecture-specific mechanisms (like routing) must be carefully considered—and that symmetry between training and inference conditions is paramount.

---

_Markdown view of https://picx.dev/p/kiCFGu, served by PicX — AI-generated visual whiteboard summaries of research papers._