"""
train_gpt_kernel.py — Parameter-Golf training pipeline informed by the
formal Lean 4 theorems in formal-lean/CriticalEigenvalue.lean.

Three mathematical objects from the Kernel formalization are wired
directly into the model architecture:

  1. Coherence activation  C(r) = 2r / (1 + r²)
     Defined in §5 of CriticalEigenvalue.lean and proved to satisfy:
       • C(r) ≤ 1  for r ≥ 0             (AM–GM bound)
       • C(1) = 1  (unique maximum)
       • C(−r) = −C(r)                   (odd symmetry)
       • C(r)² + ((r²−1)/(1+r²))² = 1   (Pythagorean identity §18)
       • C(exp λ) = sech λ               (Lyapunov duality §10)
     Used here as the nonlinearity inside every MLP block, replacing the
     standard relu² activation from the baseline.

  2. Silver ratio  δS = 1 + √2 ≈ 2.414
     Defined in §7 (Proposition 4) and proved to satisfy:
       • δS = 2 + 1/δS                   (self-similarity §20)
       • δS is the positive root of x²−2x−1 = 0
     Used here as the MLP hidden-dimension multiplier, so the MLP
     expands from d_model to ⌊δS · d_model⌋ neurons.

  3. Z/8Z rotational memory  (μ^8 = 1, §2 and §15)
     The critical eigenvalue μ = exp(3πi/4) generates an exact 8-cycle.
     The default number of attention heads is set to 8, so that each head
     occupies one slot of the cyclic group, providing uniform coverage of
     the 8 distinct phase positions proven distinct in §3.

Hard stop: to stay readable for newcomers, keep this file under 1500 lines.
"""

from __future__ import annotations

import copy
import glob
import io
import math
import os
import random
import subprocess
import sys
import time
import uuid
import zlib
from pathlib import Path

import numpy as np
import sentencepiece as spm
import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch import Tensor, nn
from torch.nn.parallel import DistributedDataParallel as DDP

# ─────────────────────────────────────────────────────────────────────────────
# KERNEL CONSTANTS  (from formal-lean/CriticalEigenvalue.lean)
# ─────────────────────────────────────────────────────────────────────────────

# Silver ratio δS = 1 + √2  (CriticalEigenvalue.lean §7, Proposition 4).
# Self-similar: δS = 2 + 1/δS.  Positive root of x² − 2x − 1 = 0.
SILVER_RATIO: float = 1.0 + math.sqrt(2)  # ≈ 2.4142135…

# Critical eigenvalue angle θ = 3π/4, so μ = exp(I·θ).
# μ^8 = 1 (CriticalEigenvalue.lean §2) and gcd(3,8)=1 (§3).
MU_ANGLE: float = 3.0 * math.pi / 4.0  # 135°

# Number of distinct μ-orbit slots: Z/8Z (§15).
MU_ORBIT_SIZE: int = 8

# Slow-precession period D (CriticalEigenvalue.lean §14, precession_period_factor).
# Lean theorem: 9 × D = 123456789  (nine copies of D fill the palindrome number).
# Also: gcd(8, D) = 1 (precession_gcd_one) — fast 8-cycle and slow precession are coprime.
PRECESSION_PERIOD: int = 13_717_421  # D = 123456789 / 9

# Phase increment per step: ΔΦ₀ = 2π / D  (CriticalEigenvalue.lean §21, phase_full_cycle).
# After PRECESSION_PERIOD steps the phase accumulation returns exactly to 2π.
PRECESSION_DELTA_PHI: float = 2.0 * math.pi / PRECESSION_PERIOD

# ─────────────────────────────────────────────────────────────────────────────
# PALINDROME PRECESSION LR FORMULA  (CriticalEigenvalue.lean §9 + §16)
# ─────────────────────────────────────────────────────────────────────────────

def palindrome_precession_scale(r: float) -> float:
    """Learning-rate multiplier from the palindrome precession formula.

    Evaluates the coherence function C(r) = 2r / (1 + r²) at
    r = (remaining_steps / warmdown_steps) ∈ [0, 1].

    Lean machine-checked properties carried over (CriticalEigenvalue.lean):
      • C(r) ≤ 1 for r ≥ 0             (AM–GM bound, §5 coherence_le_one)
      • C(1) = 1                         (unique maximum, §5 coherence_eq_one_iff)
      • C(r) = C(1/r)                    (even symmetry, §8 coherence_symm)
      • R(r) = 0 ↔ C(r) = 1            (palindrome equilibrium, §11)
      • R(1/r) = −R(r)                   (anti-symmetry, §9 palindrome_residual_antisymm)
        — ensures the warmdown profile mirrors a warmup under r ↦ 1/r,
          making the full LR arc a palindrome in the Lean-theorem sense.
      • precession_preserves_coherence   (§16) — precession is zero-overhead;
        the scale C(r) is invariant under β ↦ e^{iθ}·β for any real θ.

    Args:
        r: ratio of remaining steps to warmdown window, clipped to [0, 1].

    Returns:
        LR multiplier in [0, 1].  Returns 1.0 for r ≥ 1 (outside warmdown),
        0.0 for r ≤ 0 (training finished).
    """
    if r >= 1.0:
        return 1.0
    if r <= 0.0:
        return 0.0
    return 2.0 * r / (1.0 + r * r)


def _selfcheck_palindrome_precession() -> None:
    """Self-check: verify palindrome precession formula against Lean theorem properties.

    This mirrors the machine-checked proofs in CriticalEigenvalue.lean and
    confirms that palindrome_precession_scale satisfies all expected invariants.
    Runs at import time; raises AssertionError on any violation.
    """
    _eps = 1e-10
    _delta_s = SILVER_RATIO

    def _C(r: float) -> float:
        return 2.0 * r / (1.0 + r * r)

    def _Res(r: float) -> float:
        return (r - 1.0 / r) / _delta_s

    # C(1) = 1  (coherence_eq_one_iff, §5)
    assert abs(_C(1.0) - 1.0) < _eps, "C(1) must equal 1 (coherence_eq_one_iff)"
    # C(0) = 0
    assert abs(_C(0.0)) < _eps, "C(0) must equal 0"
    # C(r) ≤ 1 for r ≥ 0  (coherence_le_one, §5)
    for _r in (0.1, 0.5, 1.0, 1.5, 2.0, 10.0):
        assert _C(_r) <= 1.0 + _eps, f"C({_r}) must be ≤ 1 (coherence_le_one)"
    # C(−r) = −C(r): odd / anti-symmetric  (§5)
    for _r in (0.3, 0.7, 1.0, 2.0, 5.0):
        assert abs(_C(-_r) + _C(_r)) < _eps, (
            f"C must be odd: C(-{_r}) ≠ -C({_r}) (coherence odd symmetry §5)"
        )
    # C(r) = C(1/r): reciprocal symmetry  (coherence_symm, §8)
    for _r in (0.5, 2.0, 3.0):
        assert abs(_C(_r) - _C(1.0 / _r)) < _eps, (
            f"C must satisfy C({_r}) = C(1/{_r}) (coherence_symm §8)"
        )
    # Monotonicity: C strictly increasing on (0,1], strictly decreasing on [1,∞)  (§13)
    _mono_inc = (0.2, 0.4, 0.6, 0.8, 1.0)
    for _a, _b in zip(_mono_inc, _mono_inc[1:]):
        assert _C(_a) < _C(_b), (
            f"C must be strictly increasing on (0,1]: C({_a})={_C(_a):.6f} ≥ C({_b})={_C(_b):.6f}"
        )
    _mono_dec = (1.0, 1.5, 2.0, 3.0, 5.0)
    for _a, _b in zip(_mono_dec, _mono_dec[1:]):
        assert _C(_a) > _C(_b), (
            f"C must be strictly decreasing on [1,∞): C({_a})={_C(_a):.6f} ≤ C({_b})={_C(_b):.6f}"
        )
    # Lyapunov–coherence duality: C(exp λ) = sech λ = 1/cosh λ  (Theorem 14, §10)
    for _lam in (-2.0, -1.0, 0.0, 0.5, 1.0, 2.0):
        _expected = 1.0 / math.cosh(_lam)
        _got = _C(math.exp(_lam))
        assert abs(_got - _expected) < _eps, (
            f"Lyapunov duality violated at λ={_lam}: C(exp λ)={_got:.10f} ≠ sech λ={_expected:.10f}"
        )
    # R(1) = 0  (palindrome_residual_zero_iff, §9)
    assert abs(_Res(1.0)) < _eps, "Res(1) must equal 0 (palindrome_residual_zero_iff)"
    # R(1/r) = −R(r)  (palindrome_residual_antisymm, §9)
    for _r in (0.5, 2.0, 3.0):
        assert abs(_Res(1.0 / _r) + _Res(_r)) < _eps, (
            f"palindrome_residual_antisymm violated at r={_r}"
        )
    # palindrome_precession_scale boundary values
    assert abs(palindrome_precession_scale(1.0) - 1.0) < _eps, "scale(1) must be 1"
    assert abs(palindrome_precession_scale(0.0)) < _eps, "scale(0) must be 0"
    assert palindrome_precession_scale(1.5) == 1.0, "scale(r>=1) must be 1"
    assert palindrome_precession_scale(-0.1) == 0.0, "scale(r<=0) must be 0"
    # LR warmdown: scale is non-decreasing as r increases toward 1 on [0,1]
    _prev = 0.0
    for _r in (0.1, 0.3, 0.5, 0.7, 0.9, 1.0):
        _s = palindrome_precession_scale(_r)
        assert _s >= _prev - _eps, (
            f"palindrome_precession_scale must be non-decreasing on [0,1]: "
            f"scale({_r})={_s:.6f} < scale(prev)={_prev:.6f}"
        )
        _prev = _s
    # Pythagorean identity: C(r)² + ((r²-1)/(1+r²))² = 1  (§18)
    for _r in (0.5, 1.0, 2.0, 3.0):
        _c = _C(_r)
        _s = (_r * _r - 1.0) / (1.0 + _r * _r)
        assert abs(_c * _c + _s * _s - 1.0) < _eps, (
            f"Pythagorean identity violated at r={_r}"
        )


_selfcheck_palindrome_precession()


def _selfcheck_kernel_constants() -> None:
    """Validate global kernel constants against Lean-proved identities.

    All equalities are machine-checked in formal-lean/CriticalEigenvalue.lean.
    Runs at import time; raises AssertionError with an actionable message on
    any violation.  Safe to run on all ranks: no I/O or distributed calls.
    """
    _eps = 1e-12

    # ── Palindrome arithmetic (§14) ───────────────────────────────────────────

    # 9 × D = 123456789  (palindrome_period_factor, §14)
    # The palindrome number 123456789 is exactly 9 copies of the period D.
    _nine_d = 9 * PRECESSION_PERIOD
    assert _nine_d == 123_456_789, (
        f"9 * PRECESSION_PERIOD must be 123456789; got {_nine_d}. "
        f"Check PRECESSION_PERIOD={PRECESSION_PERIOD}."
    )

    # gcd(8, D) = 1  (precession_gcd_one, §14)
    # The 8-cycle and the slow precession are coprime — they share no common
    # period and therefore cover the full Z/8Z × Z/D torus without collision.
    _g = math.gcd(8, PRECESSION_PERIOD)
    assert _g == 1, (
        f"gcd(8, PRECESSION_PERIOD) must be 1; got {_g}. "
        f"Fast 8-cycle and slow precession must be coprime (precession_gcd_one §14)."
    )

    # ── Phase accumulation (§21) ──────────────────────────────────────────────

    # D × ΔΦ₀ = 2π  (phase_full_cycle, §21)
    # After exactly PRECESSION_PERIOD steps the accumulated phase returns to 2π.
    _full_cycle = PRECESSION_PERIOD * PRECESSION_DELTA_PHI
    assert abs(_full_cycle - 2.0 * math.pi) < _eps, (
        f"PRECESSION_PERIOD * PRECESSION_DELTA_PHI must equal 2π; "
        f"got {_full_cycle:.15f}, expected {2.0 * math.pi:.15f} "
        f"(phase_full_cycle §21)."
    )

    # ── Silver ratio (§7, §20) ────────────────────────────────────────────────

    # δS = 2 + 1/δS  (silver_ratio_self_similar, §20)
    assert abs(SILVER_RATIO - (2.0 + 1.0 / SILVER_RATIO)) < _eps, (
        f"SILVER_RATIO must satisfy δS = 2 + 1/δS; "
        f"got δS={SILVER_RATIO}, 2+1/δS={2.0 + 1.0 / SILVER_RATIO} "
        f"(silver_ratio_self_similar §20)."
    )

    # δS is the positive root of x² − 2x − 1 = 0  (§7 Proposition 4)
    _poly = SILVER_RATIO ** 2 - 2.0 * SILVER_RATIO - 1.0
    assert abs(_poly) < _eps, (
        f"SILVER_RATIO must satisfy x²−2x−1=0; residual={_poly:.2e} "
        f"(silver_ratio minimal polynomial §7)."
    )

    # ── Critical eigenvalue (§1, §2, §15) ────────────────────────────────────

    # μ angle = 3π/4  (§1 critical eigenvalue definition)
    assert abs(MU_ANGLE - 3.0 * math.pi / 4.0) < _eps, (
        f"MU_ANGLE must be 3π/4={3.0 * math.pi / 4.0:.15f}; got {MU_ANGLE}."
    )

    # Z/8Z orbit size = 8  (mu_pow_eight §2, rotational_memory §15)
    assert MU_ORBIT_SIZE == 8, (
        f"MU_ORBIT_SIZE must be 8 (mu_pow_eight §2, Z/8Z §15); got {MU_ORBIT_SIZE}."
    )

    # |μ| = 1 on the unit circle: cos²(3π/4) + sin²(3π/4) = 1  (mu_abs_one §2)
    _mu_abs_sq = math.cos(MU_ANGLE) ** 2 + math.sin(MU_ANGLE) ** 2
    assert abs(_mu_abs_sq - 1.0) < _eps, (
        f"|μ|² must equal 1; got {_mu_abs_sq:.15f} (mu_abs_one §2)."
    )

    # gcd(3, 8) = 1: μ = ζ^3 is a primitive 8th root (mu_powers_distinct §3)
    assert math.gcd(3, MU_ORBIT_SIZE) == 1, (
        f"gcd(3, MU_ORBIT_SIZE) must be 1 for μ to generate all 8 distinct "
        f"orbit slots (mu_powers_distinct §3); got gcd(3,{MU_ORBIT_SIZE})={math.gcd(3, MU_ORBIT_SIZE)}."
    )


_selfcheck_kernel_constants()

# ─────────────────────────────────────────────────────────────────────────────
# FEATURE FLAGS  (env-var-gated; all default to off — baseline behavior unchanged)
# ─────────────────────────────────────────────────────────────────────────────

# Lyapunov gain: replace unbounded q_gain with sech(λ)=C(exp λ)  (§10 duality).
_USE_LYAPUNOV_GAIN: bool = os.environ.get("USE_LYAPUNOV_GAIN", "0") == "1"
# μ-head phase binding: rotate q,k of head h by h·MU_ANGLE  (§15 Z/8Z slots).
_USE_MU_PHASE: bool = os.environ.get("USE_MU_PHASE", "0") == "1"
# Slow precession: advance head phases by DELTA_PHI each step (requires USE_MU_PHASE).
_USE_PRECESSION: bool = os.environ.get("USE_PRECESSION", "0") == "1"

# ─────────────────────────────────────────────────────────────────────────────
# HYPERPARAMETERS
# ─────────────────────────────────────────────────────────────────────────────

class Hyperparameters:
    # Data paths are shard globs produced by the existing preprocessing pipeline.
    data_path = os.environ.get("DATA_PATH", "./data/datasets/fineweb10B_sp1024")
    train_files = os.path.join(data_path, "fineweb_train_*.bin")
    val_files = os.path.join(data_path, "fineweb_val_*.bin")
    tokenizer_path = os.environ.get("TOKENIZER_PATH", "./data/tokenizers/fineweb_1024_bpe.model")
    run_id = os.environ.get("RUN_ID", str(uuid.uuid4()))
    seed = int(os.environ.get("SEED", 1337))

    # Validation cadence and batch size.
    val_batch_size = int(os.environ.get("VAL_BATCH_SIZE", 524_288))
    val_loss_every = int(os.environ.get("VAL_LOSS_EVERY", 1000))
    train_log_every = int(os.environ.get("TRAIN_LOG_EVERY", 200))

    # Training length.
    iterations = int(os.environ.get("ITERATIONS", 20000))
    warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 1200))
    warmup_steps = int(os.environ.get("WARMUP_STEPS", 20))
    train_batch_tokens = int(os.environ.get("TRAIN_BATCH_TOKENS", 524_288))
    train_seq_len = int(os.environ.get("TRAIN_SEQ_LEN", 1024))
    max_wallclock_seconds = float(os.environ.get("MAX_WALLCLOCK_SECONDS", 600.0))
    qk_gain_init = float(os.environ.get("QK_GAIN_INIT", 1.5))

    # Model shape.
    # num_heads=8: one head per slot of the Z/8Z orbit (CriticalEigenvalue.lean §15).
    vocab_size = int(os.environ.get("VOCAB_SIZE", 1024))
    num_layers = int(os.environ.get("NUM_LAYERS", 9))
    num_kv_heads = int(os.environ.get("NUM_KV_HEADS", 4))
    model_dim = int(os.environ.get("MODEL_DIM", 512))
    num_heads = int(os.environ.get("NUM_HEADS", MU_ORBIT_SIZE))  # 8 → Z/8Z
    tie_embeddings = bool(int(os.environ.get("TIE_EMBEDDINGS", "1")))
    rope_base = float(os.environ.get("ROPE_BASE", 10000.0))
    logit_softcap = float(os.environ.get("LOGIT_SOFTCAP", 30.0))

    # MLP hidden width is ⌊δS · model_dim⌋ unless overridden.
    # δS ≈ 2.414 (CriticalEigenvalue.lean §7).
    # Re-read MODEL_DIM from env so that overrides are honoured even when MLP_HIDDEN is unset.
    mlp_hidden = int(os.environ.get("MLP_HIDDEN", round(SILVER_RATIO * int(os.environ.get("MODEL_DIM", 512)))))

    # Optimizer hyperparameters.
    embed_lr = float(os.environ.get("EMBED_LR", 0.6))
    head_lr = float(os.environ.get("HEAD_LR", 0.008))
    tied_embed_lr = float(os.environ.get("TIED_EMBED_LR", 0.05))
    tied_embed_init_std = float(os.environ.get("TIED_EMBED_INIT_STD", 0.005))
    matrix_lr = float(os.environ.get("MATRIX_LR", 0.04))
    scalar_lr = float(os.environ.get("SCALAR_LR", 0.04))
    muon_momentum = float(os.environ.get("MUON_MOMENTUM", 0.95))
    muon_backend_steps = int(os.environ.get("MUON_BACKEND_STEPS", 5))
    muon_momentum_warmup_start = float(os.environ.get("MUON_MOMENTUM_WARMUP_START", 0.85))
    muon_momentum_warmup_steps = int(os.environ.get("MUON_MOMENTUM_WARMUP_STEPS", 500))
    beta1 = float(os.environ.get("BETA1", 0.9))
    beta2 = float(os.environ.get("BETA2", 0.95))
    adam_eps = float(os.environ.get("ADAM_EPS", 1e-8))
    grad_clip_norm = float(os.environ.get("GRAD_CLIP_NORM", 0.0))

# ─────────────────────────────────────────────────────────────────────────────
# MUON OPTIMIZER  (unchanged from baseline train_gpt.py)
# ─────────────────────────────────────────────────────────────────────────────

def zeropower_via_newtonschulz5(G: Tensor, steps: int = 10, eps: float = 1e-7) -> Tensor:
    a, b, c = (3.4445, -4.7750, 2.0315)
    X = G.bfloat16()
    X /= X.norm() + eps
    transposed = G.size(0) > G.size(1)
    if transposed:
        X = X.T
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * A @ A
        X = a * X + B @ X
    return X.T if transposed else X


class Muon(torch.optim.Optimizer):
    def __init__(self, params, lr: float, momentum: float, backend_steps: int, nesterov: bool = True):
        super().__init__(
            params,
            dict(lr=lr, momentum=momentum, backend_steps=backend_steps, nesterov=nesterov),
        )

    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        distributed = dist.is_available() and dist.is_initialized()
        world_size = dist.get_world_size() if distributed else 1
        rank = dist.get_rank() if distributed else 0

        for group in self.param_groups:
            params = group["params"]
            if not params:
                continue
            lr = group["lr"]
            momentum = group["momentum"]
            backend_steps = group["backend_steps"]
            nesterov = group["nesterov"]

            total_params = sum(int(p.numel()) for p in params)
            updates_flat = torch.zeros(total_params, device=params[0].device, dtype=torch.bfloat16)

            curr = 0
            for i, p in enumerate(params):
                if i % world_size == rank and p.grad is not None:
                    g = p.grad
                    state = self.state[p]
                    if "momentum_buffer" not in state:
                        state["momentum_buffer"] = torch.zeros_like(g)
                    buf = state["momentum_buffer"]
                    buf.mul_(momentum).add_(g)
                    if nesterov:
                        g = g.add(buf, alpha=momentum)
                    g = zeropower_via_newtonschulz5(g, steps=backend_steps)
                    g *= max(1, g.size(0) / g.size(1)) ** 0.5
                    updates_flat[curr : curr + p.numel()] = g.reshape(-1)
                curr += p.numel()

            if distributed:
                dist.all_reduce(updates_flat, op=dist.ReduceOp.SUM)

            curr = 0
            for p in params:
                g = updates_flat[curr : curr + p.numel()].view_as(p).to(dtype=p.dtype)
                p.add_(g, alpha=-lr)
                curr += p.numel()

        return loss

# ─────────────────────────────────────────────────────────────────────────────
# TOKENIZER-AGNOSTIC EVALUATION  (unchanged from baseline)
# ─────────────────────────────────────────────────────────────────────────────

def build_sentencepiece_luts(
    sp: spm.SentencePieceProcessor, vocab_size: int, device: torch.device
) -> tuple[Tensor, Tensor, Tensor]:
    sp_vocab_size = int(sp.vocab_size())
    table_size = max(sp_vocab_size, vocab_size)
    base_bytes_np = np.zeros((table_size,), dtype=np.int16)
    has_leading_space_np = np.zeros((table_size,), dtype=np.bool_)
    is_boundary_token_np = np.ones((table_size,), dtype=np.bool_)
    for token_id in range(sp_vocab_size):
        if sp.is_control(token_id) or sp.is_unknown(token_id) or sp.is_unused(token_id):
            continue
        is_boundary_token_np[token_id] = False
        if sp.is_byte(token_id):
            base_bytes_np[token_id] = 1
            continue
        piece = sp.id_to_piece(token_id)
        if piece.startswith("▁"):
            has_leading_space_np[token_id] = True
            piece = piece[1:]
        base_bytes_np[token_id] = len(piece.encode("utf-8"))
    return (
        torch.tensor(base_bytes_np, dtype=torch.int16, device=device),
        torch.tensor(has_leading_space_np, dtype=torch.bool, device=device),
        torch.tensor(is_boundary_token_np, dtype=torch.bool, device=device),
    )


def load_validation_tokens(pattern: str, seq_len: int) -> Tensor:
    files = [Path(p) for p in sorted(glob.glob(pattern))]
    if not files:
        raise FileNotFoundError(f"No files found for pattern: {pattern}")
    tokens = torch.cat([load_data_shard(file) for file in files]).contiguous()
    usable = ((tokens.numel() - 1) // seq_len) * seq_len
    if usable <= 0:
        raise ValueError(f"Validation split is too short for TRAIN_SEQ_LEN={seq_len}")
    return tokens[: usable + 1]


def eval_val(
    args: Hyperparameters,
    model: nn.Module,
    rank: int,
    world_size: int,
    device: torch.device,
    grad_accum_steps: int,
    val_tokens: Tensor,
    base_bytes_lut: Tensor,
    has_leading_space_lut: Tensor,
    is_boundary_token_lut: Tensor,
) -> tuple[float, float]:
    local_batch_tokens = args.val_batch_size // (world_size * grad_accum_steps)
    if local_batch_tokens < args.train_seq_len:
        raise ValueError(
            "VAL_BATCH_SIZE must provide at least one sequence per rank; "
            f"got VAL_BATCH_SIZE={args.val_batch_size}, WORLD_SIZE={world_size}, "
            f"GRAD_ACCUM_STEPS={grad_accum_steps}, TRAIN_SEQ_LEN={args.train_seq_len}"
        )
    local_batch_seqs = local_batch_tokens // args.train_seq_len
    total_seqs = (val_tokens.numel() - 1) // args.train_seq_len
    seq_start = (total_seqs * rank) // world_size
    seq_end = (total_seqs * (rank + 1)) // world_size
    val_loss_sum = torch.zeros((), device=device, dtype=torch.float64)
    val_token_count = torch.zeros((), device=device, dtype=torch.float64)
    val_byte_count = torch.zeros((), device=device, dtype=torch.float64)

    model.eval()
    with torch.inference_mode():
        for batch_seq_start in range(seq_start, seq_end, local_batch_seqs):
            batch_seq_end = min(batch_seq_start + local_batch_seqs, seq_end)
            raw_start = batch_seq_start * args.train_seq_len
            raw_end = batch_seq_end * args.train_seq_len + 1
            local = val_tokens[raw_start:raw_end].to(device=device, dtype=torch.int64, non_blocking=True)
            x = local[:-1].reshape(-1, args.train_seq_len)
            y = local[1:].reshape(-1, args.train_seq_len)
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
                batch_loss = model(x, y).detach()
            batch_token_count = float(y.numel())
            val_loss_sum += batch_loss.to(torch.float64) * batch_token_count
            val_token_count += batch_token_count
            prev_ids = x.reshape(-1)
            tgt_ids = y.reshape(-1)
            token_bytes = base_bytes_lut[tgt_ids].to(dtype=torch.int16)
            token_bytes += (has_leading_space_lut[tgt_ids] & ~is_boundary_token_lut[prev_ids]).to(dtype=torch.int16)
            val_byte_count += token_bytes.to(torch.float64).sum()

    if dist.is_available() and dist.is_initialized():
        dist.all_reduce(val_loss_sum, op=dist.ReduceOp.SUM)
        dist.all_reduce(val_token_count, op=dist.ReduceOp.SUM)
        dist.all_reduce(val_byte_count, op=dist.ReduceOp.SUM)

    val_loss = val_loss_sum / val_token_count
    bits_per_token = val_loss.item() / math.log(2.0)
    tokens_per_byte = val_token_count.item() / val_byte_count.item()
    model.train()
    return float(val_loss.item()), float(bits_per_token * tokens_per_byte)

# ─────────────────────────────────────────────────────────────────────────────
# POST-TRAINING QUANTIZATION  (unchanged from baseline)
# ─────────────────────────────────────────────────────────────────────────────

CONTROL_TENSOR_NAME_PATTERNS = tuple(
    pattern
    for pattern in os.environ.get(
        "CONTROL_TENSOR_NAME_PATTERNS",
        "attn_scale,attn_scales,mlp_scale,mlp_scales,resid_mix,resid_mixes,q_gain,q_lambda,mu_phase_gate,skip_weight,skip_weights",
    ).split(",")
    if pattern
)
INT8_KEEP_FLOAT_FP32_NAME_PATTERNS = tuple(
    pattern
    for pattern in os.environ.get(
        "INT8_KEEP_FLOAT_FP32_NAME_PATTERNS",
        ",".join(CONTROL_TENSOR_NAME_PATTERNS),
    ).split(",")
    if pattern
)
INT8_KEEP_FLOAT_MAX_NUMEL = 65_536
INT8_KEEP_FLOAT_STORE_DTYPE = torch.float16
INT8_PER_ROW_SCALE_DTYPE = torch.float16
INT8_CLIP_PERCENTILE = 99.99984
INT8_CLIP_Q = INT8_CLIP_PERCENTILE / 100.0


def tensor_nbytes(t: Tensor) -> int:
    return int(t.numel()) * int(t.element_size())


def keep_float_tensor(name: str, t: Tensor, passthrough_orig_dtypes: dict[str, str]) -> Tensor:
    if any(pattern in name for pattern in INT8_KEEP_FLOAT_FP32_NAME_PATTERNS):
        return t.float().contiguous()
    if t.dtype in {torch.float32, torch.bfloat16}:
        passthrough_orig_dtypes[name] = str(t.dtype).removeprefix("torch.")
        return t.to(dtype=INT8_KEEP_FLOAT_STORE_DTYPE).contiguous()
    return t


def quantize_float_tensor(t: Tensor) -> tuple[Tensor, Tensor]:
    t32 = t.float()
    if t32.ndim == 2:
        clip_abs = (
            torch.quantile(t32.abs(), INT8_CLIP_Q, dim=1)
            if t32.numel()
            else torch.empty((t32.shape[0],), dtype=torch.float32)
        )
        clipped = torch.maximum(torch.minimum(t32, clip_abs[:, None]), -clip_abs[:, None])
        scale = (clip_abs / 127.0).clamp_min(1.0 / 127.0)
        q = torch.clamp(torch.round(clipped / scale[:, None]), -127, 127).to(torch.int8).contiguous()
        return q, scale.to(dtype=INT8_PER_ROW_SCALE_DTYPE).contiguous()
    clip_abs = float(torch.quantile(t32.abs().flatten(), INT8_CLIP_Q).item()) if t32.numel() else 0.0
    scale = torch.tensor(clip_abs / 127.0 if clip_abs > 0 else 1.0, dtype=torch.float32)
    q = torch.clamp(torch.round(torch.clamp(t32, -clip_abs, clip_abs) / scale), -127, 127).to(torch.int8).contiguous()
    return q, scale


def quantize_state_dict_int8(state_dict: dict[str, Tensor]):
    quantized: dict[str, Tensor] = {}
    scales: dict[str, Tensor] = {}
    dtypes: dict[str, str] = {}
    passthrough: dict[str, Tensor] = {}
    passthrough_orig_dtypes: dict[str, str] = {}
    qmeta: dict[str, dict[str, object]] = {}
    stats = dict.fromkeys(
        ("param_count", "num_tensors", "num_float_tensors", "num_nonfloat_tensors", "baseline_tensor_bytes", "int8_payload_bytes"),
        0,
    )

    for name, tensor in state_dict.items():
        t = tensor.detach().to("cpu").contiguous()
        stats["param_count"] += int(t.numel())
        stats["num_tensors"] += 1
        stats["baseline_tensor_bytes"] += tensor_nbytes(t)

        if not t.is_floating_point():
            stats["num_nonfloat_tensors"] += 1
            passthrough[name] = t
            stats["int8_payload_bytes"] += tensor_nbytes(t)
            continue

        if t.numel() <= INT8_KEEP_FLOAT_MAX_NUMEL:
            kept = keep_float_tensor(name, t, passthrough_orig_dtypes)
            passthrough[name] = kept
            stats["int8_payload_bytes"] += tensor_nbytes(kept)
            continue

        stats["num_float_tensors"] += 1
        q, s = quantize_float_tensor(t)
        if s.ndim > 0:
            qmeta[name] = {"scheme": "per_row", "axis": 0}
        quantized[name] = q
        scales[name] = s
        dtypes[name] = str(t.dtype).removeprefix("torch.")
        stats["int8_payload_bytes"] += tensor_nbytes(q) + tensor_nbytes(s)

    obj: dict[str, object] = {
        "__quant_format__": "int8_clean_per_row_v1",
        "quantized": quantized,
        "scales": scales,
        "dtypes": dtypes,
        "passthrough": passthrough,
    }
    if qmeta:
        obj["qmeta"] = qmeta
    if passthrough_orig_dtypes:
        obj["passthrough_orig_dtypes"] = passthrough_orig_dtypes
    return obj, stats


def dequantize_state_dict_int8(obj: dict[str, object]) -> dict[str, Tensor]:
    out: dict[str, Tensor] = {}
    qmeta = obj.get("qmeta", {})
    passthrough_orig_dtypes = obj.get("passthrough_orig_dtypes", {})
    for name, q in obj["quantized"].items():
        dtype = getattr(torch, obj["dtypes"][name])
        s = obj["scales"][name]
        if qmeta.get(name, {}).get("scheme") == "per_row" or s.ndim > 0:
            s = s.to(dtype=torch.float32)
            out[name] = (q.float() * s.view(q.shape[0], *([1] * (q.ndim - 1)))).to(dtype=dtype).contiguous()
        else:
            scale = float(s.item())
            out[name] = (q.float() * scale).to(dtype=dtype).contiguous()
    for name, t in obj["passthrough"].items():
        out_t = t.detach().to("cpu").contiguous()
        orig_dtype = passthrough_orig_dtypes.get(name)
        if isinstance(orig_dtype, str):
            out_t = out_t.to(dtype=getattr(torch, orig_dtype)).contiguous()
        out[name] = out_t
    return out

# ─────────────────────────────────────────────────────────────────────────────
# DATA LOADING  (unchanged from baseline)
# ─────────────────────────────────────────────────────────────────────────────

def load_data_shard(file: Path) -> Tensor:
    header_bytes = 256 * np.dtype("<i4").itemsize
    token_bytes = np.dtype("<u2").itemsize
    header = np.fromfile(file, dtype="<i4", count=256)
    if header.size != 256 or int(header[0]) != 20240520 or int(header[1]) != 1:
        raise ValueError(f"Unexpected shard header for {file}")
    num_tokens = int(header[2])
    expected_size = header_bytes + num_tokens * token_bytes
    if file.stat().st_size != expected_size:
        raise ValueError(f"Shard size mismatch for {file}: expected {expected_size} bytes")
    tokens_np = np.fromfile(file, dtype="<u2", count=num_tokens, offset=header_bytes)
    if tokens_np.size != num_tokens:
        raise ValueError(f"Short read for {file}")
    return torch.from_numpy(tokens_np.astype(np.uint16, copy=False))


class TokenStream:
    def __init__(self, pattern: str):
        self.files = [Path(p) for p in sorted(glob.glob(pattern))]
        if not self.files:
            raise FileNotFoundError(f"No files found for pattern: {pattern}")
        self.file_idx = 0
        self.tokens = load_data_shard(self.files[0])
        self.pos = 0

    def _advance_file(self) -> None:
        self.file_idx = (self.file_idx + 1) % len(self.files)
        self.tokens = load_data_shard(self.files[self.file_idx])
        self.pos = 0

    def take(self, n: int) -> Tensor:
        chunks: list[Tensor] = []
        remaining = n
        while remaining > 0:
            avail = self.tokens.numel() - self.pos
            if avail <= 0:
                self._advance_file()
                continue
            k = min(remaining, avail)
            chunks.append(self.tokens[self.pos : self.pos + k])
            self.pos += k
            remaining -= k
        return chunks[0] if len(chunks) == 1 else torch.cat(chunks)


class DistributedTokenLoader:
    def __init__(self, pattern: str, rank: int, world_size: int, device: torch.device):
        self.rank = rank
        self.world_size = world_size
        self.device = device
        self.stream = TokenStream(pattern)

    def next_batch(self, global_tokens: int, seq_len: int, grad_accum_steps: int) -> tuple[Tensor, Tensor]:
        local_tokens = global_tokens // (self.world_size * grad_accum_steps)
        per_rank_span = local_tokens + 1
        chunk = self.stream.take(per_rank_span * self.world_size)
        start = self.rank * per_rank_span
        local = chunk[start : start + per_rank_span].to(dtype=torch.int64)
        x = local[:-1].reshape(-1, seq_len)
        y = local[1:].reshape(-1, seq_len)
        return x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True)

# ─────────────────────────────────────────────────────────────────────────────
# TRANSFORMER MODULES — KERNEL EDITION
# ─────────────────────────────────────────────────────────────────────────────

class RMSNorm(nn.Module):
    def __init__(self, eps: float | None = None):
        super().__init__()
        self.eps = eps

    def forward(self, x: Tensor) -> Tensor:
        return F.rms_norm(x, (x.size(-1),), eps=self.eps)


class CastedLinear(nn.Linear):
    def forward(self, x: Tensor) -> Tensor:
        bias = self.bias.to(x.dtype) if self.bias is not None else None
        return F.linear(x, self.weight.to(x.dtype), bias)


def restore_low_dim_params_to_fp32(module: nn.Module) -> None:
    with torch.no_grad():
        for name, param in module.named_parameters():
            if (param.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)) and param.dtype != torch.float32:
                param.data = param.data.float()


class Rotary(nn.Module):
    def __init__(self, dim: int, base: float = 10000.0):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self._seq_len_cached = 0
        self._cos_cached: Tensor | None = None
        self._sin_cached: Tensor | None = None

    def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tuple[Tensor, Tensor]:
        if (
            self._cos_cached is None
            or self._sin_cached is None
            or self._seq_len_cached != seq_len
            or self._cos_cached.device != device
        ):
            t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
            freqs = torch.outer(t, self.inv_freq.to(device))
            self._cos_cached = freqs.cos()[None, None, :, :]
            self._sin_cached = freqs.sin()[None, None, :, :]
            self._seq_len_cached = seq_len
        return self._cos_cached.to(dtype=dtype), self._sin_cached.to(dtype=dtype)


def apply_rotary_emb(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
    half = x.size(-1) // 2
    x1, x2 = x[..., :half], x[..., half:]
    return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1)


class CausalSelfAttention(nn.Module):
    def __init__(self, dim: int, num_heads: int, num_kv_heads: int, rope_base: float, qk_gain_init: float):
        super().__init__()
        if dim % num_heads != 0:
            raise ValueError("model_dim must be divisible by num_heads")
        if num_heads % num_kv_heads != 0:
            raise ValueError("num_heads must be divisible by num_kv_heads")
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = dim // num_heads
        if self.head_dim % 2 != 0:
            raise ValueError("head_dim must be even for RoPE")
        kv_dim = self.num_kv_heads * self.head_dim
        self.c_q = CastedLinear(dim, dim, bias=False)
        self.c_k = CastedLinear(dim, kv_dim, bias=False)
        self.c_v = CastedLinear(dim, kv_dim, bias=False)
        self.proj = CastedLinear(dim, dim, bias=False)
        self.proj._zero_init = True
        if _USE_LYAPUNOV_GAIN:
            # Lyapunov duality §10: C(exp λ) = sech λ — bounded gain in (0, 1].
            self.q_lambda = nn.Parameter(torch.zeros(num_heads, dtype=torch.float32))
        else:
            self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32))
        if _USE_MU_PHASE:
            # Z/8Z §15: bind head h to orbit slot h·MU_ANGLE via 2-D rotation.
            _h = torch.arange(num_heads, dtype=torch.float32) * MU_ANGLE
            _kv = torch.arange(num_kv_heads, dtype=torch.float32) * MU_ANGLE
            self.register_buffer("_mu_q_cos", torch.cos(_h).reshape(1, num_heads, 1, 1), persistent=False)
            self.register_buffer("_mu_q_sin", torch.sin(_h).reshape(1, num_heads, 1, 1), persistent=False)
            self.register_buffer("_mu_k_cos", torch.cos(_kv).reshape(1, num_kv_heads, 1, 1), persistent=False)
            self.register_buffer("_mu_k_sin", torch.sin(_kv).reshape(1, num_kv_heads, 1, 1), persistent=False)
            self.mu_phase_gate = nn.Parameter(torch.tensor(0.0))
            if _USE_PRECESSION:
                # Slow precession §14/§21: integer step stored to preserve float32 precision.
                self.register_buffer("_precession_step", torch.tensor(0, dtype=torch.int64), persistent=True)
        self.rotary = Rotary(self.head_dim, base=rope_base)

    def forward(self, x: Tensor) -> Tensor:
        bsz, seqlen, dim = x.shape
        q = self.c_q(x).reshape(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.c_k(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.c_v(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2)
        q = F.rms_norm(q, (q.size(-1),))
        k = F.rms_norm(k, (k.size(-1),))
        cos, sin = self.rotary(seqlen, x.device, q.dtype)
        q = apply_rotary_emb(q, cos, sin)
        k = apply_rotary_emb(k, cos, sin)
        if _USE_LYAPUNOV_GAIN:
            # sech(λ) = C(exp λ): bounded gain ∈ (0, 1] via Lyapunov duality §10.
            q = q * (1.0 / torch.cosh(self.q_lambda)).to(dtype=q.dtype)[None, :, None, None]
        else:
            q = q * self.q_gain.to(dtype=q.dtype)[None, :, None, None]
        if _USE_MU_PHASE:
            # Z/8Z §15: gated 2-D rotation per head; gate init=0 → identity by default.
            _g = self.mu_phase_gate.to(dtype=q.dtype)
            _half = self.head_dim // 2
            if _USE_PRECESSION:
                _ph = self._precession_step.float() * PRECESSION_DELTA_PHI
                _cp, _sp = torch.cos(_ph), torch.sin(_ph)
                _cq = (self._mu_q_cos * _cp - self._mu_q_sin * _sp).to(dtype=q.dtype)
                _sq = (self._mu_q_sin * _cp + self._mu_q_cos * _sp).to(dtype=q.dtype)
                _ck = (self._mu_k_cos * _cp - self._mu_k_sin * _sp).to(dtype=k.dtype)
                _sk = (self._mu_k_sin * _cp + self._mu_k_cos * _sp).to(dtype=k.dtype)
            else:
                _cq, _sq = self._mu_q_cos.to(dtype=q.dtype), self._mu_q_sin.to(dtype=q.dtype)
                _ck, _sk = self._mu_k_cos.to(dtype=k.dtype), self._mu_k_sin.to(dtype=k.dtype)
            q1, q2 = q[..., :_half], q[..., _half:]
            q = torch.lerp(q, torch.cat([q1 * _cq - q2 * _sq, q1 * _sq + q2 * _cq], -1), _g)
            k1, k2 = k[..., :_half], k[..., _half:]
            k = torch.lerp(k, torch.cat([k1 * _ck - k2 * _sk, k1 * _sk + k2 * _ck], -1), _g)
        # Expand k,v from num_kv_heads to num_heads for GQA
        k = k.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
        v = v.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
        y = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=None,
            is_causal=True,
        )
        y = y.transpose(1, 2).contiguous().reshape(bsz, seqlen, dim)
        return self.proj(y)


def coherence(x: Tensor) -> Tensor:
    """Coherence activation  C(r) = 2r / (1 + r²).

    Defined in CriticalEigenvalue.lean §5.  Machine-checked properties:
      • C(r) ≤ 1 for r ≥ 0  (AM–GM: 1 + r² ≥ 2r)
      • C(1) = 1             (unique maximum on ℝ⁺)
      • C(−r) = −C(r)        (odd / anti-symmetric)
      • Pythagorean: C(r)² + ((r²−1)/(1+r²))² = 1  (§18)
      • Lyapunov:   C(exp λ) = sech λ               (§10)

    As a neural activation this is a smooth, bounded nonlinearity with
    range (−1, 1), zero at the origin, unit gradient at zero, and graceful
    saturation for large |r| — resembling a normalised sinc or tanh but
    with a closed-form Pythagorean partner.
    """
    return 2.0 * x / (1.0 + x.square())


def _run_kernel_selftest() -> None:
    """Unit-test-like self-checks run when the env var KERNEL_SELFTEST=1.

    Exercises the coherence() tensor function and the global kernel constants
    at the PyTorch level, mirroring the scalar checks done at import time.
    All checks are cheap (CPU tensors, no CUDA required) and DDP-safe
    (no distributed calls; run on every rank independently).

    Invoke with:
        KERNEL_SELFTEST=1 python train_gpt_kernel.py
    """
    _eps = 1e-5  # tolerance for float32 tensor arithmetic

    # ── coherence() tensor invariants ─────────────────────────────────────────

    _xs = torch.linspace(-4.0, 4.0, 401)
    _cs = coherence(_xs)

    # Boundedness: |C(x)| < 1 strictly for x ≠ ±1  (coherence_le_one §5)
    assert (_cs.abs() <= 1.0 + _eps).all(), (
        f"coherence() out of [-1,1]: max |C|={_cs.abs().max().item():.6f}"
    )

    # Odd symmetry: C(-x) = -C(x)  (§5)
    _diff_odd = (coherence(-_xs) + _cs).abs().max().item()
    assert _diff_odd < _eps, (
        f"coherence() odd-symmetry violated: max |C(-x)+C(x)|={_diff_odd:.2e}"
    )

    # Unique maximum C(1) = 1  (coherence_eq_one_iff §5)
    _c1 = coherence(torch.tensor(1.0)).item()
    assert abs(_c1 - 1.0) < _eps, f"coherence(1) must equal 1; got {_c1:.8f}"

    # Pythagorean identity: C(r)² + ((r²-1)/(1+r²))² = 1  (§18)
    _r = torch.linspace(0.05, 5.0, 100)
    _c = coherence(_r)
    _s = (_r ** 2 - 1.0) / (1.0 + _r ** 2)
    _pyth_err = (_c ** 2 + _s ** 2 - 1.0).abs().max().item()
    assert _pyth_err < _eps, (
        f"Pythagorean identity violated in tensor: max err={_pyth_err:.2e}"
    )

    # Lyapunov duality: C(exp λ) = sech λ  (Theorem 14, §10)
    _lam = torch.linspace(-2.0, 2.0, 41)
    _lyap_err = (coherence(torch.exp(_lam)) - 1.0 / torch.cosh(_lam)).abs().max().item()
    assert _lyap_err < _eps, (
        f"Lyapunov duality violated in tensor: max err={_lyap_err:.2e}"
    )

    # Monotonicity on (0,1]: strictly increasing  (coherence_monotone §13)
    _mon_inc = coherence(torch.linspace(0.05, 1.0, 20))
    assert (_mon_inc[1:] > _mon_inc[:-1]).all(), (
        "coherence() must be strictly increasing on (0,1] (coherence_monotone §13)"
    )

    # Monotonicity on [1,∞): strictly decreasing  (coherence_monotone §13)
    _mon_dec = coherence(torch.linspace(1.0, 5.0, 20))
    assert (_mon_dec[1:] < _mon_dec[:-1]).all(), (
        "coherence() must be strictly decreasing on [1,∞) (coherence_monotone §13)"
    )

    # ── kernel constants ──────────────────────────────────────────────────────

    # 9 × D = 123456789 and gcd(8, D) = 1  (§14)
    assert 9 * PRECESSION_PERIOD == 123_456_789, (
        f"KERNEL_SELFTEST: 9*PRECESSION_PERIOD must be 123456789; "
        f"got {9 * PRECESSION_PERIOD}"
    )
    assert math.gcd(8, PRECESSION_PERIOD) == 1, (
        f"KERNEL_SELFTEST: gcd(8, PRECESSION_PERIOD) must be 1; "
        f"got {math.gcd(8, PRECESSION_PERIOD)}"
    )

    # Z/8Z orbit: num_heads default == MU_ORBIT_SIZE  (§15)
    _default_heads = int(os.environ.get("NUM_HEADS", MU_ORBIT_SIZE))
    assert _default_heads == MU_ORBIT_SIZE, (
        f"KERNEL_SELFTEST: default NUM_HEADS must equal MU_ORBIT_SIZE={MU_ORBIT_SIZE}; "
        f"got {_default_heads}"
    )

    # MLP hidden width ratio ≈ δS  (§7)
    _model_dim = int(os.environ.get("MODEL_DIM", 512))
    _mlp_hidden = round(SILVER_RATIO * _model_dim)
    _ratio = _mlp_hidden / _model_dim
    assert abs(_ratio - SILVER_RATIO) < 1.0 / _model_dim + 1e-9, (
        f"KERNEL_SELFTEST: mlp_hidden/model_dim must be ~δS={SILVER_RATIO:.6f}; "
        f"got {_ratio:.6f}"
    )

    # μ-phase rotation preserves pair norms per Z/8Z slot (§2, §15)
    _h_angles = torch.arange(MU_ORBIT_SIZE, dtype=torch.float32) * MU_ANGLE
    _vecs = torch.randn(MU_ORBIT_SIZE, 8)
    _half_d = 4
    _ch, _sh = torch.cos(_h_angles)[:, None], torch.sin(_h_angles)[:, None]
    _v1, _v2 = _vecs[:, :_half_d], _vecs[:, _half_d:]
    _rv1, _rv2 = _v1 * _ch - _v2 * _sh, _v1 * _sh + _v2 * _ch
    _norm_err = ((_v1 ** 2 + _v2 ** 2) - (_rv1 ** 2 + _rv2 ** 2)).abs().max().item()
    assert _norm_err < _eps, f"μ-phase rotation must preserve pair norms; max err={_norm_err:.2e}"

    print("kernel_selftest: all checks PASSED ✓")


if os.environ.get("KERNEL_SELFTEST", "0") == "1":
    _run_kernel_selftest()


class CoherenceMLP(nn.Module):
    """MLP block whose hidden width is ⌊δS·dim⌋ and activation is C(r).

    Both choices come directly from formal-lean/CriticalEigenvalue.lean:
      • δS = 1+√2 (silver ratio, §7 Proposition 4) determines hidden width.
      • C(r) = 2r/(1+r²) (coherence function, §5) is the nonlinearity.

    The silver ratio satisfies δS = 2 + 1/δS (self-similarity §20), so the
    expansion is slightly larger than 2× but smaller than 3×, giving a
    natural intermediate width that is mathematically well-motivated.
    """

    def __init__(self, dim: int, hidden: int):
        super().__init__()
        self.fc = CastedLinear(dim, hidden, bias=False)
        self.proj = CastedLinear(hidden, dim, bias=False)
        self.proj._zero_init = True

    def forward(self, x: Tensor) -> Tensor:
        return self.proj(coherence(self.fc(x)))


class KernelBlock(nn.Module):
    """Transformer block using CoherenceMLP in place of the baseline relu² MLP."""

    def __init__(self, dim: int, num_heads: int, num_kv_heads: int, mlp_hidden: int, rope_base: float, qk_gain_init: float):
        super().__init__()
        self.attn_norm = RMSNorm()
        self.mlp_norm = RMSNorm()
        self.attn = CausalSelfAttention(dim, num_heads, num_kv_heads, rope_base, qk_gain_init)
        self.mlp = CoherenceMLP(dim, mlp_hidden)
        self.attn_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32))
        self.mlp_scale = nn.Parameter(torch.ones(dim, dtype=torch.float32))
        self.resid_mix = nn.Parameter(torch.stack((torch.ones(dim), torch.zeros(dim))).float())

    def forward(self, x: Tensor, x0: Tensor) -> Tensor:
        mix = self.resid_mix.to(dtype=x.dtype)
        x = mix[0][None, None, :] * x + mix[1][None, None, :] * x0
        attn_out = self.attn(self.attn_norm(x))
        x = x + self.attn_scale.to(dtype=x.dtype)[None, None, :] * attn_out
        x = x + self.mlp_scale.to(dtype=x.dtype)[None, None, :] * self.mlp(self.mlp_norm(x))
        return x


class KernelGPT(nn.Module):
    """GPT variant implementing the Kernel eigenvalue formalization.

    Architectural choices derived from CriticalEigenvalue.lean:
      • num_heads = 8  (one slot per Z/8Z orbit element, §15)
      • MLP hidden  = ⌊δS · model_dim⌋  (silver ratio expansion, §7)
      • MLP activation = C(r) = 2r/(1+r²)  (coherence function, §5)
    """

    def __init__(
        self,
        vocab_size: int,
        num_layers: int,
        model_dim: int,
        num_heads: int,
        num_kv_heads: int,
        mlp_hidden: int,
        tie_embeddings: bool,
        tied_embed_init_std: float,
        logit_softcap: float,
        rope_base: float,
        qk_gain_init: float,
    ):
        super().__init__()
        if logit_softcap <= 0.0:
            raise ValueError(f"logit_softcap must be positive, got {logit_softcap}")
        self.tie_embeddings = tie_embeddings
        self.tied_embed_init_std = tied_embed_init_std
        self.logit_softcap = logit_softcap
        self.tok_emb = nn.Embedding(vocab_size, model_dim)
        self.num_encoder_layers = num_layers // 2
        self.num_decoder_layers = num_layers - self.num_encoder_layers
        self.num_skip_weights = min(self.num_encoder_layers, self.num_decoder_layers)
        self.skip_weights = nn.Parameter(torch.ones(self.num_skip_weights, model_dim, dtype=torch.float32))
        self.blocks = nn.ModuleList(
            [
                KernelBlock(model_dim, num_heads, num_kv_heads, mlp_hidden, rope_base, qk_gain_init)
                for _ in range(num_layers)
            ]
        )
        self.final_norm = RMSNorm()
        self.lm_head = None if tie_embeddings else CastedLinear(model_dim, vocab_size, bias=False)
        if self.lm_head is not None:
            self.lm_head._zero_init = True
        self._init_weights()

    def _init_weights(self) -> None:
        if self.tie_embeddings:
            nn.init.normal_(self.tok_emb.weight, mean=0.0, std=self.tied_embed_init_std)
        for module in self.modules():
            if isinstance(module, nn.Linear) and getattr(module, "_zero_init", False):
                nn.init.zeros_(module.weight)

    def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor:
        x = self.tok_emb(input_ids)
        x = F.rms_norm(x, (x.size(-1),))
        x0 = x
        skips: list[Tensor] = []

        for i in range(self.num_encoder_layers):
            x = self.blocks[i](x, x0)
            skips.append(x)
        for i in range(self.num_decoder_layers):
            if skips:
                x = x + self.skip_weights[i].to(dtype=x.dtype)[None, None, :] * skips.pop()
            x = self.blocks[self.num_encoder_layers + i](x, x0)

        x = self.final_norm(x).reshape(-1, x.size(-1))
        targets = target_ids.reshape(-1)
        if self.tie_embeddings:
            logits_proj = F.linear(x, self.tok_emb.weight)
        else:
            if self.lm_head is None:
                raise RuntimeError("lm_head is required when tie_embeddings=False")
            logits_proj = self.lm_head(x)
        logits = self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap)
        return F.cross_entropy(logits.float(), targets, reduction="mean")

# ─────────────────────────────────────────────────────────────────────────────
# TRAINING
# ─────────────────────────────────────────────────────────────────────────────

def main() -> None:
    global zeropower_via_newtonschulz5

    code = Path(__file__).read_text(encoding="utf-8")
    args = Hyperparameters()
    zeropower_via_newtonschulz5 = torch.compile(zeropower_via_newtonschulz5)

    # ── Distributed + CUDA setup ──────────────────────────────────────────────

    distributed = "RANK" in os.environ and "WORLD_SIZE" in os.environ
    rank = int(os.environ.get("RANK", "0"))
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    if world_size <= 0:
        raise ValueError(f"WORLD_SIZE must be positive, got {world_size}")
    if 8 % world_size != 0:
        raise ValueError(f"WORLD_SIZE={world_size} must divide 8 so grad_accum_steps stays integral")
    grad_accum_steps = 8 // world_size
    grad_scale = 1.0 / grad_accum_steps
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required")
    device = torch.device("cuda", local_rank)
    torch.cuda.set_device(device)
    if distributed:
        dist.init_process_group(backend="nccl", device_id=device)
        dist.barrier()
    master_process = rank == 0

    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    from torch.backends.cuda import enable_cudnn_sdp, enable_flash_sdp, enable_math_sdp, enable_mem_efficient_sdp
    # Flash SDP is significantly faster than the math fallback for seq_len >= 512.
    # Default train_seq_len=1024 in this repo (Hyperparameters.train_seq_len).
    # cudnn/mem-efficient SDP are off; flash is the sole back-end.
    enable_cudnn_sdp(False)
    enable_flash_sdp(True)
    enable_mem_efficient_sdp(False)
    enable_math_sdp(False)

    logfile = None
    if master_process:
        os.makedirs("logs", exist_ok=True)
        logfile = f"logs/{args.run_id}.txt"
        print(logfile)

    def log0(msg: str, console: bool = True) -> None:
        if not master_process:
            return
        if console:
            print(msg)
        if logfile is not None:
            with open(logfile, "a", encoding="utf-8") as f:
                print(msg, file=f)

    log0(code, console=False)
    log0("=" * 100, console=False)
    log0(f"Running Python {sys.version}", console=False)
    log0(f"Running PyTorch {torch.__version__}", console=False)
    log0(
        subprocess.run(["nvidia-smi"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=False).stdout,
        console=False,
    )
    log0("=" * 100, console=False)

    # ── Kernel constants log ──────────────────────────────────────────────────
    log0(f"kernel:silver_ratio:{SILVER_RATIO:.6f} mu_angle_deg:{math.degrees(MU_ANGLE):.1f} orbit_size:{MU_ORBIT_SIZE}")
    log0(f"kernel:mlp_hidden:{args.mlp_hidden} (={args.mlp_hidden / args.model_dim:.4f}x model_dim)")
    log0(f"kernel:precession_period:{PRECESSION_PERIOD} delta_phi:{PRECESSION_DELTA_PHI:.2e} (formal-lean/CriticalEigenvalue.lean §14, §21)")
    log0(
        f"kernel:invariants: 9*D={9 * PRECESSION_PERIOD} (expect 123456789) "
        f"gcd(8,D)={math.gcd(8, PRECESSION_PERIOD)} (expect 1) "
        f"D*dPhi/2pi={PRECESSION_PERIOD * PRECESSION_DELTA_PHI / (2.0 * math.pi):.10f} (expect 1.0)"
    )
    if _USE_LYAPUNOV_GAIN or _USE_MU_PHASE:
        log0(f"kernel:features lyapunov_gain:{int(_USE_LYAPUNOV_GAIN)} mu_phase:{int(_USE_MU_PHASE)} precession:{int(_USE_PRECESSION)}")

    # ── Seed + tokenizer ──────────────────────────────────────────────────────

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    if not args.tokenizer_path.endswith(".model"):
        raise ValueError(f"Script only setup for SentencePiece .model file: {args.tokenizer_path}")
    sp = spm.SentencePieceProcessor(model_file=args.tokenizer_path)
    if int(sp.vocab_size()) != args.vocab_size:
        raise ValueError(
            f"VOCAB_SIZE={args.vocab_size} does not match tokenizer vocab_size={int(sp.vocab_size())}"
        )
    dataset_dir = Path(args.data_path).resolve()
    actual_train_files = len(list(dataset_dir.glob("fineweb_train_*.bin")))
    val_tokens = load_validation_tokens(args.val_files, args.train_seq_len)
    base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = build_sentencepiece_luts(
        sp, args.vocab_size, device
    )
    log0(f"val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path={args.tokenizer_path}")
    log0(f"train_loader:dataset:{dataset_dir.name} train_shards:{actual_train_files}")
    log0(f"val_loader:shards pattern={args.val_files} tokens:{val_tokens.numel() - 1}")

    # ── Model + optimizer setup ───────────────────────────────────────────────

    base_model = KernelGPT(
        vocab_size=args.vocab_size,
        num_layers=args.num_layers,
        model_dim=args.model_dim,
        num_heads=args.num_heads,
        num_kv_heads=args.num_kv_heads,
        mlp_hidden=args.mlp_hidden,
        tie_embeddings=args.tie_embeddings,
        tied_embed_init_std=args.tied_embed_init_std,
        logit_softcap=args.logit_softcap,
        rope_base=args.rope_base,
        qk_gain_init=args.qk_gain_init,
    ).to(device).bfloat16()
    for module in base_model.modules():
        if isinstance(module, CastedLinear):
            module.float()
    restore_low_dim_params_to_fp32(base_model)
    compiled_model = torch.compile(base_model, dynamic=False, fullgraph=True)
    model: nn.Module = DDP(compiled_model, device_ids=[local_rank], broadcast_buffers=False) if distributed else compiled_model

    block_named_params = list(base_model.blocks.named_parameters())
    matrix_params = [
        p
        for name, p in block_named_params
        if p.ndim == 2 and not any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)
    ]
    scalar_params = [
        p
        for name, p in block_named_params
        if p.ndim < 2 or any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS)
    ]
    if base_model.skip_weights.numel() > 0:
        scalar_params.append(base_model.skip_weights)
    token_lr = args.tied_embed_lr if args.tie_embeddings else args.embed_lr
    optimizer_tok = torch.optim.Adam(
        [{"params": [base_model.tok_emb.weight], "lr": token_lr, "base_lr": token_lr}],
        betas=(args.beta1, args.beta2), eps=args.adam_eps, fused=True,
    )
    optimizer_muon = Muon(matrix_params, lr=args.matrix_lr, momentum=args.muon_momentum, backend_steps=args.muon_backend_steps)
    for group in optimizer_muon.param_groups:
        group["base_lr"] = args.matrix_lr
    optimizer_scalar = torch.optim.Adam(
        [{"params": scalar_params, "lr": args.scalar_lr, "base_lr": args.scalar_lr}],
        betas=(args.beta1, args.beta2), eps=args.adam_eps, fused=True,
    )
    optimizers: list[torch.optim.Optimizer] = [optimizer_tok, optimizer_muon, optimizer_scalar]
    if base_model.lm_head is not None:
        optimizer_head = torch.optim.Adam(
            [{"params": [base_model.lm_head.weight], "lr": args.head_lr, "base_lr": args.head_lr}],
            betas=(args.beta1, args.beta2), eps=args.adam_eps, fused=True,
        )
        optimizers.insert(1, optimizer_head)

    n_params = sum(p.numel() for p in base_model.parameters())
    log0(f"model_params:{n_params}")
    log0(f"world_size:{world_size} grad_accum_steps:{grad_accum_steps}")
    log0("sdp_backends:cudnn=False flash=True mem_efficient=False math=False")
    log0(f"attention_mode:gqa num_heads:{args.num_heads} num_kv_heads:{args.num_kv_heads}")
    log0(
        f"tie_embeddings:{args.tie_embeddings} embed_lr:{token_lr} "
        f"head_lr:{args.head_lr if base_model.lm_head is not None else 0.0} "
        f"matrix_lr:{args.matrix_lr} scalar_lr:{args.scalar_lr}"
    )
    log0(
        f"train_batch_tokens:{args.train_batch_tokens} train_seq_len:{args.train_seq_len} "
        f"iterations:{args.iterations} warmup_steps:{args.warmup_steps} "
        f"max_wallclock_seconds:{args.max_wallclock_seconds:.3f}"
    )
    log0(f"seed:{args.seed}")

    # ── Data loader & model warmup ─────────────────────────────────────────────

    train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device)

    def zero_grad_all() -> None:
        for opt in optimizers:
            opt.zero_grad(set_to_none=True)

    max_wallclock_ms = 1000.0 * args.max_wallclock_seconds if args.max_wallclock_seconds > 0 else None

    def lr_mul(step: int, elapsed_ms: float) -> float:
        """LR multiplier using palindrome precession schedule (CriticalEigenvalue.lean §9, §16).

        Replaces the previous linear warmdown with the coherence function
        C(r) = 2r/(1+r²) evaluated at r = remaining/warmdown ∈ [0,1].
        This is strictly smoother than a linear ramp: C has zero derivative at
        r=0, so the LR tapers gracefully rather than hitting a hard zero.
        """
        if args.warmdown_iters <= 0:
            return 1.0
        if max_wallclock_ms is None:
            warmdown_start = max(args.iterations - args.warmdown_iters, 0)
            if step < warmdown_start:
                return 1.0
            r = (args.iterations - step) / max(args.warmdown_iters, 1)
        else:
            step_ms = elapsed_ms / max(step, 1)
            warmdown_ms = args.warmdown_iters * step_ms
            remaining_ms = max(max_wallclock_ms - elapsed_ms, 0.0)
            if remaining_ms > warmdown_ms:
                return 1.0
            r = remaining_ms / max(warmdown_ms, 1e-9)
        return palindrome_precession_scale(r)

    if args.warmup_steps > 0:
        initial_model_state = {name: tensor.detach().cpu().clone() for name, tensor in base_model.state_dict().items()}
        initial_optimizer_states = [copy.deepcopy(opt.state_dict()) for opt in optimizers]
        model.train()
        for warmup_step in range(args.warmup_steps):
            zero_grad_all()
            for micro_step in range(grad_accum_steps):
                if distributed:
                    model.require_backward_grad_sync = micro_step == grad_accum_steps - 1
                x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps)
                with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
                    warmup_loss = model(x, y)
                (warmup_loss * grad_scale).backward()
            for opt in optimizers:
                opt.step()
            zero_grad_all()
            if args.warmup_steps <= 20 or (warmup_step + 1) % 10 == 0 or warmup_step + 1 == args.warmup_steps:
                log0(f"warmup_step:{warmup_step + 1}/{args.warmup_steps}")
        base_model.load_state_dict(initial_model_state, strict=True)
        for opt, state in zip(optimizers, initial_optimizer_states, strict=True):
            opt.load_state_dict(state)
        zero_grad_all()
        if distributed:
            model.require_backward_grad_sync = True
        train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device)

    # ── Main training loop ────────────────────────────────────────────────────

    training_time_ms = 0.0
    stop_after_step: int | None = None
    last_step_ms: float = 0.0  # duration of the most recently completed step (ms)
    torch.cuda.synchronize()
    t0 = time.perf_counter()

    step = 0
    while True:
        last_step = step == args.iterations or (stop_after_step is not None and step >= stop_after_step)

        should_validate = last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0)
        if should_validate:
            torch.cuda.synchronize()
            training_time_ms += 1000.0 * (time.perf_counter() - t0)
            val_loss, val_bpb = eval_val(
                args, model, rank, world_size, device, grad_accum_steps,
                val_tokens, base_bytes_lut, has_leading_space_lut, is_boundary_token_lut,
            )
            log0(
                f"step:{step}/{args.iterations} val_loss:{val_loss:.4f} val_bpb:{val_bpb:.4f} "
                f"train_time:{training_time_ms:.0f}ms step_avg:{training_time_ms / max(step, 1):.2f}ms"
            )
            torch.cuda.synchronize()
            t0 = time.perf_counter()

        if last_step:
            if stop_after_step is not None and step < args.iterations:
                log0(
                    f"stopping_early: wallclock_cap train_time:{training_time_ms:.0f}ms "
                    f"step:{step}/{args.iterations}"
                )
            break

        step_t0 = time.perf_counter()
        elapsed_ms = training_time_ms + 1000.0 * (step_t0 - t0)
        scale = lr_mul(step, elapsed_ms)
        zero_grad_all()
        train_loss = torch.zeros((), device=device)
        for micro_step in range(grad_accum_steps):
            if distributed:
                model.require_backward_grad_sync = micro_step == grad_accum_steps - 1
            x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len, grad_accum_steps)
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
                loss = model(x, y)
            train_loss += loss.detach()
            (loss * grad_scale).backward()
        train_loss /= grad_accum_steps

        frac = min(step / args.muon_momentum_warmup_steps, 1.0) if args.muon_momentum_warmup_steps > 0 else 1.0
        muon_momentum = (1 - frac) * args.muon_momentum_warmup_start + frac * args.muon_momentum
        for group in optimizer_muon.param_groups:
            group["momentum"] = muon_momentum

        for opt in optimizers:
            for group in opt.param_groups:
                group["lr"] = group["base_lr"] * scale

        if args.grad_clip_norm > 0:
            torch.nn.utils.clip_grad_norm_(base_model.parameters(), args.grad_clip_norm)
        for opt in optimizers:
            opt.step()
        zero_grad_all()
        if _USE_MU_PHASE and _USE_PRECESSION:
            with torch.no_grad():
                for _blk in base_model.blocks:
                    _blk.attn._precession_step.add_(1)

        step += 1
        last_step_ms = 1000.0 * (time.perf_counter() - step_t0)
        approx_training_time_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0)
        should_log_train = (
            args.train_log_every > 0
            and (step <= 10 or step % args.train_log_every == 0 or stop_after_step is not None)
        )
        if should_log_train:
            log0(
                f"step:{step}/{args.iterations} train_loss:{train_loss.item():.4f} "
                f"step_t:{last_step_ms:.1f}ms "
                f"train_time:{approx_training_time_ms:.0f}ms step_avg:{approx_training_time_ms / step:.2f}ms"
            )

        reached_cap = max_wallclock_ms is not None and approx_training_time_ms >= max_wallclock_ms
        if distributed and max_wallclock_ms is not None:
            reached_cap_tensor = torch.tensor(int(reached_cap), device=device)
            dist.all_reduce(reached_cap_tensor, op=dist.ReduceOp.MAX)
            reached_cap = bool(reached_cap_tensor.item())
        if stop_after_step is None and reached_cap:
            stop_after_step = step

    log0(
        f"peak memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB "
        f"reserved: {torch.cuda.max_memory_reserved() // 1024 // 1024} MiB"
    )

    # ── Serialization + round-trip validation ─────────────────────────────────

    if master_process:
        torch.save(base_model.state_dict(), "final_model.pt")
        model_bytes = os.path.getsize("final_model.pt")
        code_bytes = len(code.encode("utf-8"))
        log0(f"Serialized model: {model_bytes} bytes")
        log0(f"Code size: {code_bytes} bytes")
        log0(f"Total submission size: {model_bytes + code_bytes} bytes")

    quant_obj, quant_stats = quantize_state_dict_int8(base_model.state_dict())
    quant_buf = io.BytesIO()
    torch.save(quant_obj, quant_buf)
    quant_raw = quant_buf.getvalue()
    quant_blob = zlib.compress(quant_raw, level=9)
    quant_raw_bytes = len(quant_raw)
    if master_process:
        with open("final_model.int8.ptz", "wb") as f:
            f.write(quant_blob)
        quant_file_bytes = os.path.getsize("final_model.int8.ptz")
        code_bytes = len(code.encode("utf-8"))
        ratio = quant_stats["baseline_tensor_bytes"] / max(quant_stats["int8_payload_bytes"], 1)
        log0(
            f"Serialized model int8+zlib: {quant_file_bytes} bytes "
            f"(payload:{quant_stats['int8_payload_bytes']} raw_torch:{quant_raw_bytes} payload_ratio:{ratio:.2f}x)"
        )
        log0(f"Total submission size int8+zlib: {quant_file_bytes + code_bytes} bytes")

    if distributed:
        dist.barrier()
    with open("final_model.int8.ptz", "rb") as f:
        quant_blob_disk = f.read()
    quant_state = torch.load(io.BytesIO(zlib.decompress(quant_blob_disk)), map_location="cpu")
    base_model.load_state_dict(dequantize_state_dict_int8(quant_state), strict=True)
    torch.cuda.synchronize()
    t_qeval = time.perf_counter()
    q_val_loss, q_val_bpb = eval_val(
        args, model, rank, world_size, device, grad_accum_steps,
        val_tokens, base_bytes_lut, has_leading_space_lut, is_boundary_token_lut,
    )
    torch.cuda.synchronize()
    log0(
        f"final_int8_zlib_roundtrip val_loss:{q_val_loss:.4f} val_bpb:{q_val_bpb:.4f} "
        f"eval_time:{1000.0 * (time.perf_counter() - t_qeval):.0f}ms"
    )
    log0(f"final_int8_zlib_roundtrip_exact val_loss:{q_val_loss:.8f} val_bpb:{q_val_bpb:.8f}")

    if distributed:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()