Blame for code/notes/src/0108-training-inference-mismatch.tex

\documentclass[11pt]{article}
\usepackage{common}
\begin{document}
\notehead{0108}{Training--inference mismatch: truncated IS correction and why an FP32 LM head helps}{Policy-optimisation lineage, batch M1. Companion notebook: \texttt{notebooks/0108-training-inference-mismatch.ipynb}.}

\begin{problem}
\question{What happens to the gradient when the sampling engine and the training engine do not compute the same distribution?}
RL systems sample with one engine (vLLM/SGLang: fused kernels, bf16) and differentiate with another (FSDP/Megatron). Both claim to compute the same $\pi_\theta(\cdot|s)$, yet ``these engines produce numerically different outputs'' (arXiv:2510.26788 Sec.~1). The sampler's distribution $\mu_\theta$ is therefore not the trainer's $\pi_\theta$ even at $\theta=\theta_{\rm old}$. Three consequences: (i) the policy gradient estimated from $\mu$-samples is biased; (ii) PPO/CISPO ratios $r=\pi_\theta/\pi_{\rm old}$ are contaminated if $\pi_{\rm old}$ is taken from the sampler; (iii) the optimised $\theta$ is optimal for $\pi$, not for the deployed $\mu$ (a ``deployment gap''). MiniMax-M1 reports the discrepancy ``prevented reward growth'' and fixed it by raising the LM output head to FP32, improving the train/inference probability correlation from $\approx0.9x$ to $\approx0.99x$ (arXiv:2506.13585 Sec.~3.2, Fig.~3: Pearson $0.987$ before, $0.997$ after, in the figure shown). This note bounds the bias and explains \emph{quantitatively} why high-precision logits matter.
\end{problem}

\section{Notation}
\begin{center}\small\begin{tabular}{@{}ll@{}}\toprule
symbol & meaning\\\midrule
$\pi=\pi(\cdot|\theta)$, $\mu=\mu(\cdot|\theta)$ & trainer and sampler distributions of the \emph{same} weights\\
$w=\pi(o)/\mu(o)$, $\rho_t=\pi(o_t|s_t)/\mu(o_t|s_t)$ & sequence / token mismatch weight\\
$C\ge1$ & truncation level; TIS $=\min(w,C)$, MIS $=w\,\indb{w\le C}$\\
$A=R-b$; $M_A,M_s$ & $\max|A|$ and $\max_o\|\nabla\log\pi(o)\|$\\
$u=2^{-p}$ & unit roundoff, $p$ significant bits (bf16 8, fp16 11, fp32 24)\\
$z,\ e=\mathrm{fl}(z)-z$ & logits and their rounding error, $|e_a|\le u|z_a|$\\
$\chi^2_t=\chi^2(\pi_t\|\mu_t)$ & per-position chi-square divergence\\\bottomrule\end{tabular}\end{center}

\section{Derivation}
\subsection{The mismatch biases the gradient; sequence-level IS repairs it}
With the correct objective $J(\theta)=\E_{o\sim\pi}R$ but samples from $\mu$,
\begin{equation}\E_{\mu}[A\,s_\pi]-\nabla J=\sum_o(\mu-\pi)(o)\,A(o)\,s_\pi(o)=\E_\pi\big[(\tfrac\mu\pi-1)A\,s_\pi\big]\neq0\label{eq:biasmu}\end{equation}
($s_\pi=\nabla\log\pi$; arXiv:2510.26788 Eq.~3). The IS identity of 0103 repairs it exactly: $\E_\mu[w\,A\,s_\pi]=\sum_o\mu\,\frac\pi\mu As_\pi=\nabla J$ (their Eq.~5). It inherits the variance problem of 0103: $\E_\mu[w^2]=1+\chi^2(\pi\|\mu)$, and for a sequence of $T$ positions the exact statement is the recursion of 0103, $1+\chi^2_{\rm seq}=\E_\mu[w_T^2]=\E_\mu\big[w_{T-1}^2\,(1+\chi_T^2(s_T))\big]$. This is \emph{not} $\E_\mu[\prod_t(1+\chi_t^2(s_t))]$ in general: $w_{T-1}^2$ and $\chi_T^2(s_T)$ are both functions of the prefix and are correlated. The correct product identity is a martingale statement, $\E_\mu\big[w_T^2\big/\prod_{t\le T}(1+\chi_t^2(s_t))\big]=1$; the plain product formula $\E_\mu[w_T^2]=\prod_t(1+\chi_t^2)$ holds when the per-position $\chi_t^2$ do not depend on the history (independent contexts, the setting of Sec.~2.4).
\ver{toy with large logits (bf16-rounded sampler): $\mathrm{TV}(\pi,\mu)=4.5\times10^{-4}$; Eq.~\eqref{eq:biasmu} holds to $1.9\times10^{-17}$; uncorrected bias is $0.11\%$ of $\|\nabla J\|$; the IS-corrected estimator equals $\nabla J$ to $6.9\times10^{-18}$. On a state-dependent toy ($|\mathcal V|=3$, $L=8$, 511 sequences) $\E_\mu[w^2]=1.1515$ equals the recursion to $7\times10^{-16}$ and $\E_\mu[w^2/\prod_t(1+\chi_t^2)]=1$ to $4\times10^{-16}$, whereas $\E_\mu[\prod_t(1+\chi_t^2)]=1.1417\ne1.1515$.}

\subsection{Truncated and masked IS: exact bias, bounds}
Eq.~6--7 of arXiv:2510.26788 trade bias for variance with a cap $C\ge1$:
\begin{align}
\text{TIS: }&\E_\mu[\min(w,C)As]-\nabla J=-\E_\mu\big[(w-C)_+As\big],\label{eq:tis}\\
\text{MIS: }&\E_\mu[w\,\indb{w\le C}As]-\nabla J=-\E_\pi\big[\indb{w>C}\,As\big].\label{eq:mis}
\end{align}
(Both follow by subtracting the unbiased $\E_\mu[wAs]$.) Norm bounds, using $\E_\mu[(w-C)_+]=\sum_o(\pi-C\mu)_+$:
\begin{align}
\|\text{TIS bias}\|&\le\E_\mu[(w{-}C)_+|A|\|s\|]\le M_AM_s\sum_o(\pi-C\mu)_+\le M_AM_s\min\big(\TV(\pi,\mu),\,P_\pi(w>C)\big),\nonumber\\
\|\text{MIS bias}\|&\le M_AM_s\,P_\pi(w>C).\label{eq:tisbound}
\end{align}
Reasoning for the last step: for $C\ge1$, $(\pi-C\mu)_+\le(\pi-\mu)_+$ and $\sum(\pi-\mu)_+=\TV$; also $(\pi-C\mu)_+\le\pi\,\indb{\pi>C\mu}$, whose sum is $P_\pi(w>C)$. TIS therefore has bias at most the \emph{TV distance} and every weight is $\le C$, so its variance contribution is bounded; MIS throws away whole sequences with $w>C$. In both, $C\to\infty$ recovers Eq.~(5).
\paragraph{TV grows at most linearly with length.} Hybrid argument: let $H_t$ draw the first $t$ tokens from $\pi$ and the rest from $\mu$. $H_t$ and $H_{t-1}$ differ only in the conditional of token $t$, so $\TV(H_t,H_{t-1})=\E_{s_t\sim\pi}\TV_t$, and by the triangle inequality $\TV(\pi,\mu)\le\sum_{t\le T}\E_{s_t\sim\pi}[\TV_t]\le T\max_t\TV_t$.
\ver{TIS bias equals Eq.~\eqref{eq:tis} and MIS bias equals Eq.~\eqref{eq:mis} to $\le2.6\times10^{-17}$ at $C\in\{1,1.02,1.05,1.2,2\}$; chain $0\le\|\text{bias}\|\le$ bound1 $\le$ bound2 $\le M_AM_s\TV$ holds: at $C=1$: $1.1\times10^{-4}\le1.6\times10^{-4}\le2.2\times10^{-3}$; at $C=1.02$: $6.4\times10^{-8}\le6.6\times10^{-8}\le1.7\times10^{-6}$; for $C\ge1.05$ every weight is below $C$ (max $w=1.049$) and the bias is $0$. MIS at $C=1$ discards every sequence with $w>1$: bias $0.12$ vs.\ bound $4.06$. Sequence $\TV=4.5\times10^{-4}\le\sum_t\E_\pi\TV_t=6.9\times10^{-4}$.}

\subsection{Where the mismatch comes from: rounding of the logits}
A format with $p$ significant bits rounds every real $x$ to $\mathrm{fl}(x)$ with $|\mathrm{fl}(x)-x|\le u|x|$, $u=2^{-p}$. The softmax then sees $z+e$ with $|e_a|\le u|z_a|$: \emph{the absolute error is proportional to the magnitude of the logit}. We need its effect on log-probabilities.
\paragraph{Step 1 (first order).} $\log\mathrm{softmax}(z)_a=z_a-\log\sum_be^{z_b}$, so $\Delta\log\pi_a=e_a-\sum_b\pi_be_b+O(e^2)$. For i.i.d.\ errors of variance $s^2$,
\begin{equation}\Var(\Delta\log\pi_a)=s^2\big[(1-\pi_a)^2+\textstyle\sum_{b\ne a}\pi_b^2\big]=s^2\big(1-2\pi_a+\sum_b\pi_b^2\big).\label{eq:lpvar}\end{equation}
\ver{linearisation error $5.2\times10^{-8}$ (relative) at $e\sim10^{-6}$; Eq.~\eqref{eq:lpvar} vs.\ $4\times10^5$ Monte-Carlo draws: max $z$-score $1.2$ over 7 tokens.}
\paragraph{Step 2 (divergence).} With $\delta_a=\log\pi_a-\log\mu_a\approx-(e_a-\bar e)$, $\bar e=\sum_b\pi_be_b$, the chi-square is to first order the $\pi$-variance of the logit errors: $\chi^2(\pi\|\mu)=\E_\mu[(\pi/\mu-1)^2]$ with $\pi/\mu-1=e^{\delta}-1=\delta+O(\delta^2)$, and $\mu=\pi\,(1+O(e))$, so $\chi^2=\E_\pi[\delta^2]+O(e^3)=\E_\pi[(e-\bar e)^2]+O(e^3)=\Var_\pi(e)+O(e^3)$. The bound $|e_a|\le u|z_a|$ then gives:
\begin{equation}\chi^2(\pi\|\mu)\ \approx\ \Var_{a\sim\pi}(e_a)\ \le\ \E_{a\sim\pi}[e_a^2]\ \le\ u^2\,\E_\pi[z^2].\label{eq:chi2e}\end{equation}
Errors are \emph{centred}, so a peaked distribution (one dominant token) suffers less than its logit scale suggests: on the toy the per-token $\chi^2$ grows only like $\sigma_z^{1.34}$ for bf16, not $\sigma_z^2$.
\ver{$\chi^2/\Var_\pi(e)=1.0000,\ 1.0005,\ 1.0021$ at $\sigma_z=2,10,20$; $\Var_\pi(e)/(u^2\E_\pi z^2)\le0.59,\ 0.71,\ 0.64$ (never above 1).}

\subsection{Long sequences: exact scaling, bf16 vs fp16 vs fp32}
For positions with independent contexts, $\E_\mu[\log w_T]=-\sum_t\KL(\mu_t\|\pi_t)$, $\Var(\log w_T)=\sum_t\Var_t$ and $\E_\mu[w_T^2]=\prod_t(1+\chi_t^2)$, so $\ESS/n\approx\prod_t(1+\chi_t^2)^{-1}$ at length $T$ (0103, Eq.~(4)). With vocabulary $64$ and logits $\sim\mathcal N(0,\sigma_z^2)$, rounded to format $F$:
\begin{center}\small\begin{tabular}{@{}cccccc@{}}\toprule
$\sigma_z$ & format & $\chi^2$ / token & $\ESS/n$, $T{=}100$ & $T{=}1000$ & $T{=}4000$\\\midrule
10 & bf16 & $2.9\times10^{-4}$ & 0.977 & 0.760 & 0.309\\
10 & fp16 & $4.3\times10^{-6}$ & 0.9996 & 0.996 & 0.983\\
10 & fp32 & $6.5\times10^{-14}$ & 1 & 1 & 1\\
20 & bf16 & $6.2\times10^{-4}$ & 0.958 & 0.560 & 0.083\\
20 & fp32 & $1.3\times10^{-13}$ & 1 & 1 & 1\\\bottomrule\end{tabular}\end{center}
The fp32/bf16 ratio of $\chi^2$ is $2.2\times10^{-10}$ vs.\ the theoretical $(u_{32}/u_{16})^2=2.3\times10^{-10}$. So raising the precision of the logits does not make a small improvement, it removes the effect (an $\ESS$ of $31\%$ at $4000$ tokens becomes $1$), \emph{provided the logits are the dominant error source}.
\ver{MC with $2\times10^4$ sampled $64$-token sequences: $\E_\mu\log w_{64}$ within $0.41$ s.e.\ and $\Var\log w_{64}$ within $0.54$ s.e.\ of the exact sums; $\E_\mu w=1$ per position to $6.7\times10^{-16}$; measured max relative rounding errors $3.89\times10^{-3}$, $4.88\times10^{-4}$, $5.96\times10^{-8}$ against $u=3.91\times10^{-3},\,4.88\times10^{-4},\,5.96\times10^{-8}$.}

\subsection{Why the \emph{LM head} is the place to spend precision}
Logits are $z_b=\sum_ih_iW_{bi}$ with $h$ the final hidden state. Rounding $h$ and $W$ to bf16 perturbs each product $h_iW_{bi}$ by a relative error $\lesssim2u$ (two rounded factors), so the logit error is proportional to $\sum_i|h_iW_{bi}|$, the magnitude of the \emph{terms}, which is at least $|z_b|$ (cancellation can only make it larger), and the logits feed an exponential. MiniMax localised the discrepancy layer by layer to ``high-magnitude activations in the LM head at the output layer''. In an emulation ($d=1024$, bf16-rounded inputs, exact accumulation) the logit error is $0.59\,u\cdot\mathrm{std}(z)$, \emph{independent of whether outlier hidden dimensions are present}, and grows in absolute terms with them ($\mathrm{std}(z)=1.0\to8.8$, error $2.4\times10^{-3}\to2.0\times10^{-2}$ as 8 hidden dims are scaled $100\times$); with fp32 inputs the error is $\approx6.2\times10^4$ times smaller, matching $u_{\rm bf16}/u_{\rm fp32}=65536$.
\ver{error$/\mathrm{std}(z)=2.3\times10^{-3}$ ($0.59u$) at outlier scale $1,10,100$; bf16/fp32 error ratio $6.6\times10^4,\,6.4\times10^4,\,6.2\times10^4$.}
\paragraph{Limits of this explanation.} Our model isolates the rounding of the logits. A real network also rounds every earlier layer, and kernels differ (fused attention, MoE routing, lightning attention). Qi et al.\ (arXiv:2510.26788) locate the root cause in BF16 rounding \emph{throughout} and fix it with FP16 ($u=2^{-11}$); Zhang et al.\ (arXiv:2602.01826) argue that mismatch grows with gradient noise during long runs so that IS alone can fail, and suppress it with an LR schedule. FP16 has a narrow range (max $65504$), and this project's own notes record an fp16 NaN on FunctionGemma (262K-vocabulary logits), where bf16 is required: the precision choice is model-dependent. On the toy, a bf16 head gives a probability correlation $0.99979$ and no token outside $|\pi/\mu-1|>0.2$; the larger discrepancy MiniMax measured ($0.987$) must therefore come from errors that accumulate through the network, which the FP32 head is reported to remove in their model, not something we can reproduce here.

\subsection{Interaction with PPO/CISPO ratios}
Three policies appear: $\mu$ (generated the samples), $\pi_{\rm old}$ (trainer at $\theta_{\rm old}$), $\pi_\theta$. If $\pi_{\rm old}$ is taken from the sampler the ratio is $\pi_\theta/\mu$, equal to $\pi/\mu\ne1$ at $\theta=\theta_{\rm old}$, and tokens are clipped (0104) or reweighted (0107) before any update. Splitting
\[\frac{\pi_\theta}{\mu}=\underbrace{\frac{\pi_{\rm old}}{\mu}}_{\rho,\ \text{mismatch}}\cdot\underbrace{\frac{\pi_\theta}{\pi_{\rm old}}}_{r,\ \text{policy lag}}\]
lets one clip $r$ as usual and apply a truncated $\min(\rho,C)$ for the mismatch (the structure of the TIS patch on GRPO described in arXiv:2510.26788 Sec.~1). The extra forward pass for $\pi_{\rm old}$ is estimated there at $\approx25\%$ of the training cost, and they cite Liu et al.\ for the observation that token-level TIS can remain insufficient because its gradient is biased; sequence-level correction is unbiased but high-variance.

\section{Boxed result}
\begin{result}
\begin{align*}
&\E_\mu[A\,s_\pi]-\nabla J=\E_\pi\big[(\tfrac\mu\pi-1)As_\pi\big],\qquad\E_\mu[wAs_\pi]=\nabla J\\
&\text{TIS bias}=-\E_\mu[(w-C)_+As],\quad\|\cdot\|\le M_AM_s\min\big(\TV(\pi,\mu),P_\pi(w>C)\big)\\
&\TV(\pi,\mu)\le\textstyle\sum_t\E_\pi\TV_t,\qquad|e_a|\le u|z_a|,\qquad\chi^2(\pi\|\mu)\approx\Var_\pi(e)\le u^2\E_\pi[z^2]\\
&\E_\mu[w_T^2]=\textstyle\prod_t(1+\chi_t^2)\quad(\text{history-independent }\chi_t^2;\ \text{else the recursion of 0103})\\
&\chi^2_{\rm fp32}/\chi^2_{\rm bf16}\approx(2^{-16})^2
\end{align*}
\end{result}

\section{What breaks}
\begin{breaks}\small\begin{itemize}
\item \textbf{Small-error assumption.} Eqs.~\eqref{eq:lpvar}--\eqref{eq:chi2e} are first order. Discrete failures are not small perturbations: a \emph{router flip} in a mixture-of-experts layer changes which experts run, an $O(1)$ jump in $\mu$ (the router-replay line of work, arXiv:2510.11370, targets this; vault summary only, not read in full here).
\item \textbf{Support violations.} Top-$k$/top-$p$ sampling gives $\mu=0$ where $\pi>0$: IS cannot correct it (0103), and $\TV(\pi,\mu)$ is then not small.
\item \textbf{TIS/MIS are biased} (Eq.~\eqref{eq:tisbound}) and the bias is \emph{systematic}: they remove exactly the tokens/sequences with the largest $\pi/\mu$, which are the tokens where the sampler under-weights the trainer's belief.
\item \textbf{Deployment gap remains.} Correcting $\nabla J_\pi$ does not make $\arg\max_\theta J_\pi=\arg\max_\theta J_\mu$ (Eq.~4 of arXiv:2510.26788); only removing the mismatch at the source does.
\item \textbf{FP32 head is one error source.} It removes the logit rounding (here $\chi^2\to10^{-10}$ of bf16) but not hidden-state errors; with all other layers in bf16 the residual mismatch can remain large.
\item \textbf{Diagnostics.} Pearson correlation of probabilities is dominated by high-probability tokens; monitor mean $|\Delta\log p|$, $\max|\Delta\log p|$, per-token $\chi^2$ and the fraction of tokens with $|\pi/\mu-1|>\epsilon$.
\item \textbf{Precision trade-off.} fp16 gains mantissa bits but loses range ($>65504$ overflows; project note: fp16 NaN on a 262K-vocabulary model).
\end{itemize}\end{breaks}

\nextnote{Everything so far regularised the \emph{speed} of the policy (clips, caps) rather than its \emph{distance to the reference}. Note 0109 treats the KL penalty: the three Monte-Carlo estimators $k_1,k_2,k_3$, their bias and variance as \emph{values}, what their \emph{gradients} actually optimise, and why recent prover/reasoning recipes set $\beta=0$.}

\section*{References}
{\small\begin{itemize}[leftmargin=1.2em]
\item Qi et al., \emph{Defeating the Training-Inference Mismatch via FP16}, arXiv:2510.26788: Sec.~1 (two-engine mismatch, deployment gap, $\approx25\%$ overhead), Sec.~2 Eqs.~3--4 (bias, deployment gap), Sec.~2.1 Eqs.~5--7 (IS, TIS, MIS).
\item MiniMax, arXiv:2506.13585, Sec.~3.2 ``Computational Precision Mismatch in Generation and Training'' and Fig.~3 (FP32 LM head; correlations $0.987319\to0.997135$ as printed in the figure).
\item Zhang et al., \emph{Beyond Precision}, arXiv:2602.01826 (cited at abstract level: mismatch as an optimisation problem, LR scheduling; the full text is now in the vault as \texttt{beyond-precision-training-inference-mismatch-is}, but no claim above depends on it).
\item Stabilizing MoE RL by aligning training and inference routers, arXiv:2510.11370 (cited from its abstract; the full text is now in the vault as \texttt{stabilizing-moe-reinforcement-learning-by-aligning}, but no claim above depends on it).
\item Mistral AI, Leanstral technical report (\texttt{LeanstralReport.pdf}, Sec.~2.3; not on arXiv): the CISPO ratio there is against $\pi_{\rm sample}$ and is stated to target trainer--generator mismatch; see 0107, subsection \emph{Leanstral variant}.
\item Rounding model, Eqs.~\eqref{eq:biasmu}--\eqref{eq:chi2e}, the hybrid TV bound and the long-horizon product are derived here and verified in the notebook; unit roundoffs $2^{-8},2^{-11},2^{-24}$ are checked numerically against \texttt{torch} casts.
\end{itemize}}
\end{document}