# %% [markdown]
# # 0108 - Training-inference mismatch: sampler vs trainer log-probs, truncated IS, why an FP32 LM head helps
# Companion to `notes/0108 - Training-Inference Mismatch.pdf`.  fp64 reference ("trainer" pi); the "sampler" mu uses logits rounded to bf16/fp16/fp32.
# %%
import sys, math
sys.path.insert(0, "..")
from _m1_lib import *
import numpy as np
# %% [markdown]
# ## 0. Floating-point facts: unit roundoff u = 2^-(p) with p significant bits;  |round(x) - x| <= u |x|
# bf16: 8 significant bits (u=2^-8), fp16: 11 (u=2^-11), fp32: 24 (u=2^-24).
# %%
gen = torch.Generator().manual_seed(0)
x = torch.exp(torch.rand(2_000_000, generator=gen) * math.log(1e3 / 1e-3) + math.log(1e-3))    # log-uniform in [1e-3, 1e3]
for name, dt, u in (("bf16", torch.bfloat16, 2.0 ** -8), ("fp16", torch.float16, 2.0 ** -11), ("fp32", torch.float32, 2.0 ** -24)):
    rel = ((x.to(dt).to(torch.float64) - x).abs() / x)
    print(f"   {name}: max relative rounding error {float(rel.max()):.3e}   u = {u:.3e}")
    check(f"{name}: max rel. rounding error <= u (=2^-{int(-math.log2(u))})", max(0.0, float(rel.max()) - u), 0.0)
    assert float(rel.max()) > 0.5 * u
# %% [markdown]
# ## 1. The tiny LM with LARGE logits (like an LM head with high-magnitude activations); sampler mu = same network with bf16-rounded logits
# %%
lm = TinyLM(V=4, L=4, seed=0, scale=4.0)
theta = lm.theta0.clone()
R = default_reward(lm)
def seq_logprob_rounded(theta, dtype):
    V, L = lm.V, lm.L
    tok, pos = lm._logits_tables(theta)
    prev = torch.cat([torch.full((lm.N, 1), V, dtype=torch.long), lm.tokens[:, :-1]], 1)
    logits = tok[prev] + pos[None, :, :]
    if dtype is not None: logits = logits.to(dtype).to(torch.float64)
    lp = torch.log_softmax(logits, -1).gather(-1, lm.tokens[..., None]).squeeze(-1) * lm.mask
    return lp.sum(1)
pi = seq_logprob_rounded(theta, None).exp()
mu = seq_logprob_rounded(theta, torch.bfloat16).exp()
check("trainer pi is normalised", abs(float(pi.sum()) - 1), 1e-14)
check("sampler mu is normalised (bf16 logits then softmax)", abs(float(mu.sum()) - 1), 1e-14)
b = float(pi @ R); A = R - b
S = score_matrix(lm, theta)                                  # score of the TRAINER policy
g_true = (pi[:, None] * A[:, None] * S).sum(0)               # grad J_pi
w = pi / mu
TVs = 0.5 * float((pi - mu).abs().sum())
print(f"   bf16 logits: TV(pi, mu) = {TVs:.5f}   chi2(pi||mu) = {float((mu*(w-1)**2).sum()):.3e}   max w = {float(w.max()):.4f}")
for nm, dt in (("fp16", torch.float16), ("fp32", torch.float32)):
    m_ = seq_logprob_rounded(theta, dt).exp()
    print(f"   {nm} logits: TV(pi, mu) = {0.5*float((pi-m_).abs().sum()):.3e}   chi2 = {float((m_*(pi/m_-1)**2).sum()):.3e}")
# %% [markdown]
# ## 2. Ignoring the mismatch is biased; IS repairs it exactly  (arXiv:2510.26788 Eqs. 3, 5)
# %%
g_biased = (mu[:, None] * A[:, None] * S).sum(0)             # samples from mu, trainer's score, no correction
bias_formula = ((mu - pi)[:, None] * A[:, None] * S).sum(0)  # = E_pi[(mu/pi - 1) A s]
check("bias of the uncorrected estimator == sum_o (mu-pi) A s", maxabs(g_biased - g_true, bias_formula), 1e-14)
g_is = (mu[:, None] * w[:, None] * A[:, None] * S).sum(0)
check("sequence-level IS correction E_mu[(pi/mu) A s] == grad J_pi exactly", maxabs(g_is, g_true), 1e-14)
print(f"   ||bias|| / ||grad J|| = {float((g_biased-g_true).norm()/g_true.norm()):.4f}   (uncorrected)")
assert float((g_biased - g_true).norm()) > 1e-4
# %% [markdown]
# ## 3. Truncated IS (TIS, Eq. 6) and masked IS (MIS, Eq. 7): exact bias and bounds
#   TIS: E_mu[min(w,C) A s];   bias = - E_mu[(w-C)_+ A s]       ||.|| <= M_A M_s sum_o (pi - C mu)_+  <= M_A M_s min(TV, P_pi(w>C))
#   MIS: E_mu[1(w<=C) w A s];  bias = - E_pi[1(w>C) A s]         ||.|| <= M_A M_s P_pi(w>C)
# %%
MA, Ms = float(A.abs().max()), float(S.norm(dim=1).max())
print("   C       ||TIS bias||   bound1(sum mu (w-C)+ |A| ||s||)   bound2(MA Ms sum(pi-C mu)+)   MA Ms TV    ||MIS bias||   MA Ms P_pi(w>C)")
for C in (1.0, 1.02, 1.05, 1.2, 2.0):
    wt = torch.clamp(w, max=C)
    g_tis = (mu[:, None] * wt[:, None] * A[:, None] * S).sum(0)
    bias_tis = g_tis - g_true
    formula = -((mu * (w - C).clamp(min=0))[:, None] * A[:, None] * S).sum(0)
    check(f"TIS bias == -E_mu[(w-C)_+ A s]  (C={C})", maxabs(bias_tis, formula), 1e-14)
    b1 = float((mu * (w - C).clamp(min=0) * A.abs() * S.norm(dim=1)).sum())
    ex = float((pi - C * mu).clamp(min=0).sum())
    b2 = MA * Ms * ex
    assert float(bias_tis.norm()) <= b1 + 1e-15 and b1 <= b2 + 1e-15 and ex <= TVs + 1e-15
    g_mis = (mu[:, None] * (w * (w <= C))[:, None] * A[:, None] * S).sum(0)
    bias_mis = g_mis - g_true
    mis_formula = -((pi * (w > C))[:, None] * A[:, None] * S).sum(0)
    check(f"MIS bias == -E_pi[1(w>C) A s]  (C={C})", maxabs(bias_mis, mis_formula), 1e-14)
    Pw = float((pi * (w > C)).sum())
    assert float(bias_mis.norm()) <= MA * Ms * Pw + 1e-15
    print(f"   {C:4.2f}   {float(bias_tis.norm()):12.4e}   {b1:12.4e}                  {b2:12.4e}              {MA*Ms*TVs:10.4e}  {float(bias_mis.norm()):12.4e}   {MA*Ms*Pw:12.4e}")
check("TIS: 0 <= bias <= bound1 <= bound2 <= M_A M_s TV(pi,mu) for all C", 0.0, 0.0)
check("TIS at C=inf is exact (no truncation) -> bias = 0", float((((mu[:, None] * torch.clamp(w, max=1e12)[:, None] * A[:, None] * S).sum(0)) - g_true).norm()), 1e-13)
# TV of sequence distributions <= sum over positions of E[TV_t]   (hybrid argument)
def tv_chain(theta, dtype):
    tot = 0.0
    V, L = lm.V, lm.L
    tok, pos = lm._logits_tables(theta)
    def cond(prefix, rounder):
        prev = V if len(prefix) == 0 else prefix[-1]
        z = tok[prev] + pos[len(prefix)]
        if rounder is not None: z = z.to(rounder).to(torch.float64)
        return torch.softmax(z, 0)
    def rec(prefix, prob):
        nonlocal tot
        if len(prefix) >= L: return
        cp, cm = cond(prefix, None), cond(prefix, dtype)
        tot += prob * 0.5 * float((cp - cm).abs().sum())
        for a in range(V - 1): rec(prefix + (a,), prob * float(cp[a]))
    rec((), 1.0)
    return tot
tvc = tv_chain(theta, torch.bfloat16)
print(f"   TV(sequence) = {TVs:.4f} <= sum_t E_pi[TV_t] = {tvc:.4f}")
check("sequence TV <= sum_t E[TV_t]  (max violation, must be <= 0)", max(0.0, TVs - tvc), 0.0)
# %% [markdown]
# ## 4. From logit rounding error to log-prob error:  d log pi_a = e_a - sum_b pi_b e_b  (first order);  Var = s^2 (1 - 2 pi_a + sum pi_b^2) for iid logit errors
# %%
gen = torch.Generator().manual_seed(3)
z = torch.randn(7, generator=gen) * 5.0
pi_z = torch.softmax(z, 0)
eps = torch.randn(7, generator=gen) * 1e-6
dlp_exact = torch.log_softmax(z + eps, 0) - torch.log_softmax(z, 0)
dlp_lin = eps - (pi_z * eps).sum()
check("first-order log-prob error  e_a - sum_b pi_b e_b  (relative error of linearisation, e~1e-6)", float((dlp_exact - dlp_lin).abs().max() / dlp_lin.abs().max()), 1e-5)
s_ = 0.01
E = torch.randn(400000, 7, generator=gen) * s_
dlp = torch.log_softmax(z + E, -1) - torch.log_softmax(z, 0)
var_mc = dlp.var(0); var_f = s_ ** 2 * (1 - 2 * pi_z + (pi_z ** 2).sum())
zs = float(((var_mc - var_f).abs() / (var_f * math.sqrt(2 / 400000))).max())
check("Var(d log pi_a) = s^2 (1 - 2 pi_a + sum pi^2): max z-score over 7 tokens (MC, 400k draws)", zs, 4.0, kind="mc")
# %% [markdown]
# ## 4b. Where the logit error comes from: an LM-head matmul with bf16 inputs vs an fp32 head
# z_b = sum_i h_i W_bi.  Rounding h and W to bf16 perturbs each product by a relative error <= u, so the logit error scales with the
# magnitude of the terms (not just of z).  Outlier activations (a few huge hidden dims) enlarge that magnitude.
# %%
gen = torch.Generator().manual_seed(4)
d, Vv, Tn = 1024, 64, 512
def head_errors(outlier):
    h = torch.randn(Tn, d, generator=gen)
    h[:, :8] *= outlier                                         # 8 high-magnitude hidden dimensions
    W = torch.randn(Vv, d, generator=gen) / math.sqrt(d)
    z = h @ W.T                                                 # fp64 reference
    zb = h.to(torch.bfloat16).to(torch.float64) @ W.to(torch.bfloat16).to(torch.float64).T     # bf16-rounded inputs, exact accumulation
    z32 = h.to(torch.float32).to(torch.float64) @ W.to(torch.float32).to(torch.float64).T      # fp32-rounded inputs
    return float(z.std()), float((zb - z).std()), float((z32 - z).std())
print("   outlier x    std(z)    std(err bf16)   std(err fp32)    err_bf16/std(z)   /u_bf16     ratio bf16/fp32")
for out in (1.0, 10.0, 100.0):
    sz_, eb, e32 = head_errors(out)
    print(f"   {out:7.0f}   {sz_:8.2f}   {eb:12.3e}   {e32:12.3e}      {eb/sz_:10.3e}    {eb/sz_/2**-8:6.3f}     {eb/e32:9.1f}")
sz_, eb, e32 = head_errors(100.0)
check("bf16 head: logit error / std(z) between 0.05u and 1u  (u=2^-8)", 0.0 if 0.05 * 2**-8 < eb / sz_ < 2**-8 else 1.0, 0.0)
check("fp32 head cuts the logit error by ~ u_bf16/u_fp32 = 65536 (log10 ratio within 1 decade)", abs(math.log10(eb / e32) - math.log10(2**16)), 1.0)
# %% [markdown]
# ## 5. Long horizons: positions are independent contexts, so everything is EXACT:
#   E_mu[log w_T] = -sum_t KL(mu_t||pi_t),  Var(log w_T) = sum_t Var_t,  E_mu[w_T^2] = prod_t (1 + chi2_t).   (V=64, logits ~ N(0, sigma_z^2))
# %%
def per_position_stats(sigma_z, dtype, T=4096, V=64, seed=0):
    g = torch.Generator().manual_seed(seed)
    z = torch.randn(T, V, generator=g) * sigma_z
    pi = torch.softmax(z, -1)
    zr = z if dtype is None else z.to(dtype).to(torch.float64)
    mu = torch.softmax(zr, -1)
    lw = pi.log() - mu.log()                                  # log(pi/mu) per token option
    Elw = (mu * lw).sum(-1)                                   # E_mu[log w_t] = -KL(mu||pi)
    Vlw = (mu * lw ** 2).sum(-1) - Elw ** 2
    chi2 = (mu * (pi / mu - 1) ** 2).sum(-1)
    return Elw, Vlw, chi2, pi, mu
print("   sigma_z  head      mean KL/token    mean chi2/token     T=100: ESS frac     T=1000      T=4000   (ESS frac = 1/prod(1+chi2))")
rows = {}
for sz in (2.0, 5.0, 10.0, 20.0):
    for name, dt in (("bf16", torch.bfloat16), ("fp16", torch.float16), ("fp32", torch.float32)):
        Elw, Vlw, chi2, pi_t, mu_t = per_position_stats(sz, dt)
        lESS = lambda T: float(-torch.log1p(chi2[:T]).sum())
        rows[(sz, name)] = (float((-Elw).mean()), float(chi2.mean()), lESS(100), lESS(1000), lESS(4000))
        print(f"   {sz:5.1f}   {name}   {float((-Elw).mean()):14.3e}   {float(chi2.mean()):14.3e}      {math.exp(lESS(100)):10.3e}   {math.exp(lESS(1000)):10.3e}  {math.exp(lESS(4000)):10.3e}")
# why chi2 does not grow like sigma_z^2: to first order chi2_t = Var_{a~pi}(e_a) where e = round(z) - z (the *centred* logit error)
def first_order_check(sz, dtype, u, T=2048, V=64):
    g_ = torch.Generator().manual_seed(0)
    z_ = torch.randn(T, V, generator=g_) * sz
    e_ = z_.to(dtype).to(torch.float64) - z_
    pi_ = torch.softmax(z_, -1); mu_ = torch.softmax(z_ + e_, -1)
    chi2_ = (mu_ * (pi_ / mu_ - 1) ** 2).sum(-1)
    var_e = (pi_ * e_ ** 2).sum(-1) - ((pi_ * e_).sum(-1)) ** 2
    bound = u ** 2 * (pi_ * z_ ** 2).sum(-1)
    return float(chi2_.sum() / var_e.sum()), float((var_e / bound).max()), float(chi2_.mean())
for sz in (2.0, 10.0, 20.0):
    ratio, bnd, c2 = first_order_check(sz, torch.bfloat16, 2.0 ** -8)
    print(f"   sigma_z={sz:5.1f}: chi2/Var_pi(e) = {ratio:.4f}   max Var_pi(e)/(u^2 E_pi z^2) = {bnd:.3f}")
    check(f"first-order law chi2_t = Var_pi(logit rounding error)  (sigma_z={sz}, mean ratio - 1)", abs(ratio - 1), 0.1)
    check(f"bound Var_pi(e) <= u^2 E_pi[z^2]  (sigma_z={sz}, max ratio - 1 must be <= 0)", max(0.0, bnd - 1), 0.0)
szs = np.array([2.0, 5.0, 10.0, 20.0]); chis = np.array([rows[(s, 'bf16')][1] for s in szs])
slope = np.polyfit(np.log(szs), np.log(chis), 1)[0]
print(f"   bf16 log-log slope of chi2/token vs logit scale: {slope:.3f}  (between 1 and 2: the error grows like |z| but the centred error Var_pi(e) shrinks as pi becomes peaked)")
assert 1.0 < slope < 2.0
ratio_fp32 = rows[(10.0, 'fp32')][1] / rows[(10.0, 'bf16')][1]
print(f"   fp32/bf16 chi2 ratio at sigma_z=10: {ratio_fp32:.2e}   (theory (u32/u16)^2 = {(2.0**-24/2.0**-8)**2:.2e})")
check("fp32 head reduces per-token chi^2 by ~ (u_fp32/u_bf16)^2  (|log10 ratio - log10 theory|)", abs(math.log10(ratio_fp32) - math.log10((2.0 ** -24 / 2.0 ** -8) ** 2)), 1.0)
# MC check of one exact statistic: Var(log w_T) over sampled sequences for T=64
Elw, Vlw, chi2, pi_t, mu_t = per_position_stats(10.0, torch.bfloat16, T=64)
g = torch.Generator().manual_seed(9); n = 20000
tokens = torch.stack([torch.multinomial(mu_t[t], n, replacement=True, generator=g) for t in range(64)], 1)       # [n, T]
lw_seq = (pi_t.log().gather(1, tokens.T).T - mu_t.log().gather(1, tokens.T).T).sum(1)
exact_mean, exact_var = float(Elw.sum()), float(Vlw.sum())
se_mean = math.sqrt(exact_var / n)
check("MC: E_mu[log w_64] == -sum KL  (z-score)", abs(float(lw_seq.mean()) - exact_mean) / se_mean, 4.0, kind="mc")
check("MC: Var(log w_64) == sum_t Var_t  (z-score; kurtosis ignored -> loose 6)", abs(float(lw_seq.var()) - exact_var) / (exact_var * math.sqrt(2 / n)), 6.0, kind="mc")
check("E_mu[w_T] = 1 exactly per position (sum_a pi_a = 1)", float((mu_t * (pi_t / mu_t)).sum(-1).sub(1).abs().max()), 1e-13)
# %% [markdown]
# ## 5b. Caveat: E_mu[w_T^2] = prod_t (1 + chi2_t) needs history-INDEPENDENT chi2_t.  With state-dependent chi2 the exact statements are
#   (i) the recursion  E[w_T^2] = E[w_{T-1}^2 (1 + chi2_T(s_T))],   (ii) the martingale identity  E[ w_T^2 / prod_t (1 + chi2_t(s_t)) ] = 1,
# and E[prod_t (1 + chi2_t(s_t))] is NOT E[w_T^2].  Exact enumeration on a tiny autoregressive model (sampler mu = one weight vector, trainer pi = a perturbed one).
# %%
lmx = TinyLM(V=3, L=8, seed=4, scale=1.0)
tokx = lmx.theta0[:lmx.n_tok].reshape(lmx.V + 1, lmx.V); tokx[:, lmx.EOS] -= 2.0
th_mu = lmx.theta0.clone(); th_pi = perturb(th_mu, 0.8, seed=9)
lmu, lpi = lmx.token_logprobs(th_mu), lmx.token_logprobs(th_pi)
mu_x = lmu.sum(1).exp()
chi_x = torch.zeros(lmx.N, lmx.L)                                   # chi2_t(s_t) = chi2(pi(.|s_t) || mu(.|s_t)) at the realised prefix
for i, sq in enumerate(lmx.seqs):
    for t in range(len(sq)):
        cm, cp = lmx.cond_probs(th_mu, sq[:t]), lmx.cond_probs(th_pi, sq[:t])
        chi_x[i, t] = ((cp - cm) ** 2 / cm).sum()
w2_x = (lpi - lmu).sum(1).exp() ** 2
prod_x = (1 + chi_x).prod(1)                                        # masked positions have chi2 = 0 -> factor 1
Ew2 = float((mu_x * w2_x).sum()); Eprod = float((mu_x * prod_x).sum())
w2_prev = torch.cat([torch.zeros(lmx.N, 1), (lpi - lmu).cumsum(1)[:, :-1]], 1).exp() ** 2        # w_{t-1}^2 (w_0 = 1)
rec = float((mu_x * w2_prev[:, -1] * (1 + chi_x[:, -1])).sum())
print(f"   N={lmx.N}:  E_mu[w^2] = {Ew2:.6f};  recursion E[w_(T-1)^2 (1+chi2_T)] = {rec:.6f};  E_mu[prod_t(1+chi2_t)] = {Eprod:.6f};  E_mu[w^2/prod] = {float((mu_x*w2_x/prod_x).sum()):.6f}")
check("recursion E[w_T^2] = E[w_(T-1)^2 (1+chi2_T(s_T))]  (state-dependent chi2)", abs(Ew2 - rec), 1e-14)
check("martingale identity E[w_T^2 / prod_t (1+chi2_t(s_t))] = 1", abs(float((mu_x * w2_x / prod_x).sum()) - 1.0), 1e-13)
assert abs(Ew2 - Eprod) > 1e-3          # the naive product formula is genuinely wrong when chi2 depends on the history
check("E[prod(1+chi2_t)] != E[w_T^2] on a state-dependent model (relative gap, must exceed 1e-3)", max(0.0, 1e-3 - abs(Ew2 - Eprod) / Ew2), 0.0)
# %% [markdown]
# ## 6. Consequence for PPO/CISPO ratios and for a "probability correlation" diagnostic
# If the *sampler's* log-probs are used as pi_old, then at theta = theta_old the ratio is r = pi/mu != 1: tokens are clipped before any update.
# %%
for name, dt in (("bf16", torch.bfloat16), ("fp32", torch.float32)):
    Elw, Vlw, chi2, pi_t, mu_t = per_position_stats(10.0, dt, T=2048)
    gen = torch.Generator().manual_seed(1)
    tok_ = torch.stack([torch.multinomial(mu_t[t], 64, replacement=True, generator=gen) for t in range(2048)], 0)   # [T, 64]
    r = (pi_t.gather(1, tok_) / mu_t.gather(1, tok_))
    out = float(((r - 1).abs() > 0.2).double().mean())
    p_tr, p_inf = pi_t.gather(1, tok_).reshape(-1), mu_t.gather(1, tok_).reshape(-1)
    corr = float(torch.corrcoef(torch.stack([p_tr, p_inf]))[0, 1])
    print(f"   {name} head: fraction of sampled tokens with |pi/mu - 1| > 0.2 at theta=theta_old: {out:.4f}   Pearson corr(p_train, p_sampler) = {corr:.6f}")
# %% [markdown]
# ## 7. Naive vs stable weight computation under mismatch (log-space)
# %%
Elw, Vlw, chi2, pi_t, mu_t = per_position_stats(10.0, torch.bfloat16, T=4096)
tokens = torch.stack([torch.multinomial(mu_t[t], 1, replacement=True, generator=g) for t in range(4096)], 0).squeeze(1)
lr_ = (pi_t.log().gather(1, tokens[:, None]) - mu_t.log().gather(1, tokens[:, None])).squeeze(1)
logw = float(lr_.sum())
naive32 = torch.tensor(1.0, dtype=torch.float32)
for v in lr_.exp().to(torch.float32): naive32 = naive32 * v
print(f"   T=4096 bf16-mismatch sequence: log w = {logw:.2f};  naive fp32 running product = {float(naive32)};  exp(sum log) in fp64 = {math.exp(logw):.3e}")
check("naive fp32 product vs log-space: agree only when representable (log w within fp32 range)", 0.0 if (abs(logw) > 80 or abs(float(naive32) - math.exp(logw)) < 1e-4 * math.exp(logw)) else 1.0, 0.0)
# %% [markdown]
# ## 8. Naive vs efficient TIS-corrected gradient estimate (sampled from the bf16 sampler mu)
# naive: one backward per sample;  efficient: one backward through the surrogate with DETACHED truncated weights.
# %%
lm = TinyLM(V=4, L=4, seed=0, scale=4.0)
theta = lm.theta0.clone(); R = default_reward(lm)
pi_ = seq_logprob_rounded(theta, None).exp(); mu_ = seq_logprob_rounded(theta, torch.bfloat16).exp()
bb = float(pi_ @ R)
gen = torch.Generator().manual_seed(6)
idx = torch.multinomial(mu_, 500, replacement=True, generator=gen)
Cc = 1.02
w_ = (pi_ / mu_)[idx]
acc = torch.zeros_like(theta)
for j, i in enumerate(idx.tolist()):
    t = theta.clone().requires_grad_(True)
    acc += float(min(w_[j], Cc)) * float(R[i] - bb) * torch.autograd.grad(seq_logprob_rounded(t, None)[i], t)[0]
g_naive = acc / len(idx)
t = theta.clone().requires_grad_(True)
sur = (torch.clamp(w_, max=Cc).detach() * (R[idx] - bb) * seq_logprob_rounded(t, None)[idx]).mean()
g_eff = torch.autograd.grad(sur, t)[0]
check("TIS estimate: per-sample loop == single backward through detached-weight surrogate", maxabs(g_naive, g_eff), 1e-13)
print("done; checks:", len(CHECKS))