"""Block-wise 8-bit AdamW in pure torch: codebook, quantization, the optimizer,
resume, the factory, and the 2-rank gloo/FSDP2 behaviour.

Run: .venv/bin/python -m pytest tests/test_adamw8bit.py -q -m "not gpu and not oracle" -p no:cacheprovider

Why this optimizer exists (measured, do not re-derive): every shipped full-FT
config names `adamw8bit`; fp32 AdamW does not fit a 23GB A10G and
`OffloadAdamW` overshoots its pinned-staging cap 4x at 7.25B params / 4 ranks
and projects ~105GB host RSS against ~85GB available. See adamw8bit.py.

House conventions: each regression is NAMED AFTER THE INCIDENT/FAILURE CLASS it
prevents, and each assertion is one a mutant that reintroduces the bug fails
(the mutation each test kills is stated in its docstring).
"""

from __future__ import annotations

import copy
import json
import math
import os
import socket
import sys
import traceback

import pytest
import torch
import torch.multiprocessing as mp

from align.train import quant8
from align.train.adamw8bit import (
    BlockwiseAdamW8bit,
    adamw_update_,
    stochastic_round_to_bf16,
)

HP = dict(lr=1e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1)
B = 2048


def _randn(*shape, seed=0, dtype=torch.float32):
    g = torch.Generator().manual_seed(seed)
    return torch.randn(*shape, generator=g, dtype=torch.float32).to(dtype)


def _param(*shape, seed=0, dtype=torch.float32):
    return torch.nn.Parameter(_randn(*shape, seed=seed, dtype=dtype))


def _set_grads(params, step, seed=7, scale=1.0, dtype=None):
    g = torch.Generator().manual_seed(seed + step)
    for p in params:
        gr = torch.randn(p.shape, generator=g, dtype=torch.float32) * scale
        p.grad = gr.to(dtype or p.dtype)


# ---------------------------------------------------------------------------
# the codebook
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("signed", [True, False])
def test_codebook_is_a_sorted_256_entry_map_covering_the_unit_range(signed):
    """Kills: a miscounted decade (255/257 entries -> uint8 overflow or a dead
    code), a duplicate/unsorted entry (searchsorted requires strictly
    increasing input and returns garbage otherwise), a map that does not reach
    +-1 (the block's absmax element would clip), a missing exact zero (the
    all-zero state could not be represented)."""
    cb = quant8.create_dynamic_map(signed=signed)
    assert cb.shape == (256,) and cb.dtype == torch.float32
    assert bool(torch.isfinite(cb).all())
    assert bool((cb[1:] > cb[:-1]).all()), "codebook must be STRICTLY increasing"
    assert 0.0 in cb.tolist()
    assert cb[quant8.zero_code(signed)].item() == 0.0
    assert cb.max().item() == 1.0
    if signed:
        assert cb.min().item() <= -0.99, "signed map must reach (nearly) -1"
        assert (cb < 0).sum().item() == (cb > 0).sum().item() - 1, (
            "signed map is symmetric except for the extra exact +1.0"
        )
    else:
        assert cb.min().item() == 0.0 and bool((cb >= 0).all())


def test_codebook_is_dynamic_not_linear():
    """The point of the map: relative resolution across decades. A linear grid
    has 1/255 absolute spacing, so values below ~4e-3 of absmax collapse to
    zero; the dynamic map must still resolve 1e-5 of absmax to within a few
    percent. Kills: replacing create_dynamic_map with linspace(-1, 1, 256)."""
    for signed in (True, False):
        cb = quant8.create_dynamic_map(signed)
        codes = quant8.nearest_code(torch.tensor([1e-5, 1e-3, 0.5]), cb).long()
        got = cb[codes]
        rel = ((got - torch.tensor([1e-5, 1e-3, 0.5])).abs() / torch.tensor([1e-5, 1e-3, 0.5]))
        assert bool((rel < 0.45).all()), (signed, got, rel)
        assert rel[1] < 0.1 and rel[2] < 0.02
    # at least 6 decades of dynamic range above the smallest nonzero code
    assert quant8.create_dynamic_map(False)[1].item() < 1e-6


@pytest.mark.parametrize("signed", [True, False])
def test_nearest_code_is_truly_the_nearest_including_edges_and_midpoints(signed):
    """Kills: an off-by-one in the searchsorted/midpoint arithmetic (picking
    the code ABOVE every x, or the wrong side of a midpoint), edge handling
    (x below cb[0] / above cb[-1] must resolve to the end codes, not index -1
    or 256), and NaN-through-searchsorted. Brute-force argmin is the oracle."""
    cb = quant8.create_dynamic_map(signed)
    lo = -1.5 if signed else -0.2
    g = torch.Generator().manual_seed(0)
    x = torch.cat(
        [
            torch.rand(20000, generator=g) * (1.5 - lo) + lo,
            cb,  # exactly on a code
            (cb[1:] + cb[:-1]) / 2 + 1e-9,  # just above each midpoint
            (cb[1:] + cb[:-1]) / 2 - 1e-9,  # just below each midpoint
            torch.tensor([-1.0, 0.0, 1.0, -1e-30, 1e-30]),
        ]
    )
    got = quant8.nearest_code(x, cb)
    assert got.dtype == torch.uint8
    want = (x.unsqueeze(1) - cb.unsqueeze(0)).abs().argmin(dim=1)
    # Compare by DISTANCE, not by index: at a midpoint the two candidates tie
    # to within fp32 resolution (~6e-8 absolute), and either is correct. A wrong
    # side of a midpoint is off by a whole half-bin, orders of magnitude more.
    d_got = (x - cb[got.long()]).abs()
    d_want = (x - cb[want]).abs()
    assert bool((d_got <= d_want * 1.001 + 1e-9).all()), "nearest_code returned a farther code"


# ---------------------------------------------------------------------------
# quantize -> dequantize
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("signed", [True, False])
def test_roundtrip_error_is_bounded_on_gaussian_blocks(signed):
    """Quantization error, measured as the optimizer sees it: relative to the
    block's absmax (the scale the codes are normalised by). Median < 1e-2 and
    max < 0.1 are the specified bounds; the measured values are ~2e-3 and
    ~7e-3 (the top bin is 0.014 wide), so the bound has 10x headroom and a
    broken codebook (linear, shifted, unsorted) blows through it.

    ELEMENT-wise relative error on the large elements (|x| > 0.1 absmax, where
    the optimizer's update is decided) is also pinned: median < 2e-2 (measured
    ~1.3e-2; a 64-bin decade cannot do better than ~half a bin of 0.014/0.2).
    """
    x = _randn(8 * B, seed=1)
    if not signed:
        x = x.pow(2)  # a second-moment-like, heavy-tailed, non-negative tensor
    codes, absmax = quant8.quantize_blockwise(x, B, signed)
    assert codes.dtype == torch.uint8 and codes.numel() == 8 * B and absmax.shape == (8,)
    y = quant8.dequantize_blockwise(codes, absmax, x.numel(), B, signed)
    scale = absmax.repeat_interleave(B)
    abs_err = (x - y).abs() / scale
    assert abs_err.median().item() < 1e-2
    assert abs_err.max().item() < 0.1
    big = x.abs() > 0.1 * scale
    rel = ((x - y).abs() / x.abs())[big]
    assert rel.median().item() < 2e-2 and rel.max().item() < 0.1


def test_blockwise_isolates_an_outlier_to_its_own_block():
    """The reason for BLOCK-wise: one huge element must not coarsen the other
    blocks. Kills: a single global absmax (every other block's median error
    would be ~1e4x worse)."""
    x = _randn(4 * B, seed=2)
    x[5] = 1e4
    codes, absmax = quant8.quantize_blockwise(x, B, True)
    y = quant8.dequantize_blockwise(codes, absmax, x.numel(), B, True)
    clean = slice(B, 4 * B)
    err = ((x[clean] - y[clean]).abs() / x[clean].abs().clamp_min(1e-3)).median()
    assert err.item() < 2e-2, err
    assert absmax[0].item() == 1e4 and absmax[1].item() < 6


def test_zero_blocks_and_partial_last_block_roundtrip_exactly_without_nan():
    """Kills: dividing by a zero absmax (NaN codes), and a partial final block
    that leaks its zero padding into the dequantized values or the numel."""
    x = torch.zeros(3 * B + 77)
    x[B : B + 5] = torch.tensor([1.0, -2.0, 3.0, 0.0, 0.5])
    codes, absmax = quant8.quantize_blockwise(x, B, True)
    assert codes.numel() == 4 * B
    y = quant8.dequantize_blockwise(codes, absmax, x.numel(), B, True)
    assert y.shape == x.shape and bool(torch.isfinite(y).all())
    assert torch.equal(y[:B], torch.zeros(B)) and torch.equal(y[2 * B :], torch.zeros(B + 77))
    assert absmax[0].item() == 0.0 and absmax[1].item() == 3.0


def test_dequantize_indexes_with_a_gather_not_a_uint8_mask():
    """THE UINT8 INDEX TRAP: `cb[codes]` with a uint8 `codes` is a boolean-mask
    index in torch, not a gather. dequantize must return cb[code] * absmax for
    codes 0 and 255 (the two ends), not a mask-selected junk tensor."""
    cb = quant8.codebook(True, "cpu")
    codes = torch.tensor([[0, 255, 127, 1]], dtype=torch.uint8)
    out = quant8.dequantize_rows(codes, torch.tensor([2.0]), cb)
    assert torch.equal(out, torch.stack([cb[0], cb[255], cb[127], cb[1]]).unsqueeze(0) * 2.0)


# ---------------------------------------------------------------------------
# the update rule
# ---------------------------------------------------------------------------


def test_fp32_fallback_for_small_tensors_is_bit_exact_vs_torch_adamw():
    """Tensors under `min_8bit_size` keep fp32 moments, and must match
    `torch.optim.AdamW` BITWISE over many steps (decoupled decay, betas,
    bias correction, eps all included). Kills: any reordering of the op
    sequence (lerp_ vs hand EMA, folding bias-correction differently), a
    dropped weight_decay, a mis-forwarded beta2 (0.95 vs torch's 0.999)."""
    a = [_param(4, 5, seed=s) for s in range(3)] + [_param(4095, seed=9)]
    b = [torch.nn.Parameter(p.detach().clone()) for p in a]
    mine = BlockwiseAdamW8bit(a, **HP)
    ref = torch.optim.AdamW(b, **HP)
    assert all("m" in s for s in mine._state), "all four are below min_8bit_size=4096"
    for step in range(40):
        _set_grads(a, step)
        for pa, pb in zip(a, b):
            pb.grad = pa.grad.clone()
        assert mine.step() is True
        ref.step()
    for pa, pb in zip(a, b):
        assert torch.equal(pa, pb)


def test_threshold_is_below_min_8bit_size_inclusive_at_the_boundary():
    """bnb's rule: numel < min_8bit_size stays fp32; numel == min is 8-bit.
    Kills: `<=` vs `<` drift."""
    opt = BlockwiseAdamW8bit([_param(4095), _param(4096)], lr=1e-3)
    assert "m" in opt._state[0] and "m_codes" in opt._state[1]
    assert opt._state[1]["m_codes"].dtype == torch.uint8


def _quadratic(opt_fn, n=8192, steps=300, seed=1):
    g = torch.Generator().manual_seed(seed)
    t = torch.randn(n, generator=g)
    curv = torch.rand(n, generator=g) * 9 + 1
    x = torch.nn.Parameter(torch.zeros(n))
    opt = opt_fn([x])
    loss0 = 0.5 * (curv * t**2).sum().item()
    for _ in range(steps):
        x.grad = None
        loss = 0.5 * (curv * (x - t) ** 2).sum()
        loss.backward()
        opt.step()
    return 0.5 * (curv * (x.detach() - t) ** 2).sum().item(), loss0


def test_300_steps_on_a_convex_quadratic_converge_like_torch_adamw():
    """Both reach <0.5% of the initial loss, and the 8-bit endpoint is within
    1e-3 * L0 of fp32 AdamW's. Kills: state that does not accumulate (a
    zeroed-every-step moment), a lost lr/bias-correction (never converges), a
    codebook that collapses v to zero (the update explodes, m/eps)."""
    hp = dict(lr=5e-2, betas=(0.9, 0.95), weight_decay=0.0)
    l8, l0 = _quadratic(lambda p: BlockwiseAdamW8bit(p, **hp))
    l32, _ = _quadratic(lambda p: torch.optim.AdamW(p, **hp))
    assert l32 < 5e-3 * l0, (l32, l0)
    assert l8 < 5e-3 * l0, (l8, l0)
    assert abs(l8 - l32) < 1e-3 * l0, (l8, l32, l0)


def _linreg(opt_fn, steps=300):
    g = torch.Generator().manual_seed(3)
    d_in, d_out, n = 128, 64, 256  # weight has 8192 elements -> the 8-bit path
    X = torch.randn(n, d_in, generator=g)
    W_true = torch.randn(d_out, d_in, generator=g) / math.sqrt(d_in)
    Y = X @ W_true.T
    W = torch.nn.Parameter(torch.zeros(d_out, d_in))
    bias = torch.nn.Parameter(torch.zeros(d_out))  # 64 elements -> fp32 fallback
    opt = opt_fn([W, bias])
    first = None
    for _ in range(steps):
        opt.zero_grad() if hasattr(opt, "zero_grad") else None
        loss = ((X @ W.T + bias - Y) ** 2).mean()
        first = loss.item() if first is None else first
        loss.backward()
        opt.step()
    return loss.item(), first


def test_300_steps_of_tiny_linear_regression_match_torch_adamw():
    """Mixed 8-bit (weight) + fp32-fallback (bias) parameters in one
    optimizer, on a real backward pass."""
    hp = dict(lr=1e-2, betas=(0.9, 0.95), weight_decay=0.0)
    l8, first = _linreg(lambda p: BlockwiseAdamW8bit(p, **hp))
    l32, _ = _linreg(lambda p: torch.optim.AdamW(p, **hp))
    assert l32 < 1e-3 * first and l8 < 1e-3 * first, (l8, l32, first)
    assert abs(l8 - l32) < 1e-3 * first, (l8, l32, first)


def test_first_step_is_exactly_lr_times_sign_because_bias_correction_is_applied():
    """After one step, m_hat/(sqrt(v_hat)+eps) == sign(g) (to eps), so every
    element moves by exactly lr. The update uses the freshly computed fp32
    moments, not the requantized ones, so this holds on the 8-bit path to fp32
    rounding. Kills: missing bias correction (moves 0.1/sqrt(0.05) = 0.447 lr
    with beta2=0.95), or an update computed from the dequantized-after-quantize
    moments (off by the quantization error, ~1e-3 relative)."""
    p = _param(B * 3, seed=4)
    p0 = p.detach().clone()
    p.grad = _randn(B * 3, seed=5) * 10
    opt = BlockwiseAdamW8bit([p], lr=1e-3, betas=(0.9, 0.95), weight_decay=0.0)
    assert "m_codes" in opt._state[0]
    opt.step()
    moved = (p0 - p.detach()) * p.grad.sign()
    assert torch.allclose(moved, torch.full_like(moved, 1e-3), rtol=1e-4, atol=1e-7), (
        moved.min().item(),
        moved.max().item(),
    )


def test_weight_decay_is_decoupled_from_the_gradient():
    """With a ZERO gradient the moments stay zero, so the only thing that can
    move a weight is the decay: p <- p * (1 - lr*wd), to fp32 rounding. A
    COUPLED (L2) implementation adds wd*p to the gradient and Adam then
    normalises it to a step of ~lr per element -- ~1e3x more movement. Kills:
    coupled decay, and a dropped weight_decay (p unchanged)."""
    p = _param(B * 2, seed=6)
    p0 = p.detach().clone()
    p.grad = torch.zeros_like(p)
    opt = BlockwiseAdamW8bit([p], lr=1e-2, betas=(0.9, 0.95), weight_decay=0.5)
    assert "m_codes" in opt._state[0]
    opt.step()
    assert torch.allclose(p.detach(), p0 * (1 - 1e-2 * 0.5), rtol=1e-6, atol=0.0)
    assert opt._state[0]["m_absmax"].abs().max().item() == 0.0, "decay leaked into the moments"


def test_chunking_is_invisible_to_the_result():
    """Chunks are whole blocks, so processing one block at a time must equal
    processing the tensor in one pass, BITWISE (params and every code/absmax).
    Kills: a chunk boundary that is off by a block, a partial-block mishandled
    at a chunk edge, state written to the wrong row range."""
    n = B * 5 + 123
    outs = []
    for chunk in (B, 2 * B, 10**9):
        p = _param(n, seed=8)
        opt = BlockwiseAdamW8bit([p], chunk_elems=chunk, **HP)
        for step in range(6):
            _set_grads([p], step)
            opt.step()
        outs.append((p.detach().clone(), opt.state_dict()["state"][0]))
    for pw, st in outs[1:]:
        assert torch.equal(pw, outs[0][0])
        for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"):
            assert torch.equal(st[k], outs[0][1][k]), k


def test_padding_of_a_partial_last_block_stays_the_zero_code():
    """n is not a multiple of the block: the pad tail of the last row must keep
    quantizing to exact zero forever, or garbage moments would leak into the
    last block's absmax (and thus its resolution)."""
    n = 2 * B + 100  # 4196 elements: above min_8bit_size, so it is 8-bit
    p = _param(n, seed=10)
    opt = BlockwiseAdamW8bit([p], **HP)
    for step in range(5):
        _set_grads([p], step)
        opt.step()
    st = opt._state[0]
    assert st["m_codes"].shape == (3, B)
    assert bool((st["m_codes"][2, 100:] == quant8.zero_code(True)).all())
    assert bool((st["v_codes"][2, 100:] == quant8.zero_code(False)).all())


def test_param_with_no_grad_is_skipped_and_its_step_count_does_not_advance():
    """Bias correction is per tensor: a tensor that sat out a step must not be
    treated as having taken it. Also: no gradients anywhere is an error, not a
    silent no-op."""
    a, b = _param(B, seed=1), _param(B, seed=2)
    opt = BlockwiseAdamW8bit([a, b], **HP)
    a.grad = _randn(B, seed=3)
    b0 = b.detach().clone()
    assert opt.step() is True
    assert opt._state[0]["step"] == 1 and opt._state[1]["step"] == 0
    assert torch.equal(b, b0)
    a.grad = None
    with pytest.raises(RuntimeError, match="no gradients"):
        opt.step()


def test_empty_local_shard_is_tolerated():
    """FSDP2 can hand a rank an EMPTY shard (a 1-row param over 2 ranks). It
    must construct, step and round-trip rather than dividing by zero blocks."""
    e, p = torch.nn.Parameter(torch.zeros(0)), _param(B, seed=1)
    opt = BlockwiseAdamW8bit([e, p], **HP)
    e.grad, p.grad = torch.zeros(0), _randn(B, seed=2)
    assert opt.step() is True
    opt.load_state_dict(opt.state_dict())


# ---------------------------------------------------------------------------
# grad norm, clipping, the skip gate
# ---------------------------------------------------------------------------


def test_grad_norm_is_recorded_before_clipping_and_clipping_is_applied():
    """THE PPO BLIND-TRIPWIRE INCIDENT's cousin: the norm a tripwire reads must
    be the PRE-clip norm (post-clip it is <= max by construction and a
    10,000x explosion looks like 1.0). And the clip must actually rescale the
    gradients that enter the update.

    Adam is scale-invariant on a single step, so this runs TWO steps (a big
    clipped gradient, then a small one) and compares, bitwise, against an
    unclipped optimizer fed the pre-clipped gradients by hand. Kills: norm
    recorded after clipping; clipping not applied (the two trajectories
    differ by the relative weight of step 1 in the moments); a coefficient
    other than max/(norm+1e-6)."""
    ps = [_param(B * 2, seed=1), _param(B, seed=2), _param(7, seed=3)]  # 8-bit x2 + fp32
    qs = [torch.nn.Parameter(p.detach().clone()) for p in ps]
    hp = dict(lr=1e-2, betas=(0.9, 0.95), weight_decay=0.0)
    clipped = BlockwiseAdamW8bit(ps, max_grad_norm=1.0, **hp)
    manual = BlockwiseAdamW8bit(qs, max_grad_norm=None, **hp)

    for step, scale in enumerate((100.0, 0.01)):
        _set_grads(ps, step, scale=scale)
        true_norm = float(torch.sqrt(sum(p.grad.double().pow(2).sum() for p in ps)))
        for q, p in zip(qs, ps):
            q.grad = p.grad.clone()
        if true_norm > 1.0:
            coef = 1.0 / (true_norm + 1e-6)
            for q in qs:
                q.grad = q.grad * coef
        assert clipped.step() is True
        manual.step()
        assert clipped.last_grad_norm == pytest.approx(true_norm, rel=1e-6), (
            "last_grad_norm must be the GLOBAL PRE-CLIP norm"
        )
    assert clipped.last_grad_norm < 1.0  # the second step was below the cap: not clipped
    for p, q in zip(ps, qs):
        assert torch.equal(p, q)


def test_nonfinite_gradient_skips_the_step_and_leaves_every_byte_of_state_alone():
    """A NaN/inf discovered after the moments were updated has already poisoned
    the second moment. The verdict must come BEFORE any state is touched, the
    norm must read NaN (a 0.0 would read as 'converged'), and a good step after
    the skip must work. Kills: checking finiteness after the update; counting
    the skipped step in the bias-correction step number; returning True."""
    p = _param(B * 2, seed=1)
    opt = BlockwiseAdamW8bit([p], **HP)
    _set_grads([p], 0)
    opt.step()
    snap, p0 = opt.state_dict(), p.detach().clone()
    for bad in (float("nan"), float("inf"), float("-inf")):
        p.grad = _randn(B * 2, seed=3)
        p.grad[17] = bad
        assert opt.step() is False
        assert math.isnan(opt.last_grad_norm)
        assert torch.equal(p, p0)
    assert opt.skipped_steps == 3
    now = opt.state_dict()
    assert now["state"][0]["step"] == snap["state"][0]["step"] == 1
    for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"):
        assert torch.equal(now["state"][0][k], snap["state"][0][k]), k
    _set_grads([p], 1)
    assert opt.step() is True and not torch.equal(p, p0)


def test_second_moment_underflow_cannot_produce_an_m_over_eps_update():
    """THE ZERO-SECOND-MOMENT CATASTROPHE. In a block with one huge gradient
    (1.0) the tiny ones (1e-6) have v/absmax ~ 1e-12, below the smallest
    nonzero code, so their STORED second moment is exactly 0. If the update
    were computed from stored moments, it would be m / (sqrt(0) + 1e-8) -- 100x
    lr per step, i.e. the weight explodes. The update must use the FRESH fp32
    moments (quantization is a storage format), which bounds the per-step move
    by lr / sqrt(1 - beta2) (= 4.47 lr here; measured 2.1 lr).

    Kills: computing the update after requantization; a codebook whose zero code
    swallows small v."""
    p = torch.nn.Parameter(torch.zeros(4096))
    opt = BlockwiseAdamW8bit([p], lr=1e-3, betas=(0.9, 0.95), weight_decay=0.0)
    g = torch.full((4096,), 1e-6)
    g[0] = 1.0
    worst = 0.0
    for _ in range(20):
        before = p.detach().clone()
        p.grad = g.clone()
        opt.step()
        worst = max(worst, (p.detach() - before).abs().max().item())
    assert (opt._state[0]["v_codes"][0, 1:] == 0).all(), "precondition: v really underflowed"
    assert worst <= 1e-3 / math.sqrt(1 - 0.95) * 1.05, worst / 1e-3


# ---------------------------------------------------------------------------
# bf16 parameters, no master copy
# ---------------------------------------------------------------------------


def test_stochastic_rounding_is_unbiased_and_exact_on_representable_values():
    """E[SR(x)] == x, and a value already representable in bf16 is unchanged.
    Kills: round-down (biased), round-to-nearest, a carry that corrupts the
    sign/exponent for negative numbers."""
    gen = torch.Generator().manual_seed(0)
    x = torch.full((200000,), 1.0 + 3e-4)  # between 1.0 and 1.0078125
    out = stochastic_round_to_bf16(x, gen).float()
    assert set(out.unique().tolist()) <= {1.0, 1.0078125}
    assert out.mean().item() == pytest.approx(1.0 + 3e-4, abs=3e-5)
    neg = stochastic_round_to_bf16(-x, gen).float()
    assert neg.mean().item() == pytest.approx(-(1.0 + 3e-4), abs=3e-5)
    exact = torch.tensor([0.0, 1.0, -2.5, 0.15625, 3.0e38 / 2**20]).to(torch.bfloat16).float()
    assert torch.equal(stochastic_round_to_bf16(exact, gen).float(), exact)


def test_bf16_params_without_a_master_copy_still_move_under_tiny_updates():
    """THE NO-MASTER-COPY DEFECT (the silent-no-train class). A bf16 weight of
    1.0 has spacing 7.8e-3; an Adam step of lr=1e-4 is 80x below half a ulp, so
    round-to-nearest returns 1.0 EVERY step and the run trains nothing while
    its loss and grad-norm logs look healthy. Stochastic rounding makes the
    parameter move by the right amount in expectation: 200 steps of lr 1e-4 on a
    constant gradient move the mean by ~0.02.

    Both directions: nearest rounding really is stuck (the control that makes
    this test mean something), and SR really moves. Kills: stochastic_rounding
    ignored / the write-back always rounding to nearest."""
    def run(sr):
        p = torch.nn.Parameter(torch.ones(B * 4, dtype=torch.bfloat16))
        opt = BlockwiseAdamW8bit(
            [p], lr=1e-4, betas=(0.9, 0.95), weight_decay=0.0, stochastic_rounding=sr
        )
        for _ in range(200):
            p.grad = torch.ones_like(p)
            opt.step()
        return p.detach().float().mean().item()

    assert run(False) == 1.0, "control: round-to-nearest should be stuck"
    assert run(True) == pytest.approx(1.0 - 200 * 1e-4, abs=2.5e-3)


def test_bf16_param_with_fp32_grad_is_accepted():
    """The fp32-reduce policy hands back fp32 grads for bf16 params."""
    p = torch.nn.Parameter(_randn(B * 2, seed=1, dtype=torch.bfloat16))
    if hasattr(p, "grad_dtype"):  # newer torch refuses a grad dtype != param dtype by default
        p.grad_dtype = None
    p.grad = _randn(B * 2, seed=2)
    p0 = p.detach().clone()
    BlockwiseAdamW8bit([p], lr=1e-2).step()
    assert p.dtype == torch.bfloat16 and not torch.equal(p, p0)


# ---------------------------------------------------------------------------
# resume
# ---------------------------------------------------------------------------


def _train(opt, params, steps, start=0, dtype=None):
    for step in range(start, start + steps):
        _set_grads(params, step, dtype=dtype)
        opt.step()


@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
def test_resume_from_state_dict_is_bit_exact(dtype):
    """Save at step 5, keep training the original; load into a FRESH optimizer
    over a copy of the step-5 weights and train the same 5 steps: weights and
    every code/absmax byte must be identical. bf16 covers the stochastic
    rounding RNG (its generator state is part of the checkpoint). Kills: a
    state_dict that drops the step count, the absmax, the SR RNG, or the
    hyper-parameters; a load that re-quantizes instead of copying."""
    mk = lambda: [_param(B * 3, seed=1, dtype=dtype), _param(300, seed=2, dtype=dtype),
                  _param(B + 9, seed=3, dtype=dtype)]
    a = mk()
    opt_a = BlockwiseAdamW8bit(a, **HP, seed=5)
    _train(opt_a, a, 5)
    sd = opt_a.state_dict()
    b = [torch.nn.Parameter(p.detach().clone()) for p in a]
    opt_b = BlockwiseAdamW8bit(b, **HP, seed=999)  # a different seed: must be overwritten
    opt_b.load_state_dict(sd)
    _train(opt_a, a, 5, start=5)
    _train(opt_b, b, 5, start=5)
    for pa, pb in zip(a, b):
        assert torch.equal(pa, pb)
    sa, sb = opt_a.state_dict(), opt_b.state_dict()
    for x, y in zip(sa["state"], sb["state"]):
        assert x.keys() == y.keys()
        for k in x:
            assert torch.equal(torch.as_tensor(x[k]), torch.as_tensor(y[k])), k


def test_state_dict_is_a_cpu_snapshot_not_a_live_view():
    """THE SHARED-STATE ALIASING INCIDENT: `.to('cpu')` on a CPU tensor returns
    the same object, so a 'snapshot' built that way is mutated by the next
    step -- uint8 codes that no comparison flags as drifting. And loading one
    dict into two optimizers must not make them share state. Kills: missing
    copy=True in state_dict; a load that stores references."""
    p = _param(B * 2, seed=1)
    opt = BlockwiseAdamW8bit([p], **HP)
    _train(opt, [p], 2)
    sd = opt.state_dict()
    frozen = copy.deepcopy(sd)
    for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"):
        assert sd["state"][0][k].device.type == "cpu"
        assert sd["state"][0][k].data_ptr() != opt._state[0][k].data_ptr()
    _train(opt, [p], 3, start=2)
    for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"):
        assert torch.equal(sd["state"][0][k], frozen["state"][0][k]), f"{k}: snapshot drifted"
    assert sd["state"][0]["step"] == 2

    q1, q2 = _param(B * 2, seed=1), _param(B * 2, seed=1)
    o1, o2 = BlockwiseAdamW8bit([q1], **HP), BlockwiseAdamW8bit([q2], **HP)
    o1.load_state_dict(sd)
    o2.load_state_dict(sd)
    assert o1._state[0]["m_codes"].data_ptr() != o2._state[0]["m_codes"].data_ptr()
    assert o1._state[0]["m_codes"].data_ptr() != sd["state"][0]["m_codes"].data_ptr()


def test_state_dict_survives_a_pickle_roundtrip():
    """It rides `gather_object` (pickle) inside the checkpoint: it must hold no
    live handles (generators, devices) and stay byte-identical."""
    import pickle

    p = _param(B * 2, seed=1, dtype=torch.bfloat16)
    opt = BlockwiseAdamW8bit([p], **HP)
    _train(opt, [p], 3)
    sd = opt.state_dict()
    back = pickle.loads(pickle.dumps(sd))  # round-trips this test's OWN freshly built dict; no untrusted input
    q = torch.nn.Parameter(p.detach().clone())
    other = BlockwiseAdamW8bit([q], **HP)
    other.load_state_dict(back)
    _train(opt, [p], 2, start=3)
    _train(other, [q], 2, start=3)
    assert torch.equal(p, q)


def test_resume_refuses_a_different_world_size_layout_or_format_and_changes_nothing():
    """The state is SHARD-shaped. Loading it under another world size aligns
    rank r's moments to a different slice of the parameter and training
    continues plausibly on corrupted state. A refused load must also leave the
    optimizer untouched (validate everything, then write). Kills: dropping the
    world_size check; checking numel only after partially copying."""
    p = _param(B * 2, seed=1)
    opt = BlockwiseAdamW8bit([p], **HP)
    _train(opt, [p], 2)
    sd = opt.state_dict()

    fresh_p = _param(B * 2, seed=1)
    fresh = BlockwiseAdamW8bit([fresh_p], **HP)
    before = fresh.state_dict()

    with pytest.raises(ValueError, match="world_size"):
        fresh.load_state_dict(dict(sd, world_size=sd["world_size"] + 4))
    wrong_n = copy.deepcopy(sd)
    wrong_n["state"][0]["numel"] += 1
    with pytest.raises(ValueError, match="numel"):
        fresh.load_state_dict(wrong_n)
    with pytest.raises(ValueError, match="block_size"):
        fresh.load_state_dict(dict(sd, block_size=1024))
    with pytest.raises(ValueError, match="format"):
        fresh.load_state_dict({"inner": {}, "state": {}})
    bad_shape = copy.deepcopy(sd)
    bad_shape["state"][0]["m_absmax"] = torch.zeros(5)
    with pytest.raises(ValueError, match="m_absmax"):
        fresh.load_state_dict(bad_shape)
    other = BlockwiseAdamW8bit([_param(B * 2), _param(B * 2)], **HP)
    with pytest.raises(ValueError, match="tensors"):
        other.load_state_dict(sd)

    after = fresh.state_dict()
    assert after["state"][0]["step"] == before["state"][0]["step"] == 0
    for k in ("m_codes", "v_codes", "m_absmax", "v_absmax"):
        assert torch.equal(after["state"][0][k], before["state"][0][k])


# ---------------------------------------------------------------------------
# memory budget
# ---------------------------------------------------------------------------


def test_state_bytes_match_the_2n_plus_8n_over_block_budget():
    """The RSS/VRAM projection must be HONEST: for block-aligned tensors the
    real footprint is exactly 2*n + 8*n/block bytes; small tensors are 8*n; and
    `projected_state_bytes` (usable before allocating anything) agrees with
    what was actually allocated. Kills: an absmax kept per element, fp32
    codes, or a projection that forgets the fp32 fallback tensors."""
    n_big = 100 * B
    params = [_param(n_big, seed=1), _param(300, seed=2), _param(5 * B + 17, seed=3)]
    opt = BlockwiseAdamW8bit(params, **HP)
    assert opt._state[0]["m_codes"].dtype == torch.uint8
    assert opt._state[0]["m_absmax"].dtype == torch.float32
    assert opt._state[0]["m_absmax"].numel() == 100
    big_only = BlockwiseAdamW8bit([_param(n_big)], **HP).state_bytes()
    assert big_only == 2 * n_big + 8 * n_big // B
    got = opt.state_bytes()
    want = BlockwiseAdamW8bit.projected_state_bytes([p.numel() for p in params])
    assert got == want
    assert got < 0.26 * sum(p.numel() for p in params) * 8, "must be ~4x smaller than fp32 moments"
    # the headline number from the module docstring: 7.248e9 params over 4 ranks
    rank_numel = int(7_248_023_552 / 4)
    proj = BlockwiseAdamW8bit.projected_state_bytes([rank_numel])
    assert proj / rank_numel == pytest.approx(2 + 8 / B, rel=1e-3)
    assert 3.3 < proj / 2**30 < 3.5


def test_describe_placement_reports_uint8_codes_on_the_param_device():
    opt = BlockwiseAdamW8bit([_param(B * 2)], **HP)
    ent = {e["key"]: e for e in opt.describe_placement()["state"]}
    assert ent["m_codes"]["dtype"] == str(torch.uint8) and ent["m_codes"]["device"] == "cpu"
    assert ent["v_absmax"]["dtype"] == str(torch.float32)


# ---------------------------------------------------------------------------
# validation
# ---------------------------------------------------------------------------


def test_constructor_refuses_bad_inputs():
    p = _param(B)
    with pytest.raises(ValueError, match="empty"):
        BlockwiseAdamW8bit([], lr=1e-3)
    with pytest.raises(ValueError, match="twice"):
        BlockwiseAdamW8bit([p, p], lr=1e-3)
    with pytest.raises(ValueError, match="betas"):
        BlockwiseAdamW8bit([p], lr=1e-3, betas=(0.9, 1.0))
    with pytest.raises(ValueError, match="max_grad_norm"):
        BlockwiseAdamW8bit([p], lr=1e-3, max_grad_norm=0.0)
    with pytest.raises(ValueError, match="block_size"):
        BlockwiseAdamW8bit([p], lr=1e-3, block_size=0)
    with pytest.raises(ValueError, match="floating"):
        BlockwiseAdamW8bit([torch.zeros(B, dtype=torch.int64)], lr=1e-3)
    with pytest.raises(ValueError, match="contiguous"):
        BlockwiseAdamW8bit([torch.nn.Parameter(torch.zeros(64, 64).t())], lr=1e-3)
    q = _param(B)
    opt = BlockwiseAdamW8bit([q], lr=1e-3)
    q.data = torch.zeros(B + 1)  # a re-sharded parameter: state no longer fits
    q.grad = torch.zeros(B + 1)
    with pytest.raises(ValueError, match="elements"):
        opt.step()


def test_lr_written_through_param_groups_is_what_the_step_uses():
    """`Trainer._apply_scheduled_lr` writes the scheduled lr into
    `param_groups[*]['lr']`; a step that cached the constructor lr would train
    at a constant rate while logging a decaying one (the scheduler-never-applied
    incident). Kills: reading self.lr instead of the group."""
    p, q = _param(B, seed=1), _param(B, seed=1)
    a = BlockwiseAdamW8bit([p], lr=1e-3, weight_decay=0.0)
    b = BlockwiseAdamW8bit([q], lr=7e-3, weight_decay=0.0)
    a.param_groups[0]["lr"] = 7e-3
    for o, x in ((a, p), (b, q)):
        x.grad = _randn(B, seed=2)
        o.step()
    assert torch.equal(p, q)


# ---------------------------------------------------------------------------
# the factory, the shim, the shipped configs
# ---------------------------------------------------------------------------


def _tiny_params():
    return [_param(64, 64, seed=1), _param(64, seed=2)]


@pytest.mark.parametrize("name", ["sft.toml", "pretrain.toml"])
def test_make_optimizer_builds_adamw8bit_for_the_shipped_full_ft_configs(name):
    """THE LAUNCH BLOCKER (review_trainer.md): `load_config(sft.toml) ->
    make_optimizer(...)` raised NotImplementedError, so no full-FT stage could
    start. The shipped configs must now construct, tagged for the banner with
    the 8-bit footprint, and hold no host-side master."""
    from align.config import load_config
    from align.train.offload_adamw import make_optimizer

    cfg = load_config(os.path.join(os.path.dirname(__file__), "..", "configs", name))
    assert cfg.optimizer.name == "adamw8bit"
    params = _tiny_params()
    opt = make_optimizer(cfg, params)
    assert isinstance(opt, BlockwiseAdamW8bit)
    assert opt.banner_tag.startswith("adamw8bit") and "8-bit" in opt.banner_tag
    assert opt.param_groups[0]["lr"] == cfg.optimizer.lr
    assert opt.param_groups[0]["betas"] == tuple(cfg.optimizer.betas)
    assert opt.param_groups[0]["weight_decay"] == cfg.optimizer.weight_decay
    for p in params:
        p.grad = torch.ones_like(p)
    assert opt.step() is True and opt.skipped_steps == 0


def test_make_optimizer_adamw8bit_is_wired_with_the_config_hyperparameters():
    """Kills: a factory that drops betas/eps/weight_decay on the floor (the
    parity test passes against defaults but not against the run config)."""
    import types

    from align.train.offload_adamw import make_optimizer

    o = types.SimpleNamespace(
        name="AdamW8bit", lr=3e-4, betas=[0.8, 0.9], eps=1e-6, weight_decay=0.25
    )
    opt = make_optimizer(types.SimpleNamespace(optimizer=o), _tiny_params())
    g = opt.param_groups[0]
    assert (g["lr"], g["betas"], g["eps"], g["weight_decay"]) == (3e-4, (0.8, 0.9), 1e-6, 0.25)
    assert g["max_grad_norm"] is None
    o.max_grad_norm = 2.0
    assert make_optimizer(types.SimpleNamespace(optimizer=o), _tiny_params()).param_groups[0][
        "max_grad_norm"
    ] == 2.0


def test_adamw8bit_never_silently_becomes_offload_adamw_and_typos_still_raise():
    import types

    from align.train.offload_adamw import OffloadAdamW, make_optimizer

    cfg = types.SimpleNamespace(
        optimizer=types.SimpleNamespace(
            name="adamw8bit", lr=1e-3, betas=[0.9, 0.95], eps=1e-8, weight_decay=0.0
        )
    )
    assert not isinstance(make_optimizer(cfg, _tiny_params()), OffloadAdamW)
    cfg.optimizer.name = "adamw8bti"
    with pytest.raises(ValueError, match="unknown optimizer"):
        make_optimizer(cfg, _tiny_params())
    cfg.optimizer.name = "adafactor"
    with pytest.raises(ValueError, match="BANNED"):
        make_optimizer(cfg, _tiny_params())


def test_shard_local_shim_constructs_delegates_and_round_trips():
    """The legacy class name now builds the real optimizer. Its wrapper
    state_dict (world_size / n_shards / recursive CPU clone) must round-trip
    the 8-bit state bit-exactly, and its step must surface the inner skip/norm."""
    from align.train.offload_adamw import ShardLocalAdamW8bit

    a = [_param(B * 2, seed=1), _param(40, seed=2)]
    opt = ShardLocalAdamW8bit(a, lr=1e-3, weight_decay=0.1)
    _train(opt, a, 3)
    assert opt.last_grad_norm > 0 and opt.skipped_steps == 0
    a[0].grad = torch.full_like(a[0], float("nan"))
    assert opt.step() is False and opt.skipped_steps == 1 and math.isnan(opt.last_grad_norm)
    opt.zero_grad()
    sd = opt.state_dict()
    assert sd["n_shards"] == 2 and sd["skipped_steps"] == 1
    assert sd["inner"]["state"][0]["m_codes"].device.type == "cpu"

    b = [torch.nn.Parameter(p.detach().clone()) for p in a]
    other = ShardLocalAdamW8bit(b, lr=1e-3, weight_decay=0.1)
    other.load_state_dict(sd)
    _train(opt, a, 2, start=3)
    _train(other, b, 2, start=3)
    for pa, pb in zip(a, b):
        assert torch.equal(pa, pb)
    with pytest.raises(ValueError, match="world_size"):
        other.load_state_dict(dict(sd, world_size=sd["world_size"] + 1))
    with pytest.raises(ValueError, match="empty"):
        ShardLocalAdamW8bit([], lr=1e-3)


# ---------------------------------------------------------------------------
# 2-rank gloo + FSDP2 DTensors (CPU only)
# ---------------------------------------------------------------------------

WS = 2
JOIN_TIMEOUT_S = 180


def _free_port() -> int:
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


def _body_dtensor(rank: int, run_dir: str) -> dict:
    """Child body. Linear(64, 129): 129 rows over 2 ranks = 65 + 64 rows ->
    local weights of 4160 (8-bit path) and 4096 (8-bit path, exactly at the
    threshold); the bias shards (65/64 elements) take the fp32 fallback. So one
    optimizer exercises UNEVEN shards, both state paths, and DTensor grads."""
    import torch.distributed as dist
    from torch.distributed.device_mesh import init_device_mesh
    from torch.distributed.fsdp import fully_shard

    torch.manual_seed(0)
    ref = torch.nn.Linear(64, 129)
    sharded = copy.deepcopy(ref)
    fully_shard(sharded, mesh=init_device_mesh("cpu", (WS,)))
    g = torch.Generator().manual_seed(3)
    x, y = torch.randn(256, 64, generator=g), torch.randn(256, 129, generator=g)

    hp = dict(lr=1e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.0)
    ref_opt = torch.optim.AdamW(ref.parameters(), **hp)
    opt = BlockwiseAdamW8bit(sharded.parameters(), **hp)
    # The same algorithm over PLAIN clones of this rank's shards, fed the same
    # local grads: DTensor handling may not change a single bit.
    plain = [torch.nn.Parameter(p.to_local().detach().clone()) for p in sharded.parameters()]
    plain_opt = BlockwiseAdamW8bit(plain, **hp)

    def full(m):
        return {n: p.full_tensor().detach().clone() for n, p in m.named_parameters()}

    def loss(m):
        return ((m(x) - y) ** 2).mean()

    norms, ref_norms, mean_diff, plain_equal = [], [], 0.0, True
    for _ in range(6):
        loss(ref).backward()
        ref_norms.append(float(torch.sqrt(sum(p.grad.double().pow(2).sum() for p in ref.parameters()))))
        ref_opt.step()
        ref_opt.zero_grad()

        loss(sharded).backward()
        for pl, p in zip(plain, sharded.parameters()):
            pl.grad = p.grad.to_local().detach().clone()
        assert opt.step() is True
        assert plain_opt.step() is True
        opt.zero_grad()
        norms.append(opt.last_grad_norm)
        plain_equal &= all(
            torch.equal(pl.detach(), p.to_local()) for pl, p in zip(plain, sharded.parameters())
        )
        got = full(sharded)
        mean_diff = max(
            mean_diff,
            *[float((p.detach() - q).abs().mean()) for (_, p), q in zip(ref.named_parameters(), got.values())],
        )

    # a non-finite grad on ONE rank must make EVERY rank skip
    pre = full(sharded)
    sharded(x).pow(2).mean().backward()
    if rank == 1:
        next(iter(sharded.parameters())).grad.to_local().fill_(float("nan"))
    stepped = opt.step()
    opt.zero_grad()
    post = full(sharded)
    skip_diff = max(float((pre[k] - post[k]).abs().max()) for k in pre)

    # DTensor snapshot / mutate / restore: weights are the model's, so only the
    # optimizer state is restored; the step must continue identically.
    sd = opt.state_dict()
    gathered = [None] * WS
    dist.all_gather_object(gathered, norms)
    return {
        "norms": norms,
        "ref_norms": ref_norms,
        "mean_param_diff": mean_diff,
        "plain_equal": plain_equal,
        "all_rank_norms": gathered,
        "local_numel": [int(p.to_local().numel()) for p in sharded.parameters()],
        "quantized": [bool(s["quantized"]) for s in sd["state"]],
        "stepped_on_nan": stepped,
        "skip_diff": skip_diff,
        "skipped": opt.skipped_steps,
        "sd_world_size": sd["world_size"],
    }


_BODIES = {"dtensor": _body_dtensor}


def _child(rank: int, port: int, mode: str, run_dir: str) -> None:
    import torch.distributed as dist

    os.environ.update(
        RANK=str(rank),
        WORLD_SIZE=str(WS),
        LOCAL_RANK=str(rank),
        MASTER_ADDR="127.0.0.1",
        MASTER_PORT=str(port),
    )
    dist.init_process_group("gloo", rank=rank, world_size=WS)
    code = 0
    try:
        result = _BODIES[mode](rank, run_dir)
        with open(os.path.join(run_dir, f"{mode}_{rank}.json"), "w") as f:
            json.dump(result, f)
    except BaseException:
        with open(os.path.join(run_dir, f"{mode}_{rank}.err"), "w") as f:
            f.write(traceback.format_exc())
        code = 1
    finally:
        try:
            dist.destroy_process_group()
        except Exception:
            pass
    if code:
        sys.exit(code)


def _run(mode: str, run_dir: str, monkeypatch) -> list[dict]:
    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "")  # no GPU is ever touched
    ctx = mp.get_context("spawn")
    port = _free_port()
    procs = [ctx.Process(target=_child, args=(r, port, mode, run_dir)) for r in range(WS)]
    for p in procs:
        p.start()
    for p in procs:
        p.join(JOIN_TIMEOUT_S)
    hung = [p for p in procs if p.is_alive()]
    for p in hung:
        p.kill()
        p.join()
    errs = [
        f"rank {r}:\n{open(os.path.join(run_dir, f'{mode}_{r}.err')).read()}"
        for r in range(WS)
        if os.path.isfile(os.path.join(run_dir, f"{mode}_{r}.err"))
    ]
    assert not hung, f"{len(hung)} rank(s) HUNG in mode {mode!r}\n" + "\n".join(errs)
    assert not errs, "\n".join(errs)
    return [json.load(open(os.path.join(run_dir, f"{mode}_{r}.json"))) for r in range(WS)]


def test_adamw8bit_steps_fsdp2_dtensors_with_a_global_norm_and_a_global_skip(tmp_workdir, monkeypatch):
    """THE FSDP2 DTENSOR DEFECT + THE SHARD-LOCAL NORM DEFECT, for the 8-bit
    optimizer (bnb's fused kernel died on DTensors: 'mixed torch.Tensor and
    DTensor'; OffloadAdamW died the same way until `_localize`).

    (1) Stepping the LOCAL shards of an uneven FSDP2 weight is BIT-IDENTICAL to
    the same optimizer over plain clones of those shards (DTensor handling
    changes nothing), and tracks a single-process torch AdamW on the unsharded
    model on average (mean abs diff < 2e-4 (measured 6e-5) against ~6e-3 of travel; the max is
    not bounded tightly, see the cold-start note in adamw8bit.py).
    (2) `last_grad_norm` is the GLOBAL norm and the SAME number on every rank:
    a rank-local slice norm lets one rank cross a halt tripwire while its peer
    waits in a collective. (3) A NaN on one rank's shard makes both ranks skip.
    (4) State really is shard-shaped and quantized on at least one rank."""
    rows = _run("dtensor", tmp_workdir, monkeypatch)
    for r in rows:
        assert r["plain_equal"], "stepping DTensor shards != stepping plain clones of them"
        assert r["mean_param_diff"] < 2e-4, r["mean_param_diff"]
        # step 0 sees identical weights, so the GLOBAL norm must match to fp32
        # rounding; later steps' weights differ by the 8-bit error (~1e-4 rel).
        assert r["norms"][0] == pytest.approx(r["ref_norms"][0], rel=1e-5)
        assert r["norms"] == pytest.approx(r["ref_norms"], rel=2e-3), (r["norms"], r["ref_norms"])
        assert r["all_rank_norms"][0] == r["all_rank_norms"][1], r["all_rank_norms"]
        assert r["stepped_on_nan"] is False and r["skipped"] == 1, r
        assert r["skip_diff"] == 0.0, "a rank stepped although its peer's gradient was non-finite"
        assert r["sd_world_size"] == WS
        assert r["quantized"][0] is True, "the weight shard must take the 8-bit path"
    assert rows[0]["local_numel"] != rows[1]["local_numel"], "129 rows must shard unevenly"


# ---------------------------------------------------------------------------
# the launch path: shipped configs -> make_optimizer -> a real optimizer
# ---------------------------------------------------------------------------


def _shipped_configs_naming(name):
    import glob
    import os

    from align.config import load_config
    from tests.conftest import REPO_ROOT

    out = []
    for path in sorted(glob.glob(os.path.join(REPO_ROOT, "configs", "*.toml"))):
        try:
            cfg = load_config(path)
        except Exception:
            continue  # stage-specific TOMLs (ppo_stage) are not RunConfigs
        if getattr(getattr(cfg, "optimizer", None), "name", None) == name:
            out.append((os.path.basename(path), cfg))
    return out


def test_every_shipped_adamw8bit_config_constructs_the_real_optimizer_and_steps():
    """THE LAUNCH BLOCKER (review_trainer.md): every shipped full-FT config names adamw8bit and
    make_optimizer used to raise NotImplementedError for it, so no stage could construct an
    optimizer. Build from the REAL shipped config on a tiny stub model, and take a step. Mutant:
    restore the raise (or return OffloadAdamW) -> the isinstance check fails."""
    from align.train.adamw8bit import BlockwiseAdamW8bit
    from align.train.offload_adamw import make_optimizer

    found = _shipped_configs_naming("adamw8bit")
    assert found, "no shipped config names adamw8bit; this test would be vacuous"
    assert "sft.toml" in [n for n, _ in found]
    for fname, cfg in found:
        params = [torch.nn.Parameter(torch.randn(n).bfloat16()) for n in (4096 * 3, 100)]
        opt = make_optimizer(cfg, params)
        assert isinstance(opt, BlockwiseAdamW8bit), (fname, type(opt))
        assert "adamw8bit" in opt.banner_tag and "zero host RSS" in opt.banner_tag
        before = [p.detach().clone() for p in params]
        for p in params:
            p.grad = torch.randn_like(p) * 1e-2
        assert opt.step() is True
        assert any(not torch.equal(a, b.detach()) for a, b in zip(before, params)), fname


def test_optimizer_config_max_grad_norm_reaches_the_optimizer_and_is_recorded_pre_clip():
    """`OptimizerConfig.max_grad_norm` (1.0 by default since G0; None disables clipping explicitly)
    must reach the optimizer, which clips by the GLOBAL norm and records the PRE-clip
    norm. Mutant: make_optimizer drops the field -> the param-group assertion fails (the clip
    arithmetic itself is pinned bitwise by the grad-norm test above)."""
    from align.config import OptimizerConfig, RunConfig
    from align.train.offload_adamw import make_optimizer

    assert OptimizerConfig().max_grad_norm == 1.0  # clip ON by default (G0)

    def build(clip):
        cfg = RunConfig()
        cfg.optimizer.name = "adamw8bit"
        cfg.optimizer.max_grad_norm = clip
        torch.manual_seed(0)
        p = torch.nn.Parameter(torch.randn(8192).bfloat16())
        return p, make_optimizer(cfg, [p])

    p_clip, o_clip = build(1.0)
    p_free, o_free = build(None)
    g = torch.randn(8192) * 100.0  # global norm ~9000
    for p in (p_clip, p_free):
        p.grad = g.bfloat16().clone()
    o_clip.step()
    o_free.step()
    assert o_clip.last_grad_norm == pytest.approx(o_free.last_grad_norm, rel=1e-6)
    assert o_clip.last_grad_norm > 1000.0, "must record the PRE-clip norm, not <= max_grad_norm"
    # Adam is scale-invariant on step 1, so the update cannot show the clip; the group setting can.
    assert o_clip.param_groups[0]["max_grad_norm"] == 1.0
    assert o_free.param_groups[0]["max_grad_norm"] is None


def test_train_config_per_rank_optimizer_state_is_settable_from_struct_mode_toml(tmp_path):
    """The Trainer read this flag with getattr because the field did not exist, so struct-mode TOML
    (unknown keys raise) could not turn it on. It must parse, default OFF, and round-trip."""
    from align.config import TrainConfig, load_config

    assert TrainConfig().per_rank_optimizer_state is False
    f = tmp_path / "c.toml"
    f.write_text("[train]\nper_rank_optimizer_state = true\n")
    assert load_config(str(f)).train.per_rank_optimizer_state is True


def test_train_module_default_extra_state_is_an_empty_noop():
    """Subclasses (GRPO tripwire) override the pair; the base must exist so the Trainer's resume.pt
    always carries the key and `load_extra_state` never needs a getattr guard to be safe."""
    from align.train.train_module import TrainModule

    m = TrainModule.__new__(TrainModule)
    assert m.extra_state() == {}
    assert m.load_extra_state({}) is None


# ---------------------------------------------------------------------------
# adversarial review (2026-10-12): defects found by trying to break the optimizer
# ---------------------------------------------------------------------------


def test_second_moment_overflow_cannot_poison_the_codes_with_nan():
    """THE V-OVERFLOW POISON. The finiteness gate tests the float64 NORM, which stays finite for
    a gradient whose fp32 SQUARE is inf: one element of 1e21 gives v = (1-beta2) * g^2 = inf.
    torch.optim.AdamW survives (m / (inf + eps) = 0); here inf became the block's absmax, the
    requantize divided inf/inf = NaN into the uint8 codes, and the block's weights were NaN one
    step LATER with the gate still seeing only finite gradients.

    Kills: removing the `v_full.clamp_(max=fp32_max)` before requantization."""
    p = torch.nn.Parameter(_randn(B * 4, seed=1))
    ref = torch.nn.Parameter(p.detach().clone())
    opt = BlockwiseAdamW8bit([p], lr=1e-3, betas=(0.9, 0.95))
    ropt = torch.optim.AdamW([ref], lr=1e-3, betas=(0.9, 0.95), weight_decay=0.0)
    for step in range(3):
        g = _randn(B * 4, seed=10 + step)
        if step == 0:
            g[5] = 1e21
        p.grad, ref.grad = g.clone(), g.clone()
        assert opt.step() is True
        ropt.step()
        st = opt._state[0]
        assert torch.isfinite(st["v_absmax"]).all() and torch.isfinite(st["m_absmax"]).all(), step
        assert torch.isfinite(p).all(), f"weights went non-finite at step {step}"
    # The overflowing element itself: torch freezes it (v stays inf), the clamp lets it move by a
    # bounded, finite amount (m ~ 1e20 over sqrt(fp32 max) ~ 1.8e19 is a few lr per step).
    assert abs(p[5].item() - ref[5].item()) < 50 * 1e-3


def test_load_state_dict_rewinds_the_rounding_rng_of_an_optimizer_that_already_stepped():
    """`load_state_dict` promises a byte-for-byte restore. The stochastic-rounding generator is
    lazily created, so a state saved BEFORE the first bf16 step carries no `sr_rng`; loading it
    into an optimizer that has since stepped left that optimizer's generator advanced, and the
    'rewound' run rounded differently from a never-stepped one.

    Kills: not clearing `_generators` on load."""
    torch.manual_seed(0)
    base = (torch.randn(B * 3) * 0.05).to(torch.bfloat16)

    def mk():
        return torch.nn.Parameter(base.clone())

    fresh_p = mk()
    sd0 = BlockwiseAdamW8bit([fresh_p], lr=1e-3).state_dict()
    assert sd0["sr_rng"] == {}, "precondition: nothing stepped, so no generator to save"

    used_p, ref_p = mk(), mk()
    used = BlockwiseAdamW8bit([used_p], lr=1e-3)
    ref = BlockwiseAdamW8bit([ref_p], lr=1e-3)
    for s in range(3):
        _set_grads([used_p], s)
        used.step()
    used_p.data.copy_(base)
    used.load_state_dict(sd0)
    for s in range(4):
        _set_grads([used_p], s)
        _set_grads([ref_p], s)
        used.step()
        ref.step()
    assert torch.equal(used_p, ref_p)


def test_persistent_tiny_gradient_overstep_stays_inside_the_documented_bound():
    """The second-moment underflow is PERMANENT for an element whose gradient stays below
    ~5e-4 of its block's largest (v is re-floored to 0 every step), not the transient the module
    docstring used to claim: measured against torch.optim.AdamW with constant g = 1e-4 of the
    block max, 300 steps, the 8-bit travel is 3.9x at beta2 = 0.95 and 9.9x at 0.999. What must
    hold regardless of how the owner resolves that is the per-step ceiling lr / sqrt(1 - beta2):
    the update comes from the FRESH moment, so a long run of underflowed steps can never turn
    into the m / (0 + eps) = 100 lr step.

    Kills: computing the update from the requantized (stored, zero) v, at ANY step; a floor
    that is applied to the fresh moment."""
    lr, beta2 = 1e-3, 0.95
    p = torch.nn.Parameter(torch.zeros(B * 2))
    opt = BlockwiseAdamW8bit([p], lr=lr, betas=(0.9, beta2), weight_decay=0.0)
    g = torch.full((B * 2,), 1e-4)
    g[0] = 1.0
    worst = 0.0
    for _ in range(200):
        before = p.detach().clone()
        p.grad = g.clone()
        opt.step()
        worst = max(worst, (p.detach() - before)[1:].abs().max().item())
    assert (opt._state[0]["v_codes"][0, 1:] == 0).all(), "precondition: v really is underflowed"
    assert worst <= lr / math.sqrt(1 - beta2) * 1.05, worst / lr


def _body_dtensor_clip(rank: int, run_dir: str) -> dict:
    """Child body: a 1-row layer (rank 1 holds an EMPTY shard), an uneven 129-row layer, in-optimizer
    global clipping and weight decay, against torch clip_grad_norm_ + AdamW on the unsharded model."""
    from torch.distributed.device_mesh import init_device_mesh
    from torch.distributed.fsdp import fully_shard

    class Net(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.a = torch.nn.Linear(64, 1)
            self.b = torch.nn.Linear(64, 129)

        def forward(self, x):
            return torch.cat([self.a(x), self.b(x)], 1)

    torch.manual_seed(0)
    ref = Net()
    start = copy.deepcopy(ref)
    sharded = copy.deepcopy(ref)
    mesh = init_device_mesh("cpu", (WS,))
    fully_shard(sharded.a, mesh=mesh)
    fully_shard(sharded.b, mesh=mesh)
    fully_shard(sharded, mesh=mesh)
    g = torch.Generator().manual_seed(3)
    x, y = torch.randn(128, 64, generator=g), torch.randn(128, 130, generator=g) * 10

    hp = dict(lr=1e-3, betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1)
    ref_opt = torch.optim.AdamW(ref.parameters(), **hp)
    opt = BlockwiseAdamW8bit(sharded.parameters(), max_grad_norm=0.5, **hp)
    rel_norm_err, clipped = [], 0
    for _ in range(8):
        ((ref(x) - y) ** 2).mean().backward()
        total = float(torch.nn.utils.clip_grad_norm_(ref.parameters(), 0.5))
        clipped += total > 0.5
        ref_opt.step()
        ref_opt.zero_grad()
        ((sharded(x) - y) ** 2).mean().backward()
        assert opt.step() is True
        opt.zero_grad()
        rel_norm_err.append(abs(opt.last_grad_norm - total) / total)
    with torch.no_grad():
        diff = max(
            float((p.full_tensor() - r).abs().mean()) for p, r in zip(sharded.parameters(), ref.parameters())
        )
        travel = max(float((r - r0).abs().mean()) for r, r0 in zip(ref.parameters(), start.parameters()))
    sd = opt.state_dict()
    opt.load_state_dict(sd)  # an empty-shard rank must round-trip too
    return {
        "local_numel": [int(p.to_local().numel()) for p in sharded.parameters()],
        "rel_norm_err": rel_norm_err,
        "clipped_steps": int(clipped),
        "mean_diff": diff,
        "travel": travel,
    }


_BODIES["dtensor_clip"] = _body_dtensor_clip


def test_adamw8bit_global_clip_and_empty_shard_match_torch_on_two_ranks(tmp_workdir, monkeypatch):
    """The in-optimizer clip must use the GLOBAL norm: a shard-local norm is ~1/sqrt(world) of the
    true one, so the clip would fire late (or never) and the two ranks would scale differently. A
    rank whose shard of a 1-row layer is EMPTY (0 local elements) must still join every
    collective, skip its own update, and round-trip its state.

    Kills: `coef` computed from a rank-local norm; skipping `_grad_norm`'s all-reduce; treating a
    zero-element shard as 'no gradient' (the 'no gradients on any parameter' raise would hang the
    peer in its all-reduce)."""
    rows = _run("dtensor_clip", tmp_workdir, monkeypatch)
    assert rows[1]["local_numel"][:2] == [0, 0], "rank 1 must hold the empty shard of the 1-row layer"
    for r in rows:
        assert r["clipped_steps"] >= 6, "the clip must actually be active in this test"
        assert max(r["rel_norm_err"]) < 1e-3, r["rel_norm_err"]
        assert r["mean_diff"] < 0.02 * r["travel"], (r["mean_diff"], r["travel"])