train_gpt.py

from __future__ import annotations
import copy
import glob
import io
import lzma
import math
import os
import random
import subprocess, sys
import time
import uuid
from pathlib import Path
import numpy as np, 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
try:
    from flash_attn import flash_attn_func
    _HAS_FLASH_ATTN = True
except ImportError:
    _HAS_FLASH_ATTN = False
class Hyperparameters:
    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))
    val_batch_size = int(os.environ.get("VAL_BATCH_SIZE", 524_288))
    val_loss_every = int(os.environ.get("VAL_LOSS_EVERY", 4000))
    train_log_every = int(os.environ.get("TRAIN_LOG_EVERY", 500))
    iterations = int(os.environ.get("ITERATIONS", 20000))
    warmdown_iters = int(os.environ.get("WARMDOWN_ITERS", 3500))
    warmup_steps = int(os.environ.get("WARMUP_STEPS", 20))
    train_batch_tokens = int(os.environ.get("TRAIN_BATCH_TOKENS", 786_432))
    train_seq_len = int(os.environ.get("TRAIN_SEQ_LEN", 2048))
    eval_seq_len = int(os.environ.get("EVAL_SEQ_LEN", 2048))
    eval_stride = int(os.environ.get("EVAL_STRIDE", 64))
    eval_batch_seqs = int(os.environ.get("EVAL_BATCH_SEQS", 32))
    max_wallclock_seconds = float(os.environ.get("MAX_WALLCLOCK_SECONDS", 600.0))
    qk_gain_init = float(os.environ.get("QK_GAIN_INIT", 1.5))
    vocab_size = int(os.environ.get("VOCAB_SIZE", 1024))
    num_layers = int(os.environ.get("NUM_LAYERS", 11))
    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", 8))
    mlp_mult = int(os.environ.get("MLP_MULT", 3))
    mlp_hidden = int(os.environ.get("MLP_HIDDEN", 1536))
    rope_base = float(os.environ.get("ROPE_BASE", 10000.0))
    logit_softcap = float(os.environ.get("LOGIT_SOFTCAP", 30.0))
    bigram_vocab_size = int(os.environ.get("BIGRAM_VOCAB_SIZE", 3072))
    bigram_dim = int(os.environ.get("BIGRAM_DIM", 112))
    tied_embed_lr = float(os.environ.get("TIED_EMBED_LR", 0.035))
    tied_embed_init_std = float(os.environ.get("TIED_EMBED_INIT_STD", 0.005))
    matrix_lr = float(os.environ.get("MATRIX_LR", 0.025))
    scalar_lr = float(os.environ.get("SCALAR_LR", 0.025))
    muon_momentum = float(os.environ.get("MUON_MOMENTUM", 0.99))
    muon_backend_steps = int(os.environ.get("MUON_BACKEND_STEPS", 5))
    muon_momentum_warmup_start = float(os.environ.get("MUON_MOMENTUM_WARMUP_START", 0.92))
    muon_momentum_warmup_steps = int(os.environ.get("MUON_MOMENTUM_WARMUP_STEPS", 1500))
    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))
    muon_weight_decay = float(os.environ.get("MUON_WD", 0.04))
    adam_weight_decay = float(os.environ.get("ADAM_WD", 0.04))
    grad_clip_norm = float(os.environ.get("GRAD_CLIP_NORM", 0.3))
    xsa_last_n = int(os.environ.get("XSA_LAST_N", 11))
    ema_enabled = bool(int(os.environ.get("EMA_ENABLED", 1)))
    ema_decay = float(os.environ.get("EMA_DECAY", 0.997))
    partial_rope_dims = int(os.environ.get("PARTIAL_ROPE_DIMS", 16))
    ln_scale = bool(int(os.environ.get("LN_SCALE", 1)))
    ve_enabled = bool(int(os.environ.get("VE_ENABLED", "1")))
    ve_dim = int(os.environ.get("VE_DIM", 128))
    ve_layers = os.environ.get("VE_LAYERS", "9,10")
    negative_slope = float(os.environ.get("NEGATIVE_SLOPE", 0.5))
    attn_pattern = os.environ.get("ATTN_PATTERN", "")
    local_attn_window = int(os.environ.get("LOCAL_ATTN_WINDOW", 512))
    train_loader_mode = os.environ.get("TRAIN_LOADER_MODE", "coprime_multi_shard")
    loader_shards_per_rank = int(os.environ.get("LOADER_SHARDS_PER_RANK", 4))
    gptq_enabled = bool(int(os.environ.get("GPTQ_ENABLED", "1")))
    gptq_bits = int(os.environ.get("GPTQ_BITS", 6))
    gptq_group_size = int(os.environ.get("GPTQ_GROUP_SIZE", 128))
    gptq_calibration_steps = int(os.environ.get("GPTQ_CALIBRATION_STEPS", 64))
    gptq_reserve_ms = float(os.environ.get("GPTQ_RESERVE_MS", 14000))
    gptq_guard_ms = float(os.environ.get("GPTQ_GUARD_MS", 4000))
    gptq_hessian_damp = float(os.environ.get("GPTQ_HESSIAN_DAMP", 1e-4))
    gptq_target_patterns = tuple(
        pattern for pattern in os.environ.get(
            "GPTQ_TARGET_PATTERNS",
            ".attn.c_q.weight,.attn.c_k.weight,.attn.c_v.weight,.attn.proj.weight,.mlp.fc.weight,.mlp.proj.weight,bigram.proj.weight",
        ).split(",") if pattern
    )
    selective_prune_enabled = bool(int(os.environ.get("SELECTIVE_PRUNE_ENABLED", "1")))
    target_total_bytes = int(os.environ.get("TARGET_TOTAL_BYTES", 15_900_000))
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, weight_decay: float = 0.0):
        super().__init__(
            params,
            dict(lr=lr, momentum=momentum, backend_steps=backend_steps, nesterov=nesterov, weight_decay=weight_decay),
        )
    @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)
            wd = group.get("weight_decay", 0.0)
            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)
                if wd > 0:
                    p.data.mul_(1.0 - lr * wd)
                curr += p.numel()
        return loss
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,
    eval_seq_len: int | None = None,
) -> tuple[float, float]:
    seq_len = eval_seq_len or args.train_seq_len
    local_batch_tokens = args.val_batch_size // (world_size * grad_accum_steps)
    if local_batch_tokens < 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}, seq_len={seq_len}"
        )
    local_batch_seqs = local_batch_tokens // seq_len
    total_seqs = (val_tokens.numel() - 1) // 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 * seq_len
            raw_end = batch_seq_end * seq_len + 1
            local = val_tokens[raw_start:raw_end].to(device=device, dtype=torch.int64, non_blocking=True)
            x = local[:-1].reshape(-1, seq_len)
            y = local[1:].reshape(-1, 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)
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,skip_weight,skip_weights,smear,bigram.scale,ve_layer_scales,ve_shared.scale",
    ).split(",")
    if pattern
)
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 quantize_float_tensor(t: Tensor, bits: int = 8) -> tuple[Tensor, Tensor]:
    max_val = 127 if bits == 8 else (2 ** (bits - 1)) - 1  # int6: 31, int8: 127
    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 / float(max_val)).clamp_min(1.0 / float(max_val))
        q = torch.clamp(torch.round(clipped / scale[:, None]), -max_val, max_val).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 / float(max_val) if clip_abs > 0 else 1.0, dtype=torch.float32)
    q = torch.clamp(torch.round(torch.clamp(t32, -clip_abs, clip_abs) / scale), -max_val, max_val).to(torch.int8).contiguous()
    return q, scale
def pack_nbits(q: Tensor, bits: int) -> Tensor:
    vals_per_chunk = 8 // math.gcd(8, bits)
    bytes_per_chunk = vals_per_chunk * bits // 8
    max_q = (1 << bits) - 1
    flat = q.reshape(-1).to(torch.int64)
    if flat.numel() % vals_per_chunk:
        flat = F.pad(flat, (0, vals_per_chunk - flat.numel() % vals_per_chunk))
    flat = flat.view(-1, vals_per_chunk)
    acc = torch.zeros(flat.size(0), dtype=torch.int64)
    for i in range(vals_per_chunk):
        acc |= (flat[:, i] & max_q) << (bits * i)
    return torch.stack([((acc >> (8 * i)) & 0xFF).to(torch.uint8) for i in range(bytes_per_chunk)], dim=1).reshape(-1)
def unpack_nbits(blob: Tensor, bits: int, numel: int) -> Tensor:
    vals_per_chunk = 8 // math.gcd(8, bits)
    bytes_per_chunk = vals_per_chunk * bits // 8
    max_q = (1 << bits) - 1
    packed = blob.view(-1, bytes_per_chunk).to(torch.int64)
    acc = torch.zeros(packed.size(0), dtype=torch.int64)
    for i in range(bytes_per_chunk):
        acc |= packed[:, i] << (8 * i)
    vals = torch.stack([((acc >> (bits * i)) & max_q).to(torch.uint8) for i in range(vals_per_chunk)], dim=1).reshape(-1)
    return vals[:numel]
def coprime_step(size: int, seed: int) -> int:
    if size <= 1:
        return 1
    step = max(1, seed % size)
    while math.gcd(step, size) != 1:
        step += 1
    return step
def apply_selective_pm1_prune(q_signed: Tensor, w_ref: Tensor) -> tuple[Tensor, int]:
    if q_signed.numel() == 0:
        return q_signed, 0
    row_energy = w_ref.float().square().mean(dim=1)
    cutoff = torch.quantile(row_energy, 0.25)
    prune_rows = row_energy <= cutoff
    prune_mask = prune_rows[:, None] & (q_signed.abs() <= 2) & (q_signed != 0)
    pruned = torch.where(prune_mask, q_signed.sign(), q_signed)
    return pruned, int(prune_mask.sum().item())
def quantize_gptq_tensor(
    t: Tensor,
    hessian: Tensor,
    bits: int,
    group_size: int,
    hessian_damp: float,
    selective_prune: bool,
) -> tuple[Tensor, Tensor, int]:
    maxq = (1 << (bits - 1)) - 1
    minq = -(1 << (bits - 1))
    w = t.float().contiguous().clone()
    h = hessian.float().contiguous().clone()
    if h.shape != (w.shape[1], w.shape[1]):
        raise ValueError(f"Hessian shape mismatch for GPTQ tensor: {tuple(h.shape)} vs {(w.shape[1], w.shape[1])}")
    diag = torch.diagonal(h)
    dead = diag <= 0
    if dead.any(): h[dead, dead] = 1.0; w[:, dead] = 0.0
    damp = max(float(diag.mean().item()) * hessian_damp, 1e-6)
    h = 0.5 * (h + h.T)
    eye = torch.eye(h.shape[0], dtype=h.dtype)
    h_inv = None
    for mult in (1.0, 10.0, 100.0, 1000.0):
        try:
            h_inv = torch.cholesky_inverse(torch.linalg.cholesky(h + eye * (damp * mult))).contiguous(); break
        except torch._C._LinAlgError:
            pass
    if h_inv is None:
        evals, evecs = torch.linalg.eigh(h)
        h_inv = ((evecs * evals.clamp_min(damp).reciprocal()) @ evecs.T).contiguous()
    rows, cols = w.shape
    num_groups = (cols + group_size - 1) // group_size
    scales = torch.empty((rows, num_groups), dtype=torch.float16)
    q_signed = torch.empty((rows, cols), dtype=torch.int16)
    prune_count = 0
    for gidx, start in enumerate(range(0, cols, group_size)):
        end = min(start + group_size, cols)
        w_group = w[:, start:end]
        scale = (w_group.abs().amax(dim=1) / float(maxq)).clamp_min(1.0 / float(maxq))
        scales[:, gidx] = scale.to(torch.float16)
        for col in range(start, end):
            q_col = torch.round(w[:, col] / scale).clamp(minq, maxq).to(torch.int16)
            if selective_prune:
                q_col, added = apply_selective_pm1_prune(q_col[:, None], w[:, col:col + 1])
                q_col = q_col[:, 0]
                prune_count += added
            q_signed[:, col] = q_col
            deq_col = q_col.float() * scale
            err = (w[:, col] - deq_col) / h_inv[col, col].clamp_min(1e-6)
            if col + 1 < cols:
                w[:, col + 1:] -= err[:, None] * h_inv[col, col + 1:][None, :]
    q_unsigned = (q_signed + (1 << (bits - 1))).to(torch.uint8)
    return pack_nbits(q_unsigned, bits), scales.contiguous(), prune_count
def dequantize_gptq_tensor(blob: Tensor, scales: Tensor, shape: tuple[int, int], bits: int, group_size: int) -> Tensor:
    rows, cols = shape
    q_unsigned = unpack_nbits(blob, bits, rows * cols).view(rows, cols).to(torch.int16)
    q_signed = q_unsigned - (1 << (bits - 1))
    out = torch.empty((rows, cols), dtype=torch.float32)
    for gidx, start in enumerate(range(0, cols, group_size)):
        end = min(start + group_size, cols)
        out[:, start:end] = q_signed[:, start:end].float() * scales[:, gidx].float()[:, None]
    return out
def quantize_state_dict_for_export(
    state_dict: dict[str, Tensor],
    gptq_hessians: dict[str, Tensor],
    args: Hyperparameters,
) -> tuple[dict[str, object], dict[str, int]]:
    quantized: dict[str, Tensor] = {}
    scales: dict[str, Tensor] = {}
    dtypes: dict[str, str] = {}
    passthrough: dict[str, Tensor] = {}
    qmeta: dict[str, dict[str, object]] = {}
    packed: dict[str, Tensor] = {}
    packed_scales: dict[str, Tensor] = {}
    packed_shapes: dict[str, tuple[int, int]] = {}
    packed_meta: dict[str, dict[str, int]] = {}
    stats = {
        "param_count": 0,
        "baseline_tensor_bytes": 0,
        "int8_payload_bytes": 0,
        "packed_payload_bytes": 0,
        "pruned_codes": 0,
    }
    for name, tensor in state_dict.items():
        t = tensor.detach().cpu().contiguous()
        stats["param_count"] += int(t.numel())
        stats["baseline_tensor_bytes"] += tensor_nbytes(t)
        if not t.is_floating_point():
            passthrough[name] = t
            stats["int8_payload_bytes"] += tensor_nbytes(t)
            continue
        dtypes[name] = str(t.dtype).removeprefix("torch.")
        if any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) or t.numel() <= 65_536:
            passthrough[name] = t.float().contiguous() if any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS) else t.to(torch.float16).contiguous()
            stats["int8_payload_bytes"] += tensor_nbytes(passthrough[name])
            continue
        if (
            args.gptq_enabled
            and t.ndim == 2
            and any(pattern in name for pattern in args.gptq_target_patterns)
            and name in gptq_hessians
        ):
            blob, s, pruned = quantize_gptq_tensor(
                t,
                gptq_hessians[name],
                bits=args.gptq_bits,
                group_size=args.gptq_group_size,
                hessian_damp=args.gptq_hessian_damp,
                selective_prune=args.selective_prune_enabled,
            )
            packed[name] = blob
            packed_scales[name] = s
            packed_shapes[name] = tuple(t.shape)
            packed_meta[name] = {"bits": args.gptq_bits, "group_size": args.gptq_group_size}
            stats["packed_payload_bytes"] += tensor_nbytes(blob) + tensor_nbytes(s)
            stats["pruned_codes"] += pruned
            continue
        q, s = quantize_float_tensor(t, bits=8)
        quantized[name] = q
        scales[name] = s
        if s.ndim > 0:
            qmeta[name] = {"scheme": "per_row", "axis": 0}
        stats["int8_payload_bytes"] += tensor_nbytes(q) + tensor_nbytes(s)
    obj: dict[str, object] = {
        "__quant_format__": "gptq6_lzma_v1",
        "quantized": quantized,
        "scales": scales,
        "dtypes": dtypes,
        "passthrough": passthrough,
        "packed": packed,
        "packed_scales": packed_scales,
        "packed_shapes": packed_shapes,
        "packed_meta": packed_meta,
    }
    if qmeta:
        obj["qmeta"] = qmeta
    return obj, stats
def dequantize_state_dict_for_export(obj: dict[str, object]) -> dict[str, Tensor]:
    out: dict[str, Tensor] = {}
    qmeta = obj.get("qmeta", {})
    for name, blob in obj.get("packed", {}).items():
        dtype = getattr(torch, obj["dtypes"][name])
        shape = tuple(obj["packed_shapes"][name])
        bits = int(obj["packed_meta"][name]["bits"])
        group_size = int(obj["packed_meta"][name]["group_size"])
        w = dequantize_gptq_tensor(blob, obj["packed_scales"][name], shape, bits, group_size)
        out[name] = w.to(dtype=dtype).contiguous()
    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:
            out[name] = (q.float() * s.float().view(q.shape[0], *([1] * (q.ndim - 1)))).to(dtype=dtype).contiguous()
        else:
            out[name] = (q.float() * float(s.item())).to(dtype=dtype).contiguous()
    for name, t in obj["passthrough"].items():
        dtype = getattr(torch, obj["dtypes"].get(name, str(t.dtype).removeprefix("torch.")))
        out[name] = t.to(dtype=dtype).contiguous()
    return out
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, shuffle: bool = False, seed: int = 1337):
        self.files = [Path(p) for p in sorted(glob.glob(pattern))]
        if not self.files:
            raise FileNotFoundError(f"No files found for pattern: {pattern}")
        if shuffle:
            rng = random.Random(seed)
            rng.shuffle(self.files)
        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 SequentialTokenLoader:
    mode = "sequential"
    def __init__(self, pattern: str, rank: int, world_size: int, device: torch.device, seed: int):
        self.rank = rank
        self.world_size = world_size
        self.device = device
        self.stream = TokenStream(pattern, shuffle=False, seed=seed)
    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)
    def describe(self) -> str:
        return "mode:sequential"
class CoprimeMultiShardLoader:
    mode = "coprime_multi_shard"
    def __init__(self, pattern: str, rank: int, world_size: int, device: torch.device, seed: int, shards_per_rank: int):
        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.rank = rank
        self.world_size = world_size
        self.device = device
        self.seed = seed
        self.shards_per_rank = max(1, min(shards_per_rank, len(self.files)))
        self.shard_cursor = (seed + rank) % len(self.files)
        self.shard_stride = coprime_step(len(self.files), seed + 2 * rank + 1)
        self.shard_seq_pos: dict[int, int] = {}
        self.shard_seq_stride: dict[int, int] = {}
        self.cache: dict[int, Tensor] = {}
        self.cache_order: list[int] = []
        self.max_cache = max(2 * self.shards_per_rank, 4)
    def _get_shard(self, shard_idx: int) -> Tensor:
        if shard_idx in self.cache:
            return self.cache[shard_idx]
        tokens = load_data_shard(self.files[shard_idx])
        self.cache[shard_idx] = tokens
        self.cache_order.append(shard_idx)
        while len(self.cache_order) > self.max_cache:
            old = self.cache_order.pop(0)
            self.cache.pop(old, None)
        return tokens
    def _next_seq(self, shard_idx: int, seq_len: int) -> Tensor:
        tokens = self._get_shard(shard_idx)
        total_seqs = max((tokens.numel() - 1) // seq_len, 1)
        if shard_idx not in self.shard_seq_stride:
            self.shard_seq_stride[shard_idx] = coprime_step(total_seqs, self.seed + shard_idx + self.rank + 1)
            self.shard_seq_pos[shard_idx] = (self.seed + shard_idx) % total_seqs
        pos = self.shard_seq_pos[shard_idx]
        self.shard_seq_pos[shard_idx] = (pos + self.shard_seq_stride[shard_idx]) % total_seqs
        start = pos * seq_len
        return tokens[start:start + seq_len + 1]
    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)
        local_batch_seqs = local_tokens // seq_len
        shard_ids = []
        for _ in range(self.shards_per_rank):
            shard_ids.append(self.shard_cursor)
            self.shard_cursor = (self.shard_cursor + self.shard_stride) % len(self.files)
        x = torch.empty((local_batch_seqs, seq_len), dtype=torch.int64)
        y = torch.empty((local_batch_seqs, seq_len), dtype=torch.int64)
        for i in range(local_batch_seqs):
            chunk = self._next_seq(shard_ids[i % len(shard_ids)], seq_len).to(torch.int64)
            x[i] = chunk[:-1]
            y[i] = chunk[1:]
        return x.to(self.device, non_blocking=True), y.to(self.device, non_blocking=True)
    def describe(self) -> str:
        return f"mode:coprime_multi_shard shards_per_rank:{self.shards_per_rank} shard_stride:{self.shard_stride}"
def build_train_loader(args: Hyperparameters, rank: int, world_size: int, device: torch.device):
    if args.train_loader_mode == "sequential":
        return SequentialTokenLoader(args.train_files, rank, world_size, device, seed=args.seed)
    if args.train_loader_mode == "coprime_multi_shard":
        return CoprimeMultiShardLoader(
            args.train_files,
            rank,
            world_size,
            device,
            seed=args.seed,
            shards_per_rank=args.loader_shards_per_rank,
        )
    raise ValueError(f"Unknown TRAIN_LOADER_MODE={args.train_loader_mode}")
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, train_seq_len: int = 1024):
        super().__init__()
        self.dim = dim
        self.base = base
        self.train_seq_len = train_seq_len
        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
        ):
            if seq_len > self.train_seq_len:
                scale = seq_len / self.train_seq_len
                new_base = self.base * (scale ** (self.dim / (self.dim - 2)))
                inv_freq = 1.0 / (new_base ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim))
            else:
                inv_freq = self.inv_freq.to(device)
            t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
            freqs = torch.outer(t, inv_freq)
            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, partial_dims: int = 0) -> Tensor:
    if partial_dims > 0 and partial_dims < x.size(-1):
        x_rot = x[..., :partial_dims]
        x_pass = x[..., partial_dims:]
        half = partial_dims // 2
        x1, x2 = x_rot[..., :half], x_rot[..., half:]
        x_rot = torch.cat((x1 * cos[..., :half] + x2 * sin[..., :half],
                          x1 * (-sin[..., :half]) + x2 * cos[..., :half]), dim=-1)
        return torch.cat((x_rot, x_pass), dim=-1)
    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
        self.q_gain = nn.Parameter(torch.full((num_heads,), qk_gain_init, dtype=torch.float32))
        self.partial_rope_dims = 0
        rope_dim = self.head_dim
        self.rotary = Rotary(rope_dim, base=rope_base, train_seq_len=1024)
        self.use_xsa = False
        self.attn_mode = "global"
        self.local_attn_window = 0
        self._local_mask: Tensor | None = None
        self._local_mask_seq_len = 0
    def _get_local_mask(self, seqlen: int, device: torch.device) -> Tensor:
        if self._local_mask is None or self._local_mask_seq_len != seqlen or self._local_mask.device != device:
            idx = torch.arange(seqlen, device=device)
            dist = idx[:, None] - idx[None, :]
            self._local_mask = (dist >= 0) & (dist < self.local_attn_window)
            self._local_mask_seq_len = seqlen
        return self._local_mask
    def _xsa_efficient(self, y: Tensor, v: Tensor) -> Tensor:
        B, T, H, D = y.shape
        Hkv = v.size(-2)
        group = H // Hkv
        y_g = y.reshape(B, T, Hkv, group, D)
        vn = F.normalize(v, dim=-1).unsqueeze(-2)
        proj = (y_g * vn).sum(dim=-1, keepdim=True) * vn
        return (y_g - proj).reshape(B, T, H, D)
    def forward(self, x: Tensor, v_embed: Tensor | None = None) -> 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_raw = self.c_v(x)
        if v_embed is not None:
            v_raw = v_raw + v_embed.to(dtype=v_raw.dtype)
        v = v_raw.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, partial_dims=self.partial_rope_dims)
        k = apply_rotary_emb(k, cos, sin, partial_dims=self.partial_rope_dims)
        q = q * self.q_gain.to(dtype=q.dtype)[None, :, None, None]
        if _HAS_FLASH_ATTN:
            q_fa = q.transpose(1, 2)
            k_fa = k.transpose(1, 2)
            v_fa = v.transpose(1, 2)
            y = flash_attn_func(
                q_fa, k_fa, v_fa, causal=True,
                window_size=(self.local_attn_window - 1, 0) if self.attn_mode == "local" else (-1, -1),
            )
        else:
            y = F.scaled_dot_product_attention(
                q, k, v,
                attn_mask=None if self.attn_mode == "global" else self._get_local_mask(seqlen, x.device),
                is_causal=self.attn_mode == "global",
                enable_gqa=(self.num_kv_heads != self.num_heads),
            )
            y = y.transpose(1, 2)
        if self.use_xsa:
            v_bshd = v.transpose(1, 2)  # XSA needs (B, S, Hkv, D)
            y = self._xsa_efficient(y, v_bshd)
        y = y.reshape(bsz, seqlen, dim)
        return self.proj(y)
class MLP(nn.Module):
    def __init__(self, dim: int, mlp_mult: int, mlp_hidden: int = 0, negative_slope: float = 0.5):
        super().__init__()
        hidden = mlp_hidden if mlp_hidden > 0 else mlp_mult * dim
        self.fc = CastedLinear(dim, hidden, bias=False)
        self.proj = CastedLinear(hidden, dim, bias=False)
        self.proj._zero_init = True
        self.negative_slope = negative_slope
    def forward(self, x: Tensor) -> Tensor:
        x = F.leaky_relu(self.fc(x), negative_slope=self.negative_slope)
        return self.proj(x.square())
class SmearGate(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.gate = nn.Parameter(torch.zeros(dim, dtype=torch.float32))
    def forward(self, x: Tensor) -> Tensor:
        g = torch.sigmoid(self.gate.to(dtype=x.dtype))[None, None, :]
        x_prev = torch.cat([torch.zeros_like(x[:, :1]), x[:, :-1]], dim=1)
        return (1 - g) * x + g * x_prev
class BigramHashEmbedding(nn.Module):
    def __init__(self, bigram_vocab_size: int, bigram_dim: int, model_dim: int):
        super().__init__()
        self.bigram_vocab_size = bigram_vocab_size
        self.embed = nn.Embedding(bigram_vocab_size, bigram_dim)
        nn.init.zeros_(self.embed.weight)
        self.proj = CastedLinear(bigram_dim, model_dim, bias=False) if bigram_dim != model_dim else None
        if self.proj is not None:
            nn.init.zeros_(self.proj.weight)
        self.scale = nn.Parameter(torch.tensor(0.05, dtype=torch.float32))
    def bigram_hash(self, tokens: Tensor) -> Tensor:
        t = tokens.to(torch.int32)
        mod = self.bigram_vocab_size - 1
        out = torch.empty_like(t)
        out[..., 0] = mod
        out[..., 1:] = torch.bitwise_xor(36313 * t[..., 1:], 27191 * t[..., :-1]) % mod
        return out.long()
    def forward(self, token_ids: Tensor) -> Tensor:
        h = self.embed(self.bigram_hash(token_ids))
        if self.proj is not None:
            h = self.proj(h)
        return h * self.scale.to(dtype=h.dtype)
class ValueEmbedding(nn.Module):
    def __init__(self, vocab_size: int, ve_dim: int, model_dim: int):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, ve_dim)
        nn.init.normal_(self.embed.weight, std=0.01)
        self.proj = CastedLinear(ve_dim, model_dim, bias=False) if ve_dim != model_dim else None
        if self.proj is not None:
            nn.init.zeros_(self.proj.weight)
        self.scale = nn.Parameter(torch.tensor(0.1, dtype=torch.float32))
    def forward(self, token_ids: Tensor) -> Tensor:
        h = self.embed(token_ids)
        if self.proj is not None:
            h = self.proj(h)
        return h * self.scale.to(dtype=h.dtype)
class Block(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        num_kv_heads: int,
        mlp_mult: int,
        rope_base: float,
        qk_gain_init: float,
        mlp_hidden: int = 0,
        negative_slope: float = 0.5,
    ):
        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 = MLP(dim, mlp_mult, mlp_hidden, negative_slope=negative_slope)
        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())
        self.ln_scale_factor = 1.0  # Set externally for LN Scale
    def forward(self, x: Tensor, x0: Tensor, v_embed: Tensor | None = None) -> Tensor:
        mix = self.resid_mix.to(dtype=x.dtype)
        x = mix[0][None, None, :] * x + mix[1][None, None, :] * x0
        s = self.ln_scale_factor
        attn_out = self.attn(self.attn_norm(x) * s, v_embed=v_embed)
        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) * s)
        return x
class GPT(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        num_layers: int,
        model_dim: int,
        num_heads: int,
        num_kv_heads: int,
        mlp_mult: int,
        mlp_hidden: int,
        tied_embed_init_std: float,
        logit_softcap: float,
        rope_base: float,
        qk_gain_init: float,
        bigram_vocab_size: int = 0,
        bigram_dim: int = 128,
        xsa_last_n: int = 0,
        partial_rope_dims: int = 0,
        ln_scale: bool = False,
        ve_enabled: bool = False,
        ve_dim: int = 128,
        ve_layers: str = "9,10",
        negative_slope: float = 0.5,
        attn_pattern: str = "",
        local_attn_window: int = 512,
    ):
        super().__init__()
        if logit_softcap <= 0.0:
            raise ValueError(f"logit_softcap must be positive, got {logit_softcap}")
        self.model_dim = model_dim
        self.tied_embed_init_std = tied_embed_init_std
        self.logit_softcap = logit_softcap
        self.tok_emb = nn.Embedding(vocab_size, model_dim)
        self.bigram = BigramHashEmbedding(bigram_vocab_size, bigram_dim, model_dim) if bigram_vocab_size > 0 else None
        self.smear = SmearGate(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))
        head_dim = model_dim // num_heads
        kv_dim = num_kv_heads * head_dim
        self.blocks = nn.ModuleList(
            [
                Block(
                    model_dim,
                    num_heads,
                    num_kv_heads,
                    mlp_mult,
                    rope_base,
                    qk_gain_init,
                    mlp_hidden=mlp_hidden,
                    negative_slope=negative_slope,
                )
                for _ in range(num_layers)
            ]
        )
        self.ve_layer_indices = [int(x) for x in ve_layers.split(",") if x.strip()] if ve_enabled else []
        if self.ve_layer_indices:
            self.ve_shared = ValueEmbedding(vocab_size, ve_dim, kv_dim)
            self.ve_layer_scales = nn.ParameterList(
                [nn.Parameter(torch.ones(1, dtype=torch.float32)) for _ in self.ve_layer_indices]
            )
        else:
            self.ve_shared = None
            self.ve_layer_scales = nn.ParameterList()
        attn_modes = [part.strip().lower() for part in attn_pattern.split(",") if part.strip()] if attn_pattern else []
        if not attn_modes:
            attn_modes = ["global"] * num_layers
        if len(attn_modes) != num_layers:
            raise ValueError(f"ATTN_PATTERN must have {num_layers} entries, got {len(attn_modes)}")
        for i in range(num_layers):
            if xsa_last_n > 0 and i >= num_layers - xsa_last_n:
                self.blocks[i].attn.use_xsa = True
            if partial_rope_dims > 0:
                self.blocks[i].attn.partial_rope_dims = partial_rope_dims
                self.blocks[i].attn.rotary = Rotary(partial_rope_dims, base=rope_base, train_seq_len=1024)
            if ln_scale:
                self.blocks[i].ln_scale_factor = 1.0 / math.sqrt(i + 1)
            mode = attn_modes[i]
            if mode not in {"g", "global", "l", "local"}:
                raise ValueError(f"Unsupported attention mode {mode!r} in ATTN_PATTERN")
            self.blocks[i].attn.attn_mode = "local" if mode.startswith("l") else "global"
            self.blocks[i].attn.local_attn_window = local_attn_window
        self.final_norm = RMSNorm()
        self._init_weights()
    def _get_ve(self, layer_idx: int, input_ids: Tensor, cache: dict[str, Tensor]) -> Tensor | None:
        if self.ve_shared is None or layer_idx not in self.ve_layer_indices:
            return None
        if "ve" not in cache:
            cache["ve"] = self.ve_shared(input_ids)
        scale_idx = self.ve_layer_indices.index(layer_idx)
        return cache["ve"] * self.ve_layer_scales[scale_idx].to(dtype=cache["ve"].dtype)
    def _init_weights(self) -> None:
        nn.init.normal_(self.tok_emb.weight, mean=0.0, std=self.tied_embed_init_std)
        num_layers = len(self.blocks)
        for name, module in self.named_modules():
            if isinstance(module, nn.Linear):
                if getattr(module, "_zero_init", False):
                    nn.init.zeros_(module.weight)
                elif module.weight.ndim == 2 and module.weight.shape[0] >= 64 and module.weight.shape[1] >= 64:
                    nn.init.orthogonal_(module.weight, gain=1.0)
                    if ".proj." in name or name.endswith(".proj"):
                        with torch.no_grad():
                            module.weight.mul_(1.0 / math.sqrt(2 * num_layers))
    def _project_logits(self, x: Tensor) -> Tensor:
        logits_proj = F.linear(x, self.tok_emb.weight)
        return self.logit_softcap * torch.tanh(logits_proj / self.logit_softcap)
    def _hidden(self, input_ids: Tensor) -> Tensor:
        x = self.tok_emb(input_ids)
        if self.bigram is not None:
            x = x + self.bigram(input_ids)
        x = F.rms_norm(x, (x.size(-1),))
        x = self.smear(x)
        x0 = x
        skips: list[Tensor] = []
        ve_cache: dict[str, Tensor] = {}
        for i in range(self.num_encoder_layers):
            x = self.blocks[i](x, x0, v_embed=self._get_ve(i, input_ids, ve_cache))
            skips.append(x)
        for i in range(self.num_decoder_layers):
            block_idx = self.num_encoder_layers + i
            if skips:
                x = x + self.skip_weights[i].to(dtype=x.dtype)[None, None, :] * skips.pop()
            x = self.blocks[block_idx](x, x0, v_embed=self._get_ve(block_idx, input_ids, ve_cache))
        return self.final_norm(x)
    def forward(self, input_ids: Tensor, target_ids: Tensor) -> Tensor:
        x = self._hidden(input_ids).reshape(-1, self.model_dim)
        targets = target_ids.reshape(-1)
        return F.cross_entropy(self._project_logits(x).float(), targets, reduction="mean")
    def get_logits(self, input_ids: Tensor) -> Tensor:
        return self._project_logits(self._hidden(input_ids))
def eval_val_sliding(
    args, base_model: nn.Module, rank: int, world_size: int, device: torch.device,
    val_tokens: Tensor, base_bytes_lut: Tensor, has_leading_space_lut: Tensor,
    is_boundary_token_lut: Tensor, eval_seq_len: int, eval_stride: int,
    batch_seqs: int = 32,
) -> tuple[float, float]:
    total_tokens = val_tokens.numel() - 1
    all_starts = [s for s in range(0, total_tokens - eval_stride + 1, eval_stride)
                  if min(s + eval_seq_len, total_tokens) - s >= eval_stride]
    n = len(all_starts)
    my_s = (n * rank) // world_size
    my_e = (n * (rank + 1)) // world_size
    my_starts = all_starts[my_s:my_e]
    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)
    base_model.eval()
    with torch.inference_mode():
        for bi in range(0, len(my_starts), batch_seqs):
            batch_ws = my_starts[bi:bi + batch_seqs]
            bsz = len(batch_ws)
            x_batch = torch.zeros(bsz, eval_seq_len, dtype=torch.int64, device=device)
            y_batch = torch.zeros(bsz, eval_seq_len, dtype=torch.int64, device=device)
            wlens: list[int] = []
            for i, ws in enumerate(batch_ws):
                end = min(ws + eval_seq_len, total_tokens)
                wlen = end - ws
                wlens.append(wlen)
                chunk = val_tokens[ws:end + 1].to(dtype=torch.int64, device=device)
                x_batch[i, :wlen] = chunk[:-1]
                y_batch[i, :wlen] = chunk[1:]
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                logits = base_model.get_logits(x_batch)
            nll = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)).float(),
                y_batch.reshape(-1),
                reduction="none",
            ).reshape(bsz, eval_seq_len)
            for i, ws in enumerate(batch_ws):
                wlen = wlens[i]
                score_from = max(wlen - eval_stride, 0) if ws > 0 else 0
                losses = nll[i, score_from:wlen]
                val_loss_sum += losses.to(torch.float64).sum()
                val_token_count += losses.numel()
                prev_ids = x_batch[i, score_from:wlen]
                tgt_ids = y_batch[i, score_from:wlen]
                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()
    base_model.train()
    return float(val_loss.item()), float(bits_per_token * tokens_per_byte)
class ReservedGPTQCalibrator:
    def __init__(self, args: Hyperparameters, base_model: GPT):
        self.args = args
        self.enabled = bool(args.gptq_enabled and args.gptq_calibration_steps > 0)
        self.active = False
        self.steps_seen = 0
        self.hooks = []
        self.hessians: dict[str, Tensor] = {}
        self.counts: dict[str, float] = {}
        if not self.enabled:
            self.modules = {}
            return
        self.modules = {
            f"{name}.weight": mod
            for name, mod in base_model.named_modules()
            if isinstance(mod, CastedLinear) and any(pattern in f"{name}.weight" for pattern in args.gptq_target_patterns)
        }
        for name, mod in self.modules.items():
            self.hooks.append(mod.register_forward_pre_hook(self._make_hook(name)))
    def _make_hook(self, name: str):
        def hook(_mod, inputs):
            if not self.active:
                return
            x = inputs[0].detach().float().reshape(-1, inputs[0].shape[-1])
            h = (x.T @ x).cpu()
            self.hessians[name] = h if name not in self.hessians else self.hessians[name] + h
            self.counts[name] = self.counts.get(name, 0.0) + float(x.shape[0])
        return hook
    def calibrate(
        self,
        base_model: GPT,
        loader,
        device: torch.device,
        world_size: int,
        grad_accum_steps: int,
        training_time_ms: float,
        max_wallclock_ms: float | None,
        log0,
    ) -> tuple[dict[str, Tensor], float]:
        if not self.enabled or not self.modules:
            return {}, 0.0
        reserve_boundary = None if max_wallclock_ms is None else max_wallclock_ms - self.args.gptq_reserve_ms
        if max_wallclock_ms is not None and training_time_ms >= max_wallclock_ms:
            raise RuntimeError(
                f"GPTQ calibration would start after the global wallclock cap: "
                f"train_ms={training_time_ms:.0f} max_wallclock_ms={max_wallclock_ms:.0f}"
            )
        if reserve_boundary is not None and training_time_ms > reserve_boundary:
            late_ms = training_time_ms - reserve_boundary
            remaining_ms = max(max_wallclock_ms - training_time_ms, 0.0) if max_wallclock_ms is not None else -1.0
            log0(f"gptq:late_start_warning over_by:{late_ms:.0f}ms reserve_boundary_ms:{reserve_boundary:.0f} remaining_ms:{remaining_ms:.0f}")
        log0(
            f"gptq:start reserved_train_data elapsed={training_time_ms:.0f}ms "
            f"steps={self.args.gptq_calibration_steps} bits:{self.args.gptq_bits} "
            f"group_size:{self.args.gptq_group_size}"
        )
        started = time.perf_counter()
        was_training = base_model.training
        base_model.eval()
        self.active = True
        try:
            with torch.no_grad():
                for _ in range(self.args.gptq_calibration_steps):
                    elapsed_ms = training_time_ms + 1000.0 * (time.perf_counter() - started)
                    if max_wallclock_ms is not None and elapsed_ms >= max_wallclock_ms:
                        break
                    x, y = loader.next_batch(self.args.train_batch_tokens, self.args.train_seq_len, grad_accum_steps)
                    with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
                        base_model(x, y)
                    self.steps_seen += 1
        finally:
            self.active = False
            if was_training:
                base_model.train()
        stats = self.finalize(device=device, distributed=world_size > 1)
        calib_ms = 1000.0 * (time.perf_counter() - started)
        log0(f"gptq:calibrated steps:{self.steps_seen} modules:{len(stats)} in {calib_ms / 1000.0:.1f}s")
        return stats, calib_ms
    def finalize(self, device: torch.device, distributed: bool) -> dict[str, Tensor]:
        out: dict[str, Tensor] = {}
        for name, h_cpu in self.hessians.items():
            h = h_cpu.to(device=device, dtype=torch.float32)
            count = torch.tensor(self.counts.get(name, 0.0), device=device, dtype=torch.float32)
            if distributed:
                dist.all_reduce(h, op=dist.ReduceOp.SUM)
                dist.all_reduce(count, op=dist.ReduceOp.SUM)
            if count.item() > 0:
                out[name] = (h / count.clamp_min(1.0)).cpu()
        return out
    def close(self) -> None:
        for hook in self.hooks:
            hook.remove()
        self.hooks.clear()
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 = "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
    enable_cudnn_sdp(False)
    enable_flash_sdp(_HAS_FLASH_ATTN)
    enable_mem_efficient_sdp(not _HAS_FLASH_ATTN)
    enable_math_sdp(not _HAS_FLASH_ATTN)
    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)
    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")))
    effective_eval_seq_len = args.eval_seq_len if args.eval_seq_len > 0 else args.train_seq_len
    val_seq_len = max(args.train_seq_len, effective_eval_seq_len)
    val_tokens = load_validation_tokens(args.val_files, val_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"val_loader:shards pattern={args.val_files} tokens:{val_tokens.numel() - 1}")
    def make_gpt() -> GPT:
        return GPT(
            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_mult=args.mlp_mult,
            mlp_hidden=args.mlp_hidden, 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,
            bigram_vocab_size=args.bigram_vocab_size, bigram_dim=args.bigram_dim, xsa_last_n=args.xsa_last_n,
            partial_rope_dims=args.partial_rope_dims, ln_scale=args.ln_scale, ve_enabled=args.ve_enabled,
            ve_dim=args.ve_dim, ve_layers=args.ve_layers, negative_slope=args.negative_slope,
            attn_pattern=args.attn_pattern, local_attn_window=args.local_attn_window,
        )
    base_model = make_gpt().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
    matrix_params: list[Tensor] = []
    scalar_params: list[Tensor] = []
    token_params = [base_model.tok_emb.weight]
    for name, param in base_model.named_parameters():
        if name == "tok_emb.weight":
            continue
        if name == "bigram.embed.weight" or name == "ve_shared.embed.weight":
            token_params.append(param)
            continue
        if param.ndim == 2 and not any(pattern in name for pattern in CONTROL_TENSOR_NAME_PATTERNS):
            matrix_params.append(param)
        else:
            scalar_params.append(param)
    token_lr = args.tied_embed_lr
    optimizer_tok = torch.optim.AdamW(
        [{"params": token_params, "lr": token_lr, "base_lr": token_lr}],
        betas=(args.beta1, args.beta2),
        eps=args.adam_eps,
        weight_decay=args.adam_weight_decay,
        fused=True,
    )
    optimizer_muon = Muon(
        matrix_params,
        lr=args.matrix_lr,
        momentum=args.muon_momentum,
        backend_steps=args.muon_backend_steps,
        weight_decay=args.muon_weight_decay,
    )
    for group in optimizer_muon.param_groups:
        group["base_lr"] = args.matrix_lr
    optimizer_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_weight_decay,
        fused=True,
    )
    optimizers: list[torch.optim.Optimizer] = [optimizer_tok, optimizer_muon, optimizer_scalar]
    n_params = sum(p.numel() for p in base_model.parameters())
    log0(f"model_params:{n_params}")
    xsa_layers = [i for i, b in enumerate(base_model.blocks) if b.attn.use_xsa]
    attn_pattern = ",".join("L" if b.attn.attn_mode == "local" else "G" for b in base_model.blocks)
    log0(f"XSA:last_{args.xsa_last_n} active_layers:{xsa_layers}")
    log0(f"world_size:{world_size} grad_accum_steps:{grad_accum_steps}")
    log0(f"sdp_backends:cudnn=False flash={int(_HAS_FLASH_ATTN)} mem_efficient={int(not _HAS_FLASH_ATTN)} math={int(not _HAS_FLASH_ATTN)}")
    log0(f"attention_mode:gqa num_heads:{args.num_heads} num_kv_heads:{args.num_kv_heads}")
    log0(f"attn_pattern:{attn_pattern} local_attn_window:{args.local_attn_window}")
    log0(f"tied_embed_lr:{token_lr} 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} iterations:{args.iterations} warmup_steps:{args.warmup_steps} max_wallclock_seconds:{args.max_wallclock_seconds:.3f}")
    log0(f"seed:{args.seed}")
    train_loader = build_train_loader(args, rank, world_size, device)
    log0(f"train_loader:dataset:{dataset_dir.name} train_shards:{actual_train_files} {train_loader.describe()}")
    calibrator = ReservedGPTQCalibrator(args, base_model)
    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
    train_budget_ms = None
    if max_wallclock_ms is not None:
        train_budget_ms = (
            max_wallclock_ms - args.gptq_reserve_ms - args.gptq_guard_ms
            if args.gptq_enabled else max_wallclock_ms
        )
        if args.gptq_enabled:
            log0(f"gptq:reserving reserve_ms:{args.gptq_reserve_ms:.0f} guard_ms:{args.gptq_guard_ms:.0f} train_budget_ms:{train_budget_ms:.0f} calib_steps:{args.gptq_calibration_steps} bits:{args.gptq_bits} group_size:{args.gptq_group_size}")
    def lr_mul(step: int, elapsed_ms: float) -> float:
        if args.warmdown_iters <= 0:
            return 1.0
        if train_budget_ms is None:
            warmdown_start = max(args.iterations - args.warmdown_iters, 0)
            return max((args.iterations - step) / max(args.warmdown_iters, 1), 0.0) if warmdown_start <= step < args.iterations else 1.0
        step_ms = elapsed_ms / max(step, 1)
        warmdown_ms = args.warmdown_iters * step_ms
        remaining_ms = max(train_budget_ms - elapsed_ms, 0.0)
        return remaining_ms / max(warmdown_ms, 1e-9) if remaining_ms <= warmdown_ms else 1.0
    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 = build_train_loader(args, rank, world_size, device)
    training_time_ms = 0.0
    stop_after_step: int | None = None
    ema_state: dict[str, Tensor] | None = None
    if args.ema_enabled:
        ema_state = {name: t.detach().float().clone() for name, t in base_model.state_dict().items()}
    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
        elapsed_ms = training_time_ms + 1000.0 * (time.perf_counter() - 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 ema_state is not None:
            d = args.ema_decay
            for name, t in base_model.state_dict().items():
                ema_state[name].mul_(d).add_(t.detach().float(), alpha=1.0 - d)
        step += 1
        approx_training_time_ms = training_time_ms + 1000.0 * (time.perf_counter() - t0)
        if train_budget_ms is not None and train_budget_ms - approx_training_time_ms <= max(args.gptq_guard_ms * 0.25, 1000.0):
            torch.cuda.synchronize()
            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"train_time:{approx_training_time_ms:.0f}ms step_avg:{approx_training_time_ms / step:.2f}ms"
            )
        reached_cap = train_budget_ms is not None and approx_training_time_ms >= train_budget_ms
        if distributed and train_budget_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"
    )
    if ema_state is not None:
        log0(f"ema: loading EMA weights (decay={args.ema_decay})")
        ema_loaded = {name: t.to(dtype=base_model.state_dict()[name].dtype) for name, t in ema_state.items()}
        base_model.load_state_dict(ema_loaded, strict=True)
        del ema_state
    else:
        log0("ema: disabled")
    calibration_ms = 0.0
    gptq_hessians: dict[str, Tensor] = {}
    if args.gptq_enabled:
        calib_loader = build_train_loader(args, rank, world_size, device)
        gptq_hessians, calibration_ms = calibrator.calibrate(
            base_model=base_model,
            loader=calib_loader,
            device=device,
            world_size=world_size,
            grad_accum_steps=grad_accum_steps,
            training_time_ms=training_time_ms,
            max_wallclock_ms=max_wallclock_ms,
            log0=log0,
        )
    total_train_plus_calibration_ms = training_time_ms + calibration_ms
    log0(f"train_plus_calibration_ms:{total_train_plus_calibration_ms:.0f}")
    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")
    calibrator.close()
    quant_obj, quant_stats = quantize_state_dict_for_export(base_model.state_dict(), gptq_hessians, args)
    quant_buf = io.BytesIO()
    torch.save(quant_obj, quant_buf)
    quant_raw = quant_buf.getvalue()
    quant_blob = lzma.compress(quant_raw, preset=6)
    quant_raw_bytes = len(quant_raw)
    if master_process:
        with open("final_model.int6.ptz", "wb") as f:
            f.write(quant_blob)
        quant_file_bytes = os.path.getsize("final_model.int6.ptz")
        code_bytes = len(code.encode("utf-8"))
        payload_bytes = quant_stats["int8_payload_bytes"] + quant_stats["packed_payload_bytes"]
        ratio = quant_stats["baseline_tensor_bytes"] / max(payload_bytes, 1)
        log0(
            f"Serialized model int6+lzma: {quant_file_bytes} bytes "
            f"(payload:{payload_bytes} packed:{quant_stats['packed_payload_bytes']} raw_torch:{quant_raw_bytes} "
            f"payload_ratio:{ratio:.2f}x pruned_codes:{quant_stats['pruned_codes']})"
        )
        log0(f"Total submission size int6+lzma: {quant_file_bytes + code_bytes} bytes target:{args.target_total_bytes}")
    if distributed:
        dist.barrier()
    with open("final_model.int6.ptz", "rb") as f:
        quant_blob_disk = f.read()
    quant_decompressed = lzma.decompress(quant_blob_disk)
    quant_state = torch.load(io.BytesIO(quant_decompressed), map_location="cpu", weights_only=False)
    eval_model = make_gpt().to(device).bfloat16()
    for module in eval_model.modules():
        if isinstance(module, CastedLinear):
            module.float()
    restore_low_dim_params_to_fp32(eval_model)
    eval_model.load_state_dict(dequantize_state_dict_for_export(quant_state), strict=True)
    compiled_eval = torch.compile(eval_model, dynamic=False, fullgraph=True)
    torch.cuda.synchronize()
    t_qeval = time.perf_counter()
    q_val_loss, q_val_bpb = eval_val(
        args,
        compiled_eval,
        rank,
        world_size,
        device,
        grad_accum_steps,
        val_tokens,
        base_bytes_lut,
        has_leading_space_lut,
        is_boundary_token_lut,
        eval_seq_len=effective_eval_seq_len,
    )
    torch.cuda.synchronize()
    log0(
        f"final_int6_lzma_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 "
        f"eval_seq_len:{effective_eval_seq_len}"
    )
    log0(f"final_int6_lzma_roundtrip_exact val_loss:{q_val_loss:.8f} val_bpb:{q_val_bpb:.8f}")
    if args.eval_stride > 0:
        torch.cuda.synchronize()
        t_slide = time.perf_counter()
        s_val_loss, s_val_bpb = eval_val_sliding(
            args, eval_model, rank, world_size, device,
            val_tokens, base_bytes_lut, has_leading_space_lut, is_boundary_token_lut,
            eval_seq_len=effective_eval_seq_len, eval_stride=args.eval_stride,
            batch_seqs=args.eval_batch_seqs,
        )
        torch.cuda.synchronize()
        log0(
            f"final_sliding_window val_loss:{s_val_loss:.4f} val_bpb:{s_val_bpb:.4f} "
            f"eval_time:{1000.0 * (time.perf_counter() - t_slide):.0f}ms "
            f"stride:{args.eval_stride} seq_len:{effective_eval_seq_len}"
        )
        log0(f"final_sliding_window_exact val_loss:{s_val_loss:.8f} val_bpb:{s_val_bpb:.8f}")
    if distributed:
        dist.destroy_process_group()
if __name__ == "__main__":
    main()