train_jepa_ssm.py

"""
To JEPA Or Not To JEPA, That Is Le Question
=============================================
Ciprian-Florin Ifrim - 26 March 2026

Architecture: Tokens → Embedding → [Mamba-2 SSM + ReLU²/SwiGLU MLP] × L (U-Net skips) → RMSNorm → h
    JEPA branch:  h → Projector → z → Predictor (autoregressive rollout) → MSE loss
    Decode branch: h → [optional GELU adapter] → tied/untied lm_head → logits → CE loss
    SIGReg:       z → per-timestep Gaussian regularization (Epps-Pulley characteristic function test)
Training:     JEPA (MSE latent prediction) + SIGReg + CE (token decoding) + Z-loss
Evaluation:   lm_head(h) → token logits → BPB

LeWM-style JEPA (Maes et al. 2026) + Mamba-2 SSM (Gu & Dao 2024) + byte-level I/O (Meta's BLT)
- The model learns to predict its own next latent state.
- SIGReg prevents representation collapse without EMA or stop-gradient.
- CE + lm_head decode hidden states to token probabilities.
- Supports byte (256 vocab) and BPE (8K+ vocab) tokenization.
"""
import copy
import glob
import io
import lzma
import math
import os
import random
import sys
import time
from pathlib import Path
import numpy as np
try:
    import sentencepiece as spm
except ImportError:
    spm = None  # byte mode doesn't need sentencepiece
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
# ---------------------------------------------------------------------------
# Mamba-2 (requires mamba_ssm package with CUDA kernels)
# ---------------------------------------------------------------------------
from mamba_ssm import Mamba2
# ---------------------------------------------------------------------------
# Hyperparameters
# ---------------------------------------------------------------------------
def _e(k, d, t=str):
    v = os.environ.get(k, str(d))
    if t == bool:
        return bool(int(v))
    return t(v)

class Hyperparameters:
    # Data & Tokenizer
    tokenizer      = _e("TOKENIZER", "byte")         # "byte" (256 vocab, no tokenizer) or "bpe"
    tokenizer_path = _e("TOKENIZER_PATH", "")         # path to .model file (BPE only)
    vocab_size     = _e("VOCAB_SIZE", 256, int)        # 256 for byte, 8192 for BPE, etc.
    data_path      = _e("DATA_PATH", "./data/datasets/fineweb_bytes")
    train_files    = os.path.join(data_path, "fineweb_train_*.bin")
    val_files      = os.path.join(data_path, "fineweb_val_*.bin")
    run_id         = os.environ.get("RUN_ID", f"jepa_{int(time.time())}")
    seed           = _e("SEED", 42, int)
    # Model
    model_dim      = _e("MODEL_DIM", 512, int)
    num_layers     = _e("NUM_LAYERS", 12, int)
    d_state        = _e("D_STATE", 64, int)
    d_conv         = _e("D_CONV", 4, int)
    expand         = _e("EXPAND", 2, int)
    predictor_hidden_mult = _e("PREDICTOR_HIDDEN_MULT", 2, int)
    embed_dim      = _e("EMBED_DIM", 0, int)       # 0 = same as model_dim
    mlp_mult       = _e("MLP_MULT", 3, int)        # 0 = no MLP, >0 = hidden = mlp_mult * dim
    mlp_every      = _e("MLP_EVERY", 1, int)        # 1 = every block, 2 = alternate, etc.
    activation     = _e("ACTIVATION", "relu2")      # "relu2" or "swiglu"
    tie_embeddings = _e("TIE_EMBEDDINGS", 0, int)   # 0=separate, 1=tied, 2=tied+correction, 3=tied+adapter
    projector_type = _e("PROJECTOR_TYPE", "linear")  # "linear" or "mlp"
    logit_softcap  = _e("LOGIT_SOFTCAP", 15.0, float)  # 0 = disabled
    softcap_type   = _e("SOFTCAP_TYPE", "poly")     # "poly" or "tanh"
    # JEPA / Loss
    sigreg_lambda  = _e("SIGREG_LAMBDA", 1.0, float)   # was 0.01 when loss was scaled by N
    sigreg_knots   = _e("SIGREG_KNOTS", 17, int)
    sigreg_num_proj = _e("SIGREG_NUM_PROJ", 0, int)  # 0 = auto (2x embed_dim)
    sigreg_schedule = _e("SIGREG_SCHEDULE", 0, bool)  # ramp lambda from 0 to sigreg_lambda
    ce_weight      = _e("CE_WEIGHT", 1.0, float)
    jepa_weight    = _e("JEPA_WEIGHT", 1.0, float)
    jepa_steps     = _e("JEPA_STEPS", 1, int)       # 0=disabled, 1=next step, 2+=multi-step
    detach_targets = _e("DETACH_TARGETS", 0, bool)
    # Training
    train_seq_len  = _e("TRAIN_SEQ_LEN", 4096, int)
    train_batch_tokens = _e("TRAIN_BATCH_TOKENS", 524288, int)
    iterations     = _e("ITERATIONS", 10000, int)
    warmup_steps   = _e("WARMUP_STEPS", 10, int)
    warmdown_fraction = _e("WARMDOWN_FRACTION", 0.15, float)
    max_wallclock_seconds = _e("MAX_WALLCLOCK_SECONDS", 0.0, float)
    val_batch_size = _e("VAL_BATCH_SIZE", 524288, int)
    val_loss_every = _e("VAL_LOSS_EVERY", 500, int)
    train_log_every = _e("TRAIN_LOG_EVERY", 100, int)
    # Optimizer
    matrix_lr      = _e("MATRIX_LR", 0.02, float)
    scalar_lr      = _e("SCALAR_LR", 0.01, float)
    embed_lr       = _e("EMBED_LR", 0.01, float)
    muon_momentum  = _e("MUON_MOMENTUM", 0.95, float)
    muon_backend_steps = _e("MUON_BACKEND_STEPS", 3, int)
    muon_momentum_warmup_start = _e("MUON_MOMENTUM_WARMUP_START", 0.85, float)
    muon_momentum_warmup_steps = _e("MUON_MOMENTUM_WARMUP_STEPS", 500, int)
    muon_wd        = _e("MUON_WD", 0.04, float)
    adam_wd        = _e("ADAM_WD", 0.05, float)
    beta1          = _e("BETA1", 0.9, float)
    beta2          = _e("BETA2", 0.95, float)
    adam_eps        = _e("ADAM_EPS", 1e-8, float)
    grad_clip_norm = _e("GRAD_CLIP_NORM", 1.0, float)
    # Quantization
    quant_bits     = _e("QUANT_BITS", 5, int)       # INT-N for large matrices (4, 5, 6)
    qat_fraction   = _e("QAT_FRACTION", 1.0, float) # 1.0 = QAT from step 1
    _fp_raw        = os.environ.get("FP_STORAGE", "BF16")
    fp_storage     = "fp8" if _fp_raw == "FP8" else "bf16"  # FP8 or BF16 for non-INT params
    # Eval
    sliding_eval   = _e("SLIDING_EVAL", 0, bool)
    sliding_eval_stride = _e("SLIDING_EVAL_STRIDE", 16, int)
    sliding_batch_size = _e("SLIDING_BATCH_SIZE", 64, int)
    temp_scaling   = _e("TEMP_SCALING", 1, bool)
    compile_mode   = _e("COMPILE_MODE", "default")
    # Checkpoint
    checkpoint_every = _e("CHECKPOINT_EVERY", 0, int)   # 0 = disabled
    checkpoint_dir = _e("CHECKPOINT_DIR", "./checkpoints")

# ---------------------------------------------------------------------------
# SIGReg — Sketch Isotropic Gaussian Regularizer (from LeWorldModel)
# Prevents representation collapse by enforcing Gaussian latent distribution.
# Maes, Le Lidec, Scieur, LeCun, Balestriero (2026)
# ---------------------------------------------------------------------------
class SIGReg(nn.Module):
    def __init__(self, d_model: int, knots: int = 17, num_proj: int = 1024):
        super().__init__()
        self.num_proj = num_proj
        # Integration range [0.2, 4] matching LeWorldModel paper (Appendix A)
        t = torch.linspace(0.2, 4.0, knots, dtype=torch.float32)
        dt = (4.0 - 0.2) / (knots - 1)
        weights = torch.full((knots,), 2 * dt, dtype=torch.float32)
        weights[[0, -1]] = dt
        window = torch.exp(-t.square() / 2.0)
        self.register_buffer("t", t)
        self.register_buffer("phi", window)             # Gaussian characteristic function
        self.register_buffer("weights", weights * window)
        # Fixed random projection matrix — generated once, avoids CUDA RNG overhead per step.
        A = torch.randn(d_model, num_proj)
        A = A / A.norm(dim=0, keepdim=True)
        self.register_buffer("A", A)

    def _sigreg_1d(self, emb: Tensor) -> Tensor:
        """Compute SIGReg loss for a single (N, D) batch of embeddings."""
        proj = emb @ self.A.to(dtype=emb.dtype)           # (N, num_proj)
        x_t = proj.unsqueeze(-1) * self.t                 # (N, num_proj, knots)
        err = (x_t.cos().mean(0) - self.phi).square() + x_t.sin().mean(0).square()
        return (err @ self.weights).mean()

    def forward(self, embeddings: Tensor, max_timesteps: int = 64) -> Tensor:
        """
        embeddings: (B, T, D) — apply SIGReg per-timestep across batch (matching paper).
        Each timestep's batch of embeddings is independently pushed toward Gaussian.
        Samples up to max_timesteps positions to keep compute bounded.
        Returns: scalar regularization loss (mean over sampled timesteps).
        """
        if embeddings.ndim == 2:
            return self._sigreg_1d(embeddings)
        B, T, D = embeddings.shape
        # Sample timesteps if T is large
        if T > max_timesteps:
            t_idx = torch.randint(0, T, (max_timesteps,), device=embeddings.device)
            emb = embeddings[:, t_idx, :]         # (B, S, D) where S = max_timesteps
        else:
            emb = embeddings                       # (B, T, D)
        S = emb.size(1)
        # Vectorized per-timestep SIGReg:
        # Project all timesteps at once: (B, S, D) @ (D, num_proj) -> (B, S, num_proj)
        A = self.A.to(dtype=emb.dtype)
        proj = torch.einsum("bsd,dp->bsp", emb, A)   # (B, S, num_proj)
        x_t = proj.unsqueeze(-1) * self.t             # (B, S, num_proj, knots)
        # Per-timestep mean over batch dim: mean(0) -> (S, num_proj, knots)
        cos_mean = x_t.cos().mean(0)
        sin_mean = x_t.sin().mean(0)
        # Characteristic function error per timestep
        err = (cos_mean - self.phi).square() + sin_mean.square()  # (S, num_proj, knots)
        per_t = (err @ self.weights).mean(-1)          # (S,) — loss per timestep
        return per_t.mean()

# ---------------------------------------------------------------------------
# Model components
# ---------------------------------------------------------------------------
class RMSNorm(nn.Module):
    def __init__(self, eps: float = 1e-6):
        super().__init__()
        self.eps = eps

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


class SwiGLU(nn.Module):
    """SwiGLU MLP: gate-up projection -> SiLU gate -> norm -> down projection."""
    def __init__(self, dim: int, hidden_mult: int = 2):
        super().__init__()
        hidden = dim * hidden_mult
        self.gate_up = nn.Linear(dim, hidden * 2, bias=False)
        self.norm = RMSNorm()
        self.down = nn.Linear(hidden, dim, bias=False)
        nn.init.zeros_(self.down.weight)

    def forward(self, x: Tensor) -> Tensor:
        gu = self.gate_up(x)
        gate, up = gu.chunk(2, dim=-1)
        return self.down(self.norm(F.silu(gate) * up))


class ReLU2MLP(nn.Module):
    """ReLU^2 MLP: fc -> relu^2 -> norm -> down. Fewer params than SwiGLU (no gate)."""
    def __init__(self, dim: int, hidden_mult: int = 3):
        super().__init__()
        hidden = dim * hidden_mult
        self.fc = nn.Linear(dim, hidden, bias=False)
        self.norm = RMSNorm()
        self.down = nn.Linear(hidden, dim, bias=False)
        nn.init.zeros_(self.down.weight)

    def forward(self, x: Tensor) -> Tensor:
        return self.down(self.norm(F.relu(self.fc(x)).square()))


class SSMBlock(nn.Module):
    """SSM block with pre-RMSNorm, residual scales, x0 mixing, and optional MLP."""
    def __init__(self, d_model: int, d_state: int = 64, d_conv: int = 4,
                 expand: int = 2, mlp_mult: int = 0, activation: str = "relu2"):
        super().__init__()
        self.norm = RMSNorm()
        self.ssm = Mamba2(d_model=d_model, d_state=d_state, d_conv=d_conv, expand=expand)
        # Learnable per-dim residual scales (like ternary code's attn_scale / mlp_scale)
        self.ssm_scale = nn.Parameter(torch.ones(d_model, dtype=torch.float32))
        # x0 residual mixing: mix[0]*x + mix[1]*x0 — gives every block direct
        # access to the original input representation (gradient highway)
        self.resid_mix = nn.Parameter(torch.stack((torch.ones(d_model), torch.zeros(d_model))).float())
        self.mlp = None
        if mlp_mult > 0:
            self.mlp_norm = RMSNorm()
            self.mlp_scale = nn.Parameter(torch.ones(d_model, dtype=torch.float32))
            if activation == "swiglu":
                self.mlp = SwiGLU(d_model, mlp_mult)
            else:  # relu2
                self.mlp = ReLU2MLP(d_model, mlp_mult)

    def forward(self, x: Tensor, x0: Tensor) -> Tensor:
        mix = self.resid_mix.to(dtype=x.dtype)
        x = mix[0] * x + mix[1] * x0
        x = x + self.ssm_scale.to(dtype=x.dtype) * self.ssm(self.norm(x))
        if self.mlp is not None:
            x = x + self.mlp_scale.to(dtype=x.dtype) * self.mlp(self.mlp_norm(x))
        return x


# ---------------------------------------------------------------------------
# The core model
# ---------------------------------------------------------------------------
class ProjectorMLP(nn.Module):
    """MLP projector with LayerNorm."""
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.norm = nn.LayerNorm(hidden_dim)
        self.act = nn.GELU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x: Tensor) -> Tensor:
        return self.fc2(self.act(self.norm(self.fc1(x))))


class Jepa(nn.Module):
    """JEPA with SSM backbone, aligned with LeWorldModel (Maes et al. 2026)."""
    def __init__(self, vocab_size: int = 256, model_dim: int = 512, num_layers: int = 12,
                 d_state: int = 64, d_conv: int = 4, expand: int = 2,
                 embed_dim: int = 0, predictor_hidden_mult: int = 2,
                 mlp_mult: int = 3, mlp_every: int = 1, activation: str = "relu2",
                 tie_embeddings: int = 0, projector_type: str = "linear",
                 jepa_steps: int = 1, logit_softcap: float = 15.0,
                 softcap_type: str = "poly"):
        super().__init__()
        self.vocab_size = vocab_size
        self.model_dim = model_dim
        self.jepa_steps = jepa_steps
        self.tie_embeddings = tie_embeddings
        self.logit_softcap = logit_softcap
        self.softcap_type = softcap_type
        embed_dim = embed_dim if embed_dim > 0 else model_dim
        self._embed_dim = embed_dim

        # Token embedding (input)
        self.token_embed = nn.Embedding(vocab_size, embed_dim)
        nn.init.normal_(self.token_embed.weight, std=0.02)

        # Optional embed projection (when embed_dim != model_dim)
        self.embed_proj = nn.Linear(embed_dim, model_dim, bias=False) if embed_dim != model_dim else None
        # Reverse projection for tied embeddings with dim mismatch (not needed for tie=3, adapter handles it)
        self.embed_proj_rev = nn.Linear(model_dim, embed_dim, bias=False) if (
            embed_dim != model_dim and tie_embeddings in (1, 2)) else None

        # U-Net: split layers into encoder/decoder halves with skip connections.
        # Standard U-Net pairing: encoder pushes LIFO, decoder pops LIFO, so
        # skip_weights[0] pairs with the DEEPEST encoder output (standard U-Net).
        # For odd num_layers (e.g. 13 → 6 enc + 7 dec), the first decoder block
        # has no skip partner — the if-guard silently handles this.
        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))

        # SSM encoder backbone (mlp_every controls which blocks get MLP: 1=all, 2=alternate, etc.)
        self.blocks = nn.ModuleList([
            SSMBlock(model_dim, d_state, d_conv, expand,
                     mlp_mult if (mlp_every > 0 and i % mlp_every == 0) else 0,
                     activation)
            for i in range(num_layers)
        ])
        self.norm = RMSNorm()

        # Projector: h -> z (JEPA prediction space)
        if projector_type == "mlp":
            self.projector = ProjectorMLP(model_dim, model_dim * 2, model_dim)
            self.pred_proj = ProjectorMLP(model_dim, model_dim * 2, model_dim)
        else:  # linear
            self.projector = nn.Linear(model_dim, model_dim, bias=False)
            self.pred_proj = nn.Linear(model_dim, model_dim, bias=False)

        # Predictor: z_t -> predicted z_{t+1}
        pred_hidden = model_dim * predictor_hidden_mult
        self.predictor = nn.Sequential(
            nn.Linear(model_dim, pred_hidden),
            nn.GELU(),
            nn.Linear(pred_hidden, model_dim),
        )
        nn.init.zeros_(self.predictor[-1].weight)
        nn.init.zeros_(self.predictor[-1].bias)

        # Vocab bias — learnable per-token output bias (free BPB calibration)
        self.vocab_bias = nn.Parameter(torch.zeros(vocab_size, dtype=torch.float32))

        # LM head: h -> vocab logits
        # tie_embeddings: 0=separate lm_head, 1=tied, 2=tied+correction, 3=tied+nonlinear adapter
        self.lm_head_adapter = None
        if tie_embeddings == 3:
            # Nonlinear adapter before tied projection: h -> Linear -> GELU -> F.linear(_, embed.weight)
            self.lm_head_adapter = nn.Sequential(
                nn.Linear(model_dim, embed_dim, bias=False),
                nn.GELU(),
            )
            self.lm_head = None
            self.lm_head_correction = None
        elif tie_embeddings >= 1:
            self.lm_head = None
            self.lm_head_correction = nn.Parameter(
                torch.zeros(vocab_size, embed_dim)) if tie_embeddings == 2 else None
        else:
            self.lm_head = nn.Linear(model_dim, vocab_size, bias=False)
            self.lm_head_correction = None

    def _softcap(self, logits: Tensor) -> Tensor:
        """Logit softcap: prevents overconfident predictions, stabilizes early training."""
        s = self.logit_softcap
        if s <= 0:
            return logits
        if self.softcap_type == "tanh":
            return s * torch.tanh(logits / s)
        # Polynomial approximation (faster, from ternary code)
        x_sc = torch.clamp(logits / s, -2.0, 2.0)
        x2 = x_sc * x_sc
        return s * torch.clamp(x_sc * (1.0 - x2 / 3.0 + x2 * x2 / 15.0), -1.0, 1.0)

    def _get_logits(self, h: Tensor) -> Tensor:
        """Compute logits from hidden states."""
        if self.tie_embeddings >= 1:
            if self.lm_head_adapter is not None:
                # tie=3: nonlinear adapter -> tied projection
                proj = self.lm_head_adapter(h)
            elif self.embed_proj_rev is not None:
                # tie=1,2 with dim mismatch: linear reverse projection
                proj = self.embed_proj_rev(h)
            else:
                proj = h
            weight = self.token_embed.weight
            if self.lm_head_correction is not None:
                weight = weight + self.lm_head_correction
            logits_raw = F.linear(proj, weight.to(h.dtype))
        else:
            logits_raw = self.lm_head(h)
        return self._softcap(logits_raw + self.vocab_bias.to(h.dtype))

    def encode(self, input_ids: Tensor) -> Tensor:
        """input_ids: (B, T) -> hidden states h: (B, T, D)"""
        x = self.token_embed(input_ids)
        if self.embed_proj is not None:
            x = self.embed_proj(x)
        # Normalize embedding output before entering residual stream
        x = F.rms_norm(x, (x.size(-1),))
        x0 = x

        # U-Net style encoder/decoder with skip connections
        skips = []
        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):
            bi = self.num_encoder_layers + i
            if skips:
                x = x + self.skip_weights[i].to(dtype=x.dtype) * skips.pop()
            x = self.blocks[bi](x, x0)

        return self.norm(x)

    def forward(self, input_ids: Tensor, target_ids: Tensor,
                reduction: str = "mean", temperature: float | None = None,
                *, train_mode: bool = False, sigreg: SIGReg | None = None,
                sigreg_lambda: float = 0.0, ce_weight: float = 1.0,
                jepa_weight: float = 0.0, detach_targets: bool = False,
                ) -> Tensor | tuple[Tensor, dict]:
        """
        Unified forward for both training and eval (DDP hooks into this method).

        train_mode=False: returns NLL loss (for eval / BPB)
        train_mode=True:  returns (combined_loss, diagnostics_dict) with JEPA+SIGReg+CE
        """
        h = self.encode(input_ids)

        if not train_mode:
            # --- Eval path: simple CE for BPB ---
            logits = self._get_logits(h)
            if temperature is not None and temperature != 1.0:
                logits = logits / temperature
            logits_flat = logits.reshape(-1, self.vocab_size)
            targets_flat = target_ids.reshape(-1)
            if reduction == "none":
                return F.cross_entropy(logits_flat.float(), targets_flat,
                                       reduction="none").reshape(input_ids.shape)
            # CE loss without Z-loss (Z-loss is training-only regularization)
            logits_f = logits_flat.float()
            lse = torch.logsumexp(logits_f, dim=-1)
            target_logits = logits_f.gather(1, targets_flat.unsqueeze(1)).squeeze(1)
            return (lse - target_logits).mean()

        # --- Training path: JEPA + SIGReg + CE ---
        z = self.projector(h)
        diagnostics = {}
        loss = torch.zeros((), device=h.device)

        # JEPA: predict next projected embedding(s) via autoregressive composition
        if jepa_weight > 0 and self.jepa_steps > 0:
            jepa_loss = torch.zeros((), device=h.device)
            n_steps = 0
            prev_z = z[:, :-1, :]
            for k in range(1, self.jepa_steps + 1):
                if k >= z.size(1):
                    break
                pred_z = self.pred_proj(self.predictor(prev_z))
                target_z = z[:, k:k + prev_z.size(1), :]
                if detach_targets:
                    target_z = target_z.detach()
                L = min(pred_z.size(1), target_z.size(1))
                jepa_loss = jepa_loss + F.mse_loss(pred_z[:, :L], target_z[:, :L])
                n_steps += 1
                prev_z = pred_z[:, :L]
            if n_steps > 0:
                jepa_loss = jepa_loss / n_steps
                loss = loss + jepa_weight * jepa_loss
                diagnostics["jepa"] = jepa_loss.detach()

        # SIGReg on projected embeddings
        if sigreg is not None and sigreg_lambda > 0:
            sig_loss = sigreg(z)
            loss = loss + sigreg_lambda * sig_loss
            diagnostics["sigreg"] = sig_loss.detach()

        # CE: lm_head on raw hidden states + z-loss
        if ce_weight > 0:
            logits = self._get_logits(h)
            logits_flat = logits.reshape(-1, self.vocab_size).float()
            targets_flat = target_ids.reshape(-1)
            # Fused CE + Z-loss (single logsumexp)
            lse = torch.logsumexp(logits_flat, dim=-1)
            target_logits = logits_flat.gather(1, targets_flat.unsqueeze(1)).squeeze(1)
            ce_loss = (lse - target_logits).mean() + 1e-4 * (lse ** 2).mean()
            loss = loss + ce_weight * ce_loss
            diagnostics["ce"] = ce_loss.detach()

        return loss, diagnostics

# ---------------------------------------------------------------------------
# Muon optimizer (Newton-Schulz orthogonalized momentum)
# ---------------------------------------------------------------------------
def _ns_orth(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

ns_orth = _ns_orth   # will be compiled in main()


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

    @torch.no_grad()
    def step(self, closure=None):
        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, mom = group["lr"], group["momentum"]
            steps, nesterov = group["backend_steps"], group["nesterov"]
            wd = group.get("wd", 0.0)
            total = sum(p.numel() for p in params)
            flat = torch.zeros(total, device=params[0].device, dtype=torch.bfloat16)
            cur = 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 "buf" not in state:
                        state["buf"] = torch.zeros_like(g)
                    buf = state["buf"]
                    buf.mul_(mom).add_(g)
                    if nesterov:
                        g = g.add(buf, alpha=mom)
                    g = F.rms_norm(g.float(), (g.size(-1),)).bfloat16()
                    g = ns_orth(g, steps=steps)
                    g *= max(1, g.size(0) / g.size(1)) ** 0.5
                    flat[cur:cur + p.numel()] = g.reshape(-1)
                cur += p.numel()
            if distributed:
                dist.all_reduce(flat, op=dist.ReduceOp.SUM)
            cur = 0
            for p in params:
                g = flat[cur:cur + p.numel()].view_as(p).to(p.dtype)
                if wd > 0:
                    p.mul_(1 - lr * wd)
                p.add_(g, alpha=-lr)
                cur += p.numel()


# ---------------------------------------------------------------------------
# Data loading — challenge-format shards (byte or BPE)
# ---------------------------------------------------------------------------
def load_shard(path: Path, byte_mode: bool = True) -> Tensor:
    """Load a challenge-format shard (header + uint16 values).
    byte_mode=True: assert values < 256, return uint8.
    byte_mode=False: return uint16 (BPE token IDs)."""
    header_bytes = 256 * np.dtype("<i4").itemsize
    header = np.fromfile(path, dtype="<i4", count=256)
    if header.size != 256 or int(header[0]) != 20240520:
        raise ValueError(f"Unexpected shard header for {path}")
    num_tokens = int(header[2])
    tokens = np.fromfile(path, dtype="<u2", count=num_tokens, offset=header_bytes)
    if byte_mode:
        if tokens.max() > 255:
            raise ValueError(f"Shard {path} contains token IDs > 255 (max={tokens.max()}). "
                             f"This is not byte-level data — check your dataset.")
        return torch.from_numpy(tokens.astype(np.uint8, copy=False))
    # BPE: cast to int32 for safe PyTorch interop (uint16 has limited operator support)
    return torch.from_numpy(tokens.astype(np.int32, copy=False))


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

    def _advance(self):
        self.file_idx = (self.file_idx + 1) % len(self.files)
        self.data = load_shard(self.files[self.file_idx], self.byte_mode)
        self.pos = 0

    def take(self, n: int) -> Tensor:
        chunks = []
        remaining = n
        while remaining > 0:
            avail = self.data.numel() - self.pos
            if avail <= 0:
                self._advance()
                continue
            k = min(remaining, avail)
            chunks.append(self.data[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,
                 byte_mode: bool = True):
        self.rank, self.world_size, self.device = rank, world_size, device
        self.stream = TokenStream(pattern, byte_mode)

    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]
        local = local.pin_memory().to(self.device, non_blocking=True).to(torch.int64)
        x = local[:-1].reshape(-1, seq_len)
        y = local[1:].reshape(-1, seq_len)
        return x, y


# ---------------------------------------------------------------------------
# Validation
# ---------------------------------------------------------------------------
def load_val_tokens(pattern: str, seq_len: int, byte_mode: bool = True,
                    max_tokens: int = int(os.environ.get("VAL_MAX_TOKENS", 500_000))) -> Tensor:
    files = sorted(glob.glob(pattern))
    assert files, f"No files: {pattern}"
    data = torch.cat([load_shard(Path(p), byte_mode) for p in files]).contiguous()
    if max_tokens > 0:
        data = data[:max_tokens + 1]
    u = ((data.numel() - 1) // seq_len) * seq_len
    return data[:u + 1]


def build_bpe_luts(sp, vocab_size: int, device: torch.device):
    """Build token-to-byte lookup tables for BPE BPB calculation (from ternary code)."""
    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("\u2581"):
            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 eval_val(args, model, rank, world_size, device,
             val_tokens: Tensor, temperature: float = 1.0,
             bpe_luts: tuple | None = None):
    """Evaluate BPB on validation set.
    bpe_luts: (base_bytes_lut, has_leading_space_lut, is_boundary_token_lut) for BPE BPB.
    If None, byte-level BPB = loss / ln(2)."""
    local_batch_tokens = args.val_batch_size // world_size
    local_batch_seqs = max(1, 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
    loss_sum = torch.zeros((), device=device, dtype=torch.float64)
    token_count = torch.zeros((), device=device, dtype=torch.float64)
    byte_count = torch.zeros((), device=device, dtype=torch.float64)
    model.eval()
    with torch.inference_mode():
        for batch_start in range(seq_start, seq_end, local_batch_seqs):
            batch_end = min(batch_start + local_batch_seqs, seq_end)
            raw_start = batch_start * args.train_seq_len
            raw_end = batch_end * args.train_seq_len + 1
            local = val_tokens[raw_start:raw_end].to(device=device, dtype=torch.int64)
            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):
                batch_loss = model(x, y, temperature=temperature).detach()
            n = float(y.numel())
            loss_sum += batch_loss.to(torch.float64) * n
            token_count += n
            # BPE: count actual bytes per token for accurate BPB
            if bpe_luts is not None:
                base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = bpe_luts
                prev_ids, tgt_ids = x.reshape(-1), y.reshape(-1)
                tok_bytes = base_bytes_lut[tgt_ids].to(torch.int16)
                tok_bytes += (has_leading_space_lut[tgt_ids] & ~is_boundary_token_lut[prev_ids]).to(torch.int16)
                byte_count += tok_bytes.to(torch.float64).sum()
    if dist.is_available() and dist.is_initialized():
        for t in (loss_sum, token_count, byte_count):
            dist.all_reduce(t, op=dist.ReduceOp.SUM)
    val_loss = loss_sum / token_count
    if bpe_luts is not None:
        # BPE: BPB = (loss / ln2) * (tokens / bytes)
        bpb = (val_loss.item() / math.log(2.0)) * (token_count.item() / byte_count.item())
    else:
        # Byte: each prediction is one byte
        bpb = val_loss.item() / math.log(2.0)
    model.train()
    return float(val_loss.item()), float(bpb)


def eval_val_sliding(args, model, rank, world_size, device,
                     val_tokens: Tensor, stride: int = 64, temperature: float = 1.0,
                     bpe_luts: tuple | None = None):
    """Sliding window evaluation for better BPB."""
    seq_len = args.train_seq_len
    batch_size = args.sliding_batch_size
    total_tokens = val_tokens.numel() - 1
    all_starts = list(range(0, total_tokens - seq_len, stride))
    my_starts = all_starts[rank::world_size]
    loss_sum = torch.zeros((), device=device, dtype=torch.float64)
    token_count = torch.zeros((), device=device, dtype=torch.float64)
    byte_count = torch.zeros((), device=device, dtype=torch.float64)
    model.eval()
    with torch.inference_mode():
        for i in range(0, len(my_starts), batch_size):
            batch_starts = my_starts[i:i + batch_size]
            starts_t = torch.tensor(batch_starts, dtype=torch.int64)
            offsets = torch.arange(seq_len + 1, dtype=torch.int64)
            indices = starts_t.unsqueeze(1) + offsets.unsqueeze(0)
            local_batch = val_tokens[indices].to(device=device, dtype=torch.int64, non_blocking=True)
            x, y = local_batch[:, :-1], local_batch[:, 1:]
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                per_token_loss = model(x, y, reduction="none", temperature=temperature).detach()
            for b, start in enumerate(batch_starts):
                score_from = 0 if start == 0 else seq_len - stride
                scored = per_token_loss[b, score_from:]
                sx, sy = x[b, score_from:], y[b, score_from:]
                loss_sum += scored.to(torch.float64).sum()
                token_count += scored.numel()
                if bpe_luts is not None:
                    base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = bpe_luts
                    tok_bytes = base_bytes_lut[sy].to(torch.int16)
                    tok_bytes += (has_leading_space_lut[sy] & ~is_boundary_token_lut[sx]).to(torch.int16)
                    byte_count += tok_bytes.to(torch.float64).sum()
    if dist.is_available() and dist.is_initialized():
        for t in (loss_sum, token_count, byte_count):
            dist.all_reduce(t, op=dist.ReduceOp.SUM)
    val_loss = loss_sum / token_count
    if bpe_luts is not None:
        bpb = (val_loss.item() / math.log(2.0)) * (token_count.item() / byte_count.item())
    else:
        bpb = val_loss.item() / math.log(2.0)
    model.train()
    return float(val_loss.item()), float(bpb)


# ---------------------------------------------------------------------------
# JEPA Diagnostics — verify representation quality independent of BPB
# ---------------------------------------------------------------------------
def jepa_diagnostics(model: nn.Module, val_tokens: Tensor, device: torch.device,
                     seq_len: int = 2048, max_samples: int = 50000):
    model.eval()
    base = model.module if hasattr(model, 'module') else model
    base = base._orig_mod if hasattr(base, '_orig_mod') else base

    n = min(max_samples + 1, val_tokens.numel())
    data = val_tokens[:n].to(device=device, dtype=torch.int64)
    usable = ((data.numel() - 1) // seq_len) * seq_len
    x = data[:usable].reshape(-1, seq_len)
    y = data[1:usable + 1].reshape(-1, seq_len)
    n_seqs = min(x.shape[0], max_samples // seq_len)
    x, y = x[:n_seqs], y[:n_seqs]

    cos_sims, jepa_mses = [], []
    lm_top1 = lm_top5 = lm_top10 = total = 0
    all_hidden, all_targets = [], []

    with torch.no_grad():
        for i in range(n_seqs):
            xi, yi = x[i:i+1], y[i:i+1]
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                h = base.encode(xi)
                z = base.projector(h)
                pred_z = base.pred_proj(base.predictor(z[:, :-1, :]))
            target_z = z[:, 1:, :]
            cos = F.cosine_similarity(pred_z[0].float(), target_z[0].float(), dim=-1)
            cos_sims.append(cos.mean().item())
            jepa_mses.append((pred_z - target_z).pow(2).mean().item())
            h_f = h[0].float()
            logits = base._get_logits(h[0])                  # keep bfloat16 for model layers
            targets = yi[0]
            preds = logits.argsort(dim=-1, descending=True)
            lm_top1 += (preds[:, :1] == targets.unsqueeze(1)).any(1).sum().item()
            lm_top5 += (preds[:, :5] == targets.unsqueeze(1)).any(1).sum().item()
            lm_top10 += (preds[:, :10] == targets.unsqueeze(1)).any(1).sum().item()
            total += targets.numel()
            if len(all_hidden) * seq_len < 100000:
                all_hidden.append(h_f.cpu())
                all_targets.append(targets.cpu())

    # Linear probe on frozen hidden states
    probe_bpb = float('nan')
    if all_hidden:
        H = torch.cat(all_hidden, dim=0)
        T = torch.cat(all_targets, dim=0)
        n_train = int(0.8 * H.shape[0])
        H_tr, H_val = H[:n_train].to(device), H[n_train:].to(device)
        T_tr, T_val = T[:n_train].to(device), T[n_train:].to(device)
        probe = nn.Linear(H.shape[-1], base.vocab_size).to(device)
        opt = torch.optim.Adam(probe.parameters(), lr=1e-3)
        bs = min(4096, n_train)
        for _ in range(200):
            idx = torch.randint(0, n_train, (bs,), device=device)
            loss = F.cross_entropy(probe(H_tr[idx]), T_tr[idx])
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.inference_mode():
            probe_bpb = F.cross_entropy(probe(H_val), T_val).item() / math.log(2.0)

    model.train()
    return {
        "jepa_cos_sim": sum(cos_sims) / max(len(cos_sims), 1),
        "jepa_mse": sum(jepa_mses) / max(len(jepa_mses), 1),
        "lm_top1": lm_top1 / max(total, 1),
        "lm_top5": lm_top5 / max(total, 1),
        "lm_top10": lm_top10 / max(total, 1),
        "probe_bpb": probe_bpb,
    }


# ---------------------------------------------------------------------------
# Temperature search (eval-time calibration)
# ---------------------------------------------------------------------------
def find_best_temperature(args, model, rank, world_size, device, val_tokens,
                          bpe_luts: tuple | None = None):
    """Two-phase temperature search: coarse grid then fine-grained around best."""
    best_t, best_loss = 1.0, float("inf")
    # Phase 1: coarse grid
    for t in [0.5, 0.7, 0.85, 0.95, 1.0, 1.1, 1.3, 1.5, 2.0]:
        loss, _ = eval_val(args, model, rank, world_size, device,
                           val_tokens, temperature=t, bpe_luts=bpe_luts)
        if loss < best_loss:
            best_loss = loss
            best_t = t
    # Phase 2: fine grid around best (step 0.02)
    center = best_t
    for t in [center - 0.06, center - 0.04, center - 0.02,
              center + 0.02, center + 0.04, center + 0.06]:
        if t <= 0:
            continue
        loss, _ = eval_val(args, model, rank, world_size, device,
                           val_tokens, temperature=t, bpe_luts=bpe_luts)
        if loss < best_loss:
            best_loss = loss
            best_t = t
    return best_t


# ---------------------------------------------------------------------------
# Checkpointing
# ---------------------------------------------------------------------------
def _ckpt_path(checkpoint_dir: str, step: int) -> str:
    return os.path.join(checkpoint_dir, f"ckpt_step{step:07d}.pt")

def _latest_checkpoint(checkpoint_dir: str) -> str | None:
    if not os.path.isdir(checkpoint_dir):
        return None
    ckpts = sorted(glob.glob(os.path.join(checkpoint_dir, "ckpt_step*.pt")))
    return ckpts[-1] if ckpts else None

def save_checkpoint(checkpoint_dir: str, step: int, base_model: nn.Module,
                    optimizers: list, training_time_ms: float) -> None:
    os.makedirs(checkpoint_dir, exist_ok=True)
    path = _ckpt_path(checkpoint_dir, step)
    payload = {
        "step": step,
        "training_time_ms": training_time_ms,
        "model": base_model.state_dict(),
        "optimizers": [o.state_dict() for o in optimizers],
        "rng_cpu": torch.get_rng_state(),
        "rng_cuda": torch.cuda.get_rng_state(),
    }
    tmp = path + ".tmp"
    torch.save(payload, tmp)
    os.replace(tmp, path)   # atomic on POSIX

def load_checkpoint(path: str, base_model: nn.Module, optimizers: list,
                    device: torch.device) -> tuple[int, float]:
    payload = torch.load(path, map_location=device, weights_only=False)
    base_model.load_state_dict(payload["model"], strict=True)
    for opt, sd in zip(optimizers, payload["optimizers"]):
        opt.load_state_dict(sd)
    torch.set_rng_state(payload["rng_cpu"].cpu())
    torch.cuda.set_rng_state(payload["rng_cuda"].cpu())
    return payload["step"], payload["training_time_ms"]


# ---------------------------------------------------------------------------
# Training-only modules (discarded from artifact — not needed at eval time)
# ---------------------------------------------------------------------------
TRAINING_ONLY = ("predictor.", "projector.", "pred_proj.")

def _is_training_only(name: str) -> bool:
    return any(name.startswith(p) or f".{p}" in name for p in TRAINING_ONLY)


def _is_large_matrix(name: str, tensor: Tensor, min_numel: int = 4096) -> bool:
    """Large 2D weight matrices suitable for INT quantization.
    Excludes: token_embed (lookup table), A_log (SSM log-transition),
    lm_head_correction (fine residual correction needs full precision),
    skip_weights (small U-Net skips)."""
    return (tensor.ndim == 2 and tensor.numel() >= min_numel
            and "token_embed" not in name and "A_log" not in name
            and "lm_head_correction" not in name and "skip_weight" not in name)


# ---------------------------------------------------------------------------
# QAT — Quantization-Aware Training (snap/restore for Mamba2 compatibility)
# Note: Mamba2's CUDA kernels use internal weight layouts that prevent
# standard STE injection. snap/restore computes gradients at the quantized
# point, then applies them to the full-precision weights. This is noisier
# than proper STE but necessary for external CUDA kernel modules.
# ---------------------------------------------------------------------------
def _build_qat_param_list(model: nn.Module) -> list[tuple[str, nn.Parameter]]:
    """Pre-compute list of parameters subject to INT-N QAT. Uses same filter as serialization."""
    return [(name, param) for name, param in model.named_parameters()
            if not _is_training_only(name) and _is_large_matrix(name, param)]


def _build_fp8_param_list(model: nn.Module) -> list[tuple[str, nn.Parameter]]:
    """Pre-compute list of parameters subject to FP8 QAT (non-INT 2D params like token_embed).
    Excludes Mamba2 internals (their CUDA kernels require exact dtype matching)."""
    return [(name, param) for name, param in model.named_parameters()
            if not _is_training_only(name)
            and not _is_large_matrix(name, param)
            and param.ndim == 2 and param.numel() >= 256
            and "A_log" not in name and ".ssm." not in name]


def _snap_weights_to_grid(qat_params: list[tuple[str, nn.Parameter]], bits: int = 6,
                          fp8_params: list[tuple[str, nn.Parameter]] | None = None,
                          ) -> tuple[list[Tensor], list[Tensor]]:
    """Snap large matrices to INT-N grid and optionally FP8 grid. Returns originals for restoration."""
    qmax = (1 << (bits - 1)) - 1
    int_originals = []
    with torch.no_grad():
        for _, param in qat_params:
            int_originals.append(param.data.clone())
            absmax = param.abs().amax(dim=-1, keepdim=True).clamp(min=1e-8)
            scale = absmax / qmax
            param.data.copy_(torch.clamp(torch.round(param / scale), -qmax, qmax) * scale)
    fp8_originals = []
    if fp8_params:
        with torch.no_grad():
            for _, param in fp8_params:
                fp8_originals.append(param.data.clone())
                # STE: snap to FP8 grid (same approach as ternary code's QATLinear)
                param.data.copy_(param.data.to(torch.float8_e4m3fn).to(param.dtype))
    return int_originals, fp8_originals


def _restore_weights(qat_params: list[tuple[str, nn.Parameter]], int_originals: list[Tensor],
                     fp8_params: list[tuple[str, nn.Parameter]] | None = None,
                     fp8_originals: list[Tensor] | None = None):
    """Restore original weights after QAT forward/backward."""
    with torch.no_grad():
        for (_, param), orig in zip(qat_params, int_originals):
            param.data.copy_(orig)
        if fp8_params and fp8_originals:
            for (_, param), orig in zip(fp8_params, fp8_originals):
                param.data.copy_(orig)


# ---------------------------------------------------------------------------
# Quantization + Compression Pipeline
# ---------------------------------------------------------------------------
def quantize_intN(t: Tensor, bits: int = 6) -> tuple[bytes, Tensor, list, int]:
    """Quantize a 2D tensor to N-bit integers with per-row absmax scaling.
    Returns (packed_bytes, scales_bf16, original_shape, n_vals).
    Uses np.packbits for fast bit packing. Input must be CPU tensor."""
    assert not t.is_cuda, "quantize_intN expects CPU tensor"
    t32 = t.float()
    if t32.ndim < 2:
        t32 = t32.unsqueeze(0)
    qmax = (1 << (bits - 1)) - 1
    absmax = t32.abs().amax(dim=-1, keepdim=True).clamp(min=1e-8)
    scale = absmax / qmax
    q = torch.clamp(torch.round(t32 / scale), -qmax, qmax).to(torch.int16)
    scale_bf16 = scale.bfloat16().squeeze(-1)

    # Shift to unsigned for packing
    flat = (q.reshape(-1) + qmax).numpy().astype(np.uint8)
    n_vals = len(flat)

    # Expand each value into its individual bits, flatten, then packbits
    bits_array = ((flat[:, None] >> np.arange(bits, dtype=np.uint8)) & 1).ravel()
    pad_len = (8 - len(bits_array) % 8) % 8
    if pad_len > 0:
        bits_array = np.pad(bits_array, (0, pad_len))
    packed_bytes = np.packbits(bits_array, bitorder='little').tobytes()

    return packed_bytes, scale_bf16, list(t.shape), n_vals


def dequantize_intN(packed_bytes: bytes, scale: Tensor, shape: list,
                    n_vals: int, bits: int = 6) -> Tensor:
    """Dequantize N-bit packed integers back to float. Matches quantize_intN's packbits layout."""
    qmax = (1 << (bits - 1)) - 1
    packed = np.frombuffer(packed_bytes, dtype=np.uint8)

    # Unpack all bits, then reassemble N-bit values
    all_bits = np.unpackbits(packed, bitorder='little')
    total_needed = n_vals * bits
    all_bits = all_bits[:total_needed].reshape(n_vals, bits)
    # Reconstruct unsigned values from bit planes
    powers = (1 << np.arange(bits, dtype=np.uint16))
    flat = (all_bits.astype(np.uint16) * powers).sum(axis=1).astype(np.int32) - qmax

    t = torch.from_numpy(flat.astype(np.float32))
    if len(shape) >= 2:
        t = t.reshape(-1, shape[-1])
        t = t * scale.float().unsqueeze(-1)
        return t[:shape[0], :shape[-1]].reshape(shape)
    else:
        return (t * scale.float()).reshape(shape)


# --- Compression backends ---
def _compress_lzma(data: bytes) -> bytes:
    return lzma.compress(data, preset=9)

def _decompress_lzma(data: bytes) -> bytes:
    return lzma.decompress(data)

def _compress_zstd(data: bytes) -> bytes:
    import zstandard as zstd
    return zstd.ZstdCompressor(level=22).compress(data)

def _decompress_zstd(data: bytes) -> bytes:
    import zstandard as zstd
    return zstd.ZstdDecompressor().decompress(data)

def _compress_brotli(data: bytes) -> bytes:
    import brotli
    return brotli.compress(data, quality=11)

def _decompress_brotli(data: bytes) -> bytes:
    import brotli
    return brotli.decompress(data)

COMPRESSORS = {
    "lzma": (_compress_lzma, _decompress_lzma),
    "zstd": (_compress_zstd, _decompress_zstd),
    "brotli": (_compress_brotli, _decompress_brotli),
}


# --- Full save/load pipeline ---
def save_model_quantized(model: nn.Module, code: str, path_prefix: str,
                         bits: int = 6, fp_storage: str = "bf16") -> tuple[str, int, int, dict]:
    """
    Quantize and compress the model. Tries all available compressors, picks smallest.
    Strips training-only parameters (predictor, projector, pred_proj).
    fp_storage: "fp8" stores non-INT 2D params as float8_e4m3fn (1 byte), "bf16" as bfloat16 (2 bytes).
    """
    sd = model.state_dict()
    quantized = {}
    stats = {"int_params": 0, "int_bytes": 0, "fp_params": 0, "fp_bytes": 0,
             "discarded_params": 0}
    for name, tensor in sd.items():
        t = tensor.detach().cpu().float().contiguous()
        if _is_training_only(name):
            stats["discarded_params"] += t.numel()
            continue
        if _is_large_matrix(name, t):
            packed, scale, shape, n_vals = quantize_intN(t, bits=bits)
            quantized[name] = {
                "type": f"int{bits}", "packed": packed, "scale": scale,
                "shape": shape, "n_vals": n_vals, "bits": bits,
            }
            stats["int_params"] += t.numel()
            stats["int_bytes"] += len(packed) + scale.numel() * 2
        elif fp_storage == "fp8" and t.ndim == 2:
            quantized[name] = {"type": "fp8", "data": t.to(torch.float8_e4m3fn)}
            stats["fp_params"] += t.numel()
            stats["fp_bytes"] += t.numel()
        else:
            quantized[name] = {"type": "bf16", "data": t.bfloat16()}
            stats["fp_params"] += t.numel()
            stats["fp_bytes"] += t.numel() * 2

    buf = io.BytesIO()
    torch.save(quantized, buf)
    raw_bytes = buf.getvalue()

    code_bytes = len(code.encode("utf-8"))
    results = {}
    for method, (compress_fn, _) in COMPRESSORS.items():
        try:
            compressed = compress_fn(raw_bytes)
            total = len(compressed) + code_bytes
            results[method] = (compressed, total)
        except (ImportError, Exception):
            pass

    assert results, "No compressor available"
    best_method = min(results, key=lambda m: results[m][1])
    best_compressed, best_total = results[best_method]

    best_path = f"{path_prefix}.{best_method}.ptz"
    with open(best_path, "wb") as f:
        f.write(best_compressed)

    stats["compression_results"] = {
        m: {"compressed_mb": len(c) / 1e6, "total_mb": t / 1e6}
        for m, (c, t) in results.items()
    }
    stats["best_method"] = best_method
    stats["raw_mb"] = len(raw_bytes) / 1e6

    return best_method, len(best_compressed), code_bytes, stats


def load_model_quantized(path: str, model: nn.Module,
                         device: torch.device) -> nn.Module:
    """Load quantized + compressed model. Auto-detects compression from extension."""
    for method, (_, decompress_fn) in COMPRESSORS.items():
        if f".{method}." in path:
            with open(path, "rb") as f:
                raw = decompress_fn(f.read())
            break
    else:
        with open(path, "rb") as f:
            raw = lzma.decompress(f.read())

    quantized = torch.load(io.BytesIO(raw), map_location="cpu", weights_only=False)
    out = {}
    for name, entry in quantized.items():
        if entry["type"].startswith("int"):
            bits = entry["bits"]
            t = dequantize_intN(entry["packed"], entry["scale"],
                                entry["shape"], entry["n_vals"], bits)
            out[name] = t.bfloat16()
        elif entry["type"] == "fp8":
            out[name] = entry["data"].to(torch.float32).to(torch.bfloat16)
        else:
            out[name] = entry["data"]

    model.load_state_dict(out, strict=False)
    return model.to(device)


# ---------------------------------------------------------------------------
# Training
# ---------------------------------------------------------------------------
SCALAR_PARAM_PATTERNS = ("bias", "vocab_bias", "lm_head_correction", "resid_mix", "skip_weight",
                         "ssm_scale", "mlp_scale")


def main() -> None:
    args = Hyperparameters()
    code = Path(__file__).read_text(encoding="utf-8")

    global ns_orth
    ns_orth = torch.compile(_ns_orth)

    # --- Distributed 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"))
    grad_accum_steps = max(1, 8 // world_size)
    grad_scale = 1.0 / grad_accum_steps

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA 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 = rank == 0
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    os.makedirs("logs/cuda/", exist_ok=True)
    logfile = f"logs/cuda/{args.run_id}.txt" if master else None

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

    log0(code, console=False)
    log0("=" * 100, console=False)
    log0(f"Python {sys.version}", console=False)
    log0(f"PyTorch {torch.__version__}", console=False)

    # --- Log all hyperparameters ---
    log0("--- Hyperparameters ---", console=False)
    log0(" ".join(f"{a}={getattr(args,a)}" for a in sorted(dir(args))
         if not a.startswith("_") and a not in ("train_files","val_files")
         and not callable(getattr(args,a))), console=False)

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

    # --- Tokenizer & validation data ---
    byte_mode = args.tokenizer == "byte"
    bpe_luts = None
    if not byte_mode:
        assert spm is not None, "sentencepiece required for BPE mode: pip install sentencepiece"
        sp = spm.SentencePieceProcessor(model_file=args.tokenizer_path)
        bpe_luts = build_bpe_luts(sp, args.vocab_size, device)
    val_tokens = load_val_tokens(args.val_files, args.train_seq_len, byte_mode)

    # --- Model ---
    base_model = Jepa(
        vocab_size=args.vocab_size,
        model_dim=args.model_dim,
        num_layers=args.num_layers,
        d_state=args.d_state,
        d_conv=args.d_conv,
        expand=args.expand,
        embed_dim=args.embed_dim,
        predictor_hidden_mult=args.predictor_hidden_mult,
        mlp_mult=args.mlp_mult,
        mlp_every=args.mlp_every,
        activation=args.activation,
        tie_embeddings=args.tie_embeddings,
        projector_type=args.projector_type,
        jepa_steps=args.jepa_steps,
        logit_softcap=args.logit_softcap,
        softcap_type=args.softcap_type,
    ).to(device).bfloat16()

    # Promote scalars and small params to float32 (matching ternary code pattern).
    # EXCEPTION: Mamba2's internal conv1d and SSM params must stay bfloat16 because
    # the causal_conv1d CUDA kernel requires weight and bias to have matching dtype.
    _mamba_internal = (".ssm.",)
    for name, param in base_model.named_parameters():
        if any(mi in name for mi in _mamba_internal):
            continue  # leave Mamba2 internals in bfloat16
        if param.ndim < 2 or any(pat in name for pat in SCALAR_PARAM_PATTERNS):
            param.data = param.data.float()

    # Auto-scale SIGReg projections if not set
    effective_embed_dim = args.embed_dim if args.embed_dim > 0 else args.model_dim
    sigreg_num_proj = args.sigreg_num_proj if args.sigreg_num_proj > 0 else effective_embed_dim * 2
    sigreg = SIGReg(d_model=args.model_dim, knots=args.sigreg_knots, num_proj=sigreg_num_proj).to(device)

    torch._dynamo.config.optimize_ddp = False
    compiled_model = torch.compile(base_model,
                                   mode=args.compile_mode if args.compile_mode != "default" else None)

    # DDP: static_graph when all params are used every step (default config).
    # Fall back to find_unused_parameters only if a loss term is disabled,
    # which leaves predictor/projector (jepa_weight=0) or lm_head (ce_weight=0) unused.
    _has_unused = (args.jepa_weight == 0 or args.ce_weight == 0)
    model = DDP(compiled_model, device_ids=[local_rank],
                broadcast_buffers=False, gradient_as_bucket_view=True,
                find_unused_parameters=_has_unused,
                static_graph=not _has_unused) if distributed else compiled_model

    # --- Optimizers ---
    embed_params = [base_model.token_embed.weight]
    embed_ids = {id(p) for p in embed_params}

    matrix_params = [p for n, p in base_model.named_parameters()
                     if p.ndim == 2
                     and not any(pat in n for pat in SCALAR_PARAM_PATTERNS)
                     and id(p) not in embed_ids]
    scalar_params = [p for n, p in base_model.named_parameters()
                     if (p.ndim < 2 or any(pat in n for pat in SCALAR_PARAM_PATTERNS))
                     and id(p) not in embed_ids]

    opt_muon = Muon(matrix_params, lr=args.matrix_lr, momentum=args.muon_momentum,
                    backend_steps=args.muon_backend_steps, wd=args.muon_wd)
    for g in opt_muon.param_groups:
        g["base_lr"] = args.matrix_lr

    opt_scalar = torch.optim.AdamW(
        [{"params": scalar_params, "lr": args.scalar_lr, "base_lr": args.scalar_lr}],
        betas=(args.beta1, args.beta2), eps=args.adam_eps,
        weight_decay=args.adam_wd, fused=True)

    opt_embed = torch.optim.AdamW(
        [{"params": embed_params, "lr": args.embed_lr, "base_lr": args.embed_lr}],
        betas=(args.beta1, args.beta2), eps=args.adam_eps,
        weight_decay=args.adam_wd, fused=True)

    optimizers = [opt_muon, opt_scalar, opt_embed]

    # --- Checkpoint resume ---
    _resume_step = 0
    _resume_training_ms = 0.0
    if args.checkpoint_every > 0:
        ckpt = _latest_checkpoint(args.checkpoint_dir)
        if ckpt is not None:
            log0(f"checkpoint found: {ckpt}")
            _resume_step, _resume_training_ms = load_checkpoint(
                ckpt, base_model, optimizers, device)
            log0(f"checkpoint loaded, starting from step {_resume_step} "
                 f"(accumulated_train_time:{_resume_training_ms:.0f}ms)")
        else:
            log0("no checkpoint found, starting from scratch")

    # --- Log config ---
    n_params = sum(p.numel() for p in base_model.parameters())
    n_train_only = sum(p.numel() for n, p in base_model.named_parameters()
                       if any(n.startswith(t) for t in ("predictor.", "projector.", "pred_proj.")))
    n_eval = n_params - n_train_only
    # Rough artifact size estimate
    _fp_bytes = 1 if args.fp_storage == "fp8" else 2  # FP8=1B, BF16=2B per param
    est_bytes = sum(
        p.numel() * args.quant_bits / 8 if _is_large_matrix(n, p)
        else p.numel() * (_fp_bytes if p.ndim == 2 else 2)
        for n, p in base_model.named_parameters()
        if not _is_training_only(n)
    )
    log0(f"Jepa | params:{n_params:,} (eval:{n_eval:,} train_only:{n_train_only:,}) "
         f"L:{args.num_layers} d:{args.model_dim} dS:{args.d_state} expand:{args.expand} "
         f"tok:{args.tokenizer} V:{args.vocab_size} ws:{world_size} ga:{grad_accum_steps} s:{args.seed}")
    log0(f"JEPA: jepa_w={args.jepa_weight} sigreg_l={args.sigreg_lambda} ce_w={args.ce_weight} "
         f"steps={args.jepa_steps} detach={args.detach_targets} "
         f"act={args.activation} mlp={args.mlp_mult} mlp_every={args.mlp_every} proj={args.projector_type} "
         f"tie={args.tie_embeddings} qat={args.qat_fraction} int{args.quant_bits} fp={args.fp_storage} "
         f"softcap={args.logit_softcap} ({args.softcap_type})")
    log0(f"est_raw:{est_bytes/1e6:.1f}MB (pre-compression) budget:16.00MB")

    # --- Data loader ---
    train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device, byte_mode)

    def zero_grad_all():
        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):
        """Warmdown schedule (warmup is handled by ghost warmup, not LR ramp)."""
        if args.warmdown_fraction <= 0:
            return 1.0
        if max_wallclock_ms is None:
            wd_start = int(args.iterations * (1.0 - args.warmdown_fraction))
            if step < wd_start:
                return 1.0
            return max((args.iterations - step) / max(args.iterations * args.warmdown_fraction, 1), 0.0)
        wd_ms = max_wallclock_ms * args.warmdown_fraction
        remaining = max(max_wallclock_ms - elapsed_ms, 0.0)
        return remaining / max(wd_ms, 1e-9) if remaining <= wd_ms else 1.0

    # --- Ghost warmup: run N steps to warm optimizer momentum buffers, then reset model ---
    # This gives Muon/Adam warm momentum buffers without burning real training budget.
    # Model weights are restored to their initial values after warmup, but optimizer
    # state (momentum buffers, Adam moments) are kept. Data loader is also reset so
    # training sees the same data as if warmup never happened.
    if args.warmup_steps > 0 and _resume_step == 0:
        _saved_model = {n: t.detach().cpu().clone() for n, t in base_model.state_dict().items()}
        _saved_opts = [copy.deepcopy(o.state_dict()) for o in optimizers]
        model.train()
        for ws in range(args.warmup_steps):
            zero_grad_all()
            for mi in range(grad_accum_steps):
                if distributed:
                    model.require_backward_grad_sync = mi == grad_accum_steps - 1
                x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len,
                                               grad_accum_steps)
                # SIGReg lambda for warmup
                sig_lam = args.sigreg_lambda
                torch.compiler.cudagraph_mark_step_begin()
                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                    loss, _ = model(x, y, train_mode=True, sigreg=sigreg,
                                    sigreg_lambda=sig_lam, ce_weight=args.ce_weight,
                                    jepa_weight=args.jepa_weight,
                                    detach_targets=args.detach_targets)
                (loss * grad_scale).backward()
            for o in optimizers:
                o.step()
            zero_grad_all()
            log0(f"warmup:{ws+1}/{args.warmup_steps}")

        # Restore model weights but KEEP optimizer momentum buffers
        base_model.load_state_dict(_saved_model, strict=True)
        # Restore optimizer LR/weight-decay config but keep moment buffers
        for o, s in zip(optimizers, _saved_opts):
            # Only restore param_groups config, not state (moments)
            for pg_new, pg_saved in zip(o.param_groups, s["param_groups"]):
                for k in pg_saved:
                    if k != "params":
                        pg_new[k] = pg_saved[k]
        zero_grad_all()
        # Reset Adam step counters so bias correction starts from step 1
        # (keeps warmed momentum buffers but fixes the 1-beta^t denominator)
        for o in optimizers:
            for state in o.state.values():
                if "step" in state:
                    state["step"] = state["step"] * 0  # type-agnostic (tensor or int)
        # Reset data loader so training sees the same data
        train_loader = DistributedTokenLoader(args.train_files, rank, world_size, device, byte_mode)
        # Reset compiled graph cache — weight restore can invalidate assumptions
        torch._dynamo.reset()
        del _saved_model, _saved_opts

    # --- Main training loop ---
    training_time_ms = _resume_training_ms
    stop_after_step: int | None = None
    train_loss = torch.zeros((), device=device)
    _qat_logged = False
    qat_params = _build_qat_param_list(base_model)  # cached for speed
    fp8_params = _build_fp8_param_list(base_model) if args.fp_storage == "fp8" else None
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    step = _resume_step
    steps_this_session = 0

    # Guard: budget already exhausted at resume time
    if max_wallclock_ms is None:
        if step >= args.iterations:
            log0(f"checkpoint step {step} >= iterations {args.iterations}, skipping training")
            stop_after_step = step
    else:
        if training_time_ms >= max_wallclock_ms:
            log0(f"checkpoint train_time {training_time_ms:.0f}ms >= cap {max_wallclock_ms:.0f}ms, skipping")
            stop_after_step = step

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

        # Validation
        if last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0):
            torch.cuda.synchronize()
            training_time_ms += 1000.0 * (time.perf_counter() - t0)
            val_loss, val_bpb = eval_val(args, model, rank, world_size, device,
                                         val_tokens, bpe_luts=bpe_luts)
            log0(f"step:{step}/{args.iterations} val_loss:{val_loss:.4f} val_bpb:{val_bpb:.4f} "
                 f"train_time:{training_time_ms:.0f}ms")
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            steps_this_session = 0

        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 step:{step}/{args.iterations}")
            break

        # LR scheduling
        elapsed_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0)
        scale = lr_mul(step, elapsed_ms)

        # Muon momentum warmup
        if args.muon_momentum_warmup_steps > 0:
            frac = min(step / max(args.muon_momentum_warmup_steps, 1), 1.0)
            for g in opt_muon.param_groups:
                g["momentum"] = ((1 - frac) * args.muon_momentum_warmup_start
                                 + frac * args.muon_momentum)

        # Training step
        zero_grad_all()
        train_loss.zero_()

        # Accumulate diagnostics across micro-steps
        diag_accum = {}

        # SIGReg lambda scheduling (ramp from 0 to target)
        if args.sigreg_schedule and max_wallclock_ms is not None:
            sig_lam = args.sigreg_lambda * min(elapsed_ms / (max_wallclock_ms * 0.3), 1.0)
        elif args.sigreg_schedule:
            sig_lam = args.sigreg_lambda * min(step / max(args.iterations * 0.3, 1), 1.0)
        else:
            sig_lam = args.sigreg_lambda

        # QAT: determine if active
        qat_active = False
        if args.qat_fraction > 0:
            if max_wallclock_ms is not None:
                qat_active = elapsed_ms >= max_wallclock_ms * (1.0 - args.qat_fraction)
            else:
                qat_active = step >= int(args.iterations * (1.0 - args.qat_fraction))

        # QAT: snap weights to grid before forward (save originals)
        qat_originals = None
        if qat_active:
            qat_originals = _snap_weights_to_grid(qat_params, bits=args.quant_bits,
                                                   fp8_params=fp8_params)

        for micro in range(grad_accum_steps):
            if distributed:
                model.require_backward_grad_sync = micro == grad_accum_steps - 1
            x, y = train_loader.next_batch(args.train_batch_tokens, args.train_seq_len,
                                           grad_accum_steps)

            torch.compiler.cudagraph_mark_step_begin()
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                loss, diag = model(x, y, train_mode=True, sigreg=sigreg,
                                   sigreg_lambda=sig_lam, ce_weight=args.ce_weight,
                                   jepa_weight=args.jepa_weight,
                                   detach_targets=args.detach_targets)
            train_loss.add_(loss.detach())
            (loss * grad_scale).backward()

            # Accumulate diagnostics
            for k, v in diag.items():
                if k not in diag_accum:
                    diag_accum[k] = v.clone()
                else:
                    diag_accum[k] += v

        # QAT: restore originals before optimizer step
        if qat_originals is not None:
            int_orig, fp8_orig = qat_originals
            _restore_weights(qat_params, int_orig, fp8_params=fp8_params, fp8_originals=fp8_orig)

        train_loss /= grad_accum_steps
        for k in diag_accum:
            diag_accum[k] /= grad_accum_steps

        # Gradient clipping
        if args.grad_clip_norm > 0:
            torch.nn.utils.clip_grad_norm_(base_model.parameters(), args.grad_clip_norm)

        # Optimizer step with LR scaling
        for opt in optimizers:
            for g in opt.param_groups:
                g["lr"] = g["base_lr"] * scale
            opt.step()
        zero_grad_all()
        step += 1
        steps_this_session += 1
        approx_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0)

        # Checkpoint save
        if master and args.checkpoint_every > 0 and step % args.checkpoint_every == 0:
            save_checkpoint(args.checkpoint_dir, step, base_model, optimizers, approx_ms)
            log0(f"step:{step} checkpoint saved to {_ckpt_path(args.checkpoint_dir, step)}")

        # QAT start notification (once)
        if qat_active and not _qat_logged:
            log0(f"step:{step} QAT started (INT{args.quant_bits})")
            _qat_logged = True

        if args.train_log_every > 0 and step % args.train_log_every == 0:
            session_ms = 1000.0 * (time.perf_counter() - t0)
            d_str = " ".join(f"{k}:{v.item():.4f}" for k, v in diag_accum.items())
            log0(f"step:{step}/{args.iterations} loss:{train_loss.item():.4f} "
                 f"{d_str} t:{approx_ms:.0f}ms step_avg:{session_ms/steps_this_session:.1f}ms")

        # Wallclock cap
        if stop_after_step is None and max_wallclock_ms is not None and step % 10 == 0:
            reached = approx_ms >= max_wallclock_ms
            if distributed:
                cap_t = torch.tensor(int(reached), device=device)
                dist.all_reduce(cap_t, op=dist.ReduceOp.MAX)
                reached = bool(cap_t.item())
            if reached:
                stop_after_step = step

    # --- JEPA diagnostics (BEFORE serialization — projector/predictor are discarded from artifact) ---
    if master:
        diag = jepa_diagnostics(model, val_tokens, device, seq_len=args.train_seq_len)
        log0(f"jepa_diagnostics cos_sim:{diag['jepa_cos_sim']:.4f} jepa_mse:{diag['jepa_mse']:.4f} "
             f"lm_top1:{diag['lm_top1']:.4f} lm_top5:{diag['lm_top5']:.4f} "
             f"lm_top10:{diag['lm_top10']:.4f} probe_bpb:{diag['probe_bpb']:.4f}")

    # --- Serialization (QAT already applied during training) ---
    if master:
        log0("--- Compression & Serialization ---")
        method, model_bytes, code_bytes, stats = save_model_quantized(
            base_model, code, "final_model", bits=args.quant_bits, fp_storage=args.fp_storage)
        total = model_bytes + code_bytes
        compressors = " | ".join(
            f"{m}:{r['compressed_mb']:.2f}MB" for m, r in stats["compression_results"].items())
        log0(f"  Params: {stats['int_params']+stats['fp_params']:,} | Discarded: {stats['discarded_params']:,}")
        log0(f"  Compressors: {compressors}")
        fp_label = f"+{args.fp_storage}" if args.fp_storage == "fp8" else ""
        log0(f"  INT{args.quant_bits}{fp_label}+{method}: {model_bytes/1e6:.2f}MB model + {code_bytes/1e6:.3f}MB code "
             f"= {total/1e6:.2f}/16.00MB {'FITS' if total <= 16_000_000 else 'OVER'}")

    # --- Roundtrip evaluation ---
    if distributed:
        dist.barrier()

    candidates = sorted(glob.glob("final_model.*.ptz"))
    load_path = candidates[0] if candidates else "final_model.lzma.ptz"
    load_model_quantized(load_path, base_model, device)
    torch._dynamo.reset()

    q_loss, q_bpb = eval_val(args, model, rank, world_size, device,
                              val_tokens, bpe_luts=bpe_luts)
    log0(f"final_roundtrip val_loss:{q_loss:.4f} val_bpb:{q_bpb:.4f}")

    # Temperature search (calibrate on TRAINING data, not validation)
    opt_temp = 1.0
    if args.temp_scaling:
        torch.cuda.synchronize()
        t_temp = time.perf_counter()
        # Take a small chunk of training data for calibration (same as ternary code)
        calib_tokens = train_loader.stream.take(65536)
        # Pad to seq_len alignment
        u = ((calib_tokens.numel() - 1) // args.train_seq_len) * args.train_seq_len
        calib_tokens = calib_tokens[:u + 1]
        opt_temp = find_best_temperature(args, model, rank, world_size,
                                         device, calib_tokens, bpe_luts=bpe_luts)
        torch.cuda.synchronize()
        log0(f"temp_scaling optimal_T:{opt_temp:.2f} "
             f"time:{1000.0 * (time.perf_counter() - t_temp):.0f}ms")
        # Final eval with optimal temperature (on val data — this is the actual evaluation)
        t_loss, t_bpb = eval_val(args, model, rank, world_size, device,
                                  val_tokens, temperature=opt_temp, bpe_luts=bpe_luts)
        log0(f"final_temped val_loss:{t_loss:.4f} val_bpb:{t_bpb:.4f} T={opt_temp:.2f}")

    # Sliding window eval
    if args.sliding_eval:
        torch.cuda.synchronize()
        t_sw = time.perf_counter()
        temp = opt_temp if args.temp_scaling else 1.0
        sw_loss, sw_bpb = eval_val_sliding(args, model, rank, world_size,
                                            device, val_tokens,
                                            stride=args.sliding_eval_stride,
                                            temperature=temp, bpe_luts=bpe_luts)
        torch.cuda.synchronize()
        log0(f"final_sliding val_loss:{sw_loss:.4f} val_bpb:{sw_bpb:.4f} "
             f"(stride={args.sliding_eval_stride}, T={temp:.2f}) "
             f"time:{1000.0 * (time.perf_counter() - t_sw):.0f}ms")

    if distributed:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()