train_gpt_mlx.py

#!/usr/bin/env python3
"""
The `train_gpt.py` and `train_gpt_mlx.py` scripts are intended as good launching-off points for new participants, not SOTA configs. We'll accept PRs that tune, improve, or simplify these scripts without significantly increasing complexity, but competitive submissions should stay in the `/records` folder.

Hard stop: To keep readable for newcomers, let's make sure `train_gpt.py` and `train_gpt_mlx.py` never are longer than 1500 lines.
"""
from __future__ import annotations

import glob
import json
import math
import os
import pickle
import sys
import time
import uuid
import zlib
import zstandard as _zstd_mlx
try:
    import brotli as _brotli
except ImportError:
    _brotli = None
from collections.abc import Callable
from pathlib import Path

import numpy as np
import sentencepiece as spm

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
from mlx.utils import tree_flatten, tree_unflatten

# ==============================================================================
# SHARD FORMAT + COMPUTE DTYPE
# ==============================================================================

COMPUTE_DTYPE = mx.bfloat16

# ==============================================================================
# HYPERPARAMETERS
# ==============================================================================
# Default Simple Baseline run:
# - 9 transformer blocks at width 512
# - 8 attention heads with 4 KV heads (GQA) and 2x MLP expansion
# - vocab size 1024, sequence length 1024, tied embeddings
# - 524,288 train tokens per step for 20,000 iterations with a ~10 minute cap
class Hyperparameters:
    # Data / tokenizer.
    data_path: str = os.environ.get("DATA_PATH", "./data/datasets/fineweb10B_sp1024")
    tokenizer_path: str = os.environ.get("TOKENIZER_PATH", "./data/tokenizers/fineweb_1024_bpe.model")
    run_id: str = os.environ.get("RUN_ID", str(uuid.uuid4()))
    seed: int = int(os.environ.get("SEED", 1337))

    # Training loop. These defaults now mirror train_gpt.py on a single process.
    iterations: int = int(os.environ.get("ITERATIONS", 20_000))
    val_loss_every: int = int(os.environ.get("VAL_LOSS_EVERY", 0))
    # Validation always uses the full fineweb_val split.
    val_batch_size: int = int(os.environ.get("VAL_BATCH_SIZE", 524_288))
    train_log_every: int = int(os.environ.get("TRAIN_LOG_EVERY", 200))
    train_batch_tokens: int = int(os.environ.get("TRAIN_BATCH_TOKENS", 524_288))
    grad_accum_steps: int = int(os.environ.get("GRAD_ACCUM_STEPS", 8))
    train_seq_len: int = int(os.environ.get("TRAIN_SEQ_LEN", os.environ.get("TRAIN_MAX_SEQ_LEN", 1024)))
    # Chunk each logical MLX microbatch into smaller sub-batches to reduce peak
    # memory pressure without changing the effective optimizer batch.
    mlx_max_microbatch_tokens: int = int(os.environ.get("MLX_MAX_MICROBATCH_TOKENS", 8_192))
    # Force MLX to materialize the graph after every sub-batch, preventing lazy
    # graph buildup across accumulation steps. Keeps peak memory low on 16GB machines.
    # Disable on 32GB+ unified memory for better throughput (MLX_EAGER_EVAL=0).
    mlx_eager_eval: bool = bool(int(os.environ.get("MLX_EAGER_EVAL", "1")))
    warmup_steps: int = int(os.environ.get("WARMUP_STEPS", 20))
    warmdown_iters: int = int(os.environ.get("WARMDOWN_ITERS", 3500))
    max_wallclock_seconds: float = float(os.environ.get("MAX_WALLCLOCK_SECONDS", 600.0))

    # Model (defaults match the current baseline setup).
    vocab_size: int = int(os.environ.get("VOCAB_SIZE", 1024))
    num_layers: int = int(os.environ.get("NUM_LAYERS", 11))
    model_dim: int = int(os.environ.get("MODEL_DIM", 512))
    num_heads: int = int(os.environ.get("NUM_HEADS", 8))
    num_kv_heads: int = int(os.environ.get("NUM_KV_HEADS", 4))
    mlp_mult: int = int(os.environ.get("MLP_MULT", 3))
    tie_embeddings: bool = bool(int(os.environ.get("TIE_EMBEDDINGS", "1")))
    tied_embed_init_std: float = float(os.environ.get("TIED_EMBED_INIT_STD", 0.005))
    logit_chunk_tokens: int = int(os.environ.get("LOGIT_CHUNK_TOKENS", 0))
    logit_softcap: float = float(os.environ.get("LOGIT_SOFTCAP", 30.0))
    rope_base: float = float(os.environ.get("ROPE_BASE", 10000.0))
    rope_dims: int = int(os.environ.get("ROPE_DIMS", 16))
    qk_gain_init: float = float(os.environ.get("QK_GAIN_INIT", 1.5))
    z_loss_weight: float = float(os.environ.get("Z_LOSS_WEIGHT", 1e-4))
    bigram_hash_size: int = int(os.environ.get("BIGRAM_HASH_SIZE", 4096))
    ln_scale: bool = bool(int(os.environ.get("LN_SCALE", "1")))
    mtp_num_heads: int = int(os.environ.get("MTP_NUM_HEADS", 3))
    mtp_loss_weight: float = float(os.environ.get("MTP_LOSS_WEIGHT", 0.15))

    # Optimizer. We keep the same per-group defaults as train_gpt.py.
    beta1: float = float(os.environ.get("BETA1", 0.9))
    beta2: float = float(os.environ.get("BETA2", 0.95))
    adam_eps: float = float(os.environ.get("ADAM_EPS", 1e-8))
    tied_embed_lr: float = float(os.environ.get("TIED_EMBED_LR", 0.05))
    matrix_lr: float = float(os.environ.get("MATRIX_LR", 0.04))
    scalar_lr: float = float(os.environ.get("SCALAR_LR", 0.04))
    muon_momentum: float = float(os.environ.get("MUON_MOMENTUM", 0.99))
    muon_backend_steps: int = int(os.environ.get("MUON_BACKEND_STEPS", 4))
    muon_momentum_warmup_start: float = float(os.environ.get("MUON_MOMENTUM_WARMUP_START", 0.92))
    muon_momentum_warmup_steps: int = int(os.environ.get("MUON_MOMENTUM_WARMUP_STEPS", 1500))
    grad_clip_norm: float = float(os.environ.get("GRAD_CLIP_NORM", 0.3))
    qat_start_fraction: float = float(os.environ.get("QAT_START_FRACTION", 0.15))
    eval_stride: int = int(os.environ.get("EVAL_STRIDE", 64))
    use_ngram_eval: bool = bool(int(os.environ.get("USE_NGRAM_EVAL", "0")))
    ngram_order: int = int(os.environ.get("NGRAM_ORDER", "9"))
    ema_decay: float = float(os.environ.get("EMA_DECAY", 0.997))
    slot_enabled: bool = bool(int(os.environ.get("SLOT_ENABLED", "0")))
    slot_steps: int = int(os.environ.get("SLOT_STEPS", 8))
    slot_lr: float = float(os.environ.get("SLOT_LR", 0.005))
    window_size: int = int(os.environ.get("WINDOW_SIZE", 0))
    window_attn_layers: str = os.environ.get("WINDOW_ATTN_LAYERS", "")
    eval_seq_len: int = int(os.environ.get("EVAL_SEQ_LEN", 0))

    out_dir: str = os.environ.get("OUT_DIR", "logs")

    @property
    def train_files(self) -> str:
        return f"{self.data_path}/fineweb_train_*.bin"

    @property
    def val_files(self) -> str:
        return f"{self.data_path}/fineweb_val_*.bin"

    @property
    def microbatch_tokens(self) -> int:
        return self.train_batch_tokens // self.grad_accum_steps

    def lr_mul(self, step: int, elapsed_ms: float) -> float:
        if self.warmdown_iters <= 0:
            return 1.0
        if self.max_wallclock_seconds <= 0:
            warmdown_start = max(self.iterations - self.warmdown_iters, 0)
            return max((self.iterations - step) / max(self.warmdown_iters, 1), 0.0) if warmdown_start <= step < self.iterations else 1.0
        step_ms = elapsed_ms / max(step, 1)
        warmdown_ms = self.warmdown_iters * step_ms
        remaining_ms = max(1000.0 * self.max_wallclock_seconds - elapsed_ms, 0.0)
        return remaining_ms / max(warmdown_ms, 1e-9) if remaining_ms <= warmdown_ms else 1.0


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",
    ).split(",")
    if pattern
)
INT8_KEEP_FLOAT_FP32_NAME_PATTERNS = tuple(
    pattern
    for pattern in os.environ.get(
        "INT8_KEEP_FLOAT_FP32_NAME_PATTERNS",
        ",".join(CONTROL_TENSOR_NAME_PATTERNS),
    ).split(",")
    if pattern
)


def token_chunks(total_tokens: int, seq_len: int, max_chunk_tokens: int) -> list[int]:
    usable_total = (total_tokens // seq_len) * seq_len
    if usable_total <= 0:
        raise ValueError(f"token budget too small for seq_len={seq_len}")
    usable_chunk = max((max_chunk_tokens // seq_len) * seq_len, seq_len)
    chunks: list[int] = []
    remaining = usable_total
    while remaining > 0:
        chunk = min(remaining, usable_chunk)
        chunks.append(chunk)
        remaining -= chunk
    return chunks


def accumulate_flat_grads(
    accum: dict[str, mx.array] | None,
    grads_tree: dict,
    scale: float,
) -> dict[str, mx.array]:
    flat = dict(tree_flatten(grads_tree))
    if accum is None:
        return {k: g * scale for k, g in flat.items()}
    for k, g in flat.items():
        accum[k] = accum[k] + g * scale
    return accum


# ==============================================================================
# MATH HELPERS
# ==============================================================================

def rms_norm(x: mx.array, eps: float = 1e-6) -> mx.array:
    return (x * mx.rsqrt(mx.mean(x * x, axis=-1, keepdims=True) + eps)).astype(x.dtype)


# Polar Express coefficients: minimax-optimal per-iteration quintic polynomials
# for Newton-Schulz orthogonalization. Source: arXiv:2505.16932 / modded-nanogpt.
# Each (a,b,c) gives: x = a*x + (b*A + c*A²)@x where A = x@xᵀ
_PE_COEFFS = [
    (8.156554524902461,  -22.48329292557795,  15.878769915207462),
    (4.042929935166739,  -2.808917465908714,   0.5000178451051316),
    (3.8916678022926607, -2.772484153217685,   0.5060648178503393),
    (3.285753657755655,  -2.3681294933425376,  0.46449024233003106),
    (2.3465413258596377, -1.7097828382687081,  0.42323551169305323),
]

def zeropower_newtonschulz5(g: mx.array, steps: int, eps: float = 1e-7) -> mx.array:
    # Polar Express orthogonalization: adaptive per-step minimax quintic.
    # 4 steps with these coefficients ≈ quality of 5 steps with fixed coefficients.
    x = g.astype(mx.float32)
    x = x / (mx.sqrt(mx.sum(x * x)) + eps)
    transposed = x.shape[0] > x.shape[1]
    if transposed:
        x = x.T
    for a, b, c in _PE_COEFFS[:steps]:
        a_mat = x @ x.T
        x = a * x + (b * a_mat + c * (a_mat @ a_mat)) @ x
    if transposed:
        x = x.T
    return x.astype(g.dtype)


def load_data_shard(path: Path) -> np.ndarray:
    header_bytes = 256 * np.dtype("<i4").itemsize
    token_bytes = np.dtype("<u2").itemsize
    header = np.fromfile(path, 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 {path}")
    num_tokens = int(header[2])
    if path.stat().st_size != header_bytes + num_tokens * token_bytes:
        raise ValueError(f"Shard size mismatch for {path}")
    tokens = np.fromfile(path, dtype="<u2", count=num_tokens, offset=header_bytes)
    if tokens.size != num_tokens:
        raise ValueError(f"Short read for {path}")
    return tokens.astype(np.int32, copy=False)


# ==============================================================================
# TOKEN STREAMING / BATCHING
# ==============================================================================


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

    def next_file(self) -> None:
        self.file_idx = (self.file_idx + 1) % len(self.files)
        if self.file_idx == 0:
            self.epoch += 1
            if self.log_fn is not None:
                self.log_fn(
                    f"WARNING: starting epoch:{self.epoch} "
                    f"dataset:{self.dataset_name} train_shards:{len(self.files)}"
                )
        self.tokens = load_data_shard(self.files[self.file_idx])
        self.pos = 0

    def take(self, n: int) -> np.ndarray:
        chunks: list[np.ndarray] = []
        left = n
        while left > 0:
            if self.pos >= self.tokens.size:
                self.next_file()
            k = min(left, int(self.tokens.size - self.pos))
            chunks.append(self.tokens[self.pos : self.pos + k])
            self.pos += k
            left -= k
        return chunks[0] if len(chunks) == 1 else np.concatenate(chunks, axis=0)


class TokenLoader:
    def __init__(
        self,
        pattern: str,
        log_fn: Callable[[str], None] | None = None,
        dataset_name: str = "",
    ):
        self.stream = TokenStream(pattern, log_fn=log_fn, dataset_name=dataset_name)

    def next_batch(self, batch_tokens: int, seq_len: int) -> tuple[mx.array, mx.array]:
        usable = (batch_tokens // seq_len) * seq_len
        if usable <= 0:
            raise ValueError(f"token budget too small for seq_len={seq_len}")
        chunk = self.stream.take(usable + 1)
        x = chunk[:-1].reshape(-1, seq_len)
        y = chunk[1:].reshape(-1, seq_len)
        return mx.array(x, dtype=mx.int32), mx.array(y, dtype=mx.int32)


# ==============================================================================
# MODEL BLOCKS
# ==============================================================================

class CastedLinear(nn.Module):
    _qat_enabled: bool = False

    def __init__(self, in_dim: int, out_dim: int):
        super().__init__()
        self.weight = nn.Linear(in_dim, out_dim, bias=False).weight.astype(mx.float32)

    def __call__(self, x: mx.array) -> mx.array:
        w = self.weight.astype(x.dtype)
        if CastedLinear._qat_enabled and w.ndim == 2:
            w32 = self.weight.astype(mx.float32)
            scale = mx.maximum(mx.abs(w32).max(axis=1) / 31.0, mx.array(1.0 / 31.0))
            w_q = (mx.clip(mx.round(w32 / scale[:, None]), -31, 31) * scale[:, None]).astype(x.dtype)
            w = w + mx.stop_gradient(w_q - w)  # STE: forward sees w_q, backward sees w
        return x @ w.T


class RMSNormNoWeight(nn.Module):
    # MLX module wrapper around the functional RMSNorm helper so it composes nicely in blocks.
    def __call__(self, x: mx.array) -> mx.array:
        return rms_norm(x)


class CausalSelfAttention(nn.Module):
    def __init__(self, dim: int, num_heads: int, num_kv_heads: int,
                 rope_base: float, qk_gain_init: float, rope_dims: int = 0):
        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")
        self.rope_dims = rope_dims if rope_dims > 0 else self.head_dim
        kv_dim = self.num_kv_heads * self.head_dim
        self.c_q = CastedLinear(dim, dim)
        self.c_k = CastedLinear(dim, kv_dim)
        self.c_v = CastedLinear(dim, kv_dim)
        self.proj = CastedLinear(dim, dim)
        self.q_gain = mx.ones((num_heads,), dtype=mx.float32) * qk_gain_init
        self.rope = nn.RoPE(self.rope_dims, traditional=False, base=rope_base)
        self.scale = self.head_dim ** -0.5
        self.use_xsa: bool = False
        self.window_size: int = 0  # 0 = full attention

    def _xsa(self, y: mx.array, v: mx.array) -> mx.array:
        """XSA: subtract the v-aligned component from each query output head."""
        B, T, H, D = y.shape
        Hkv = v.shape[-2]
        g = H // Hkv
        y_g = y.reshape(B, T, Hkv, g, D)
        v_norm = v / mx.sqrt(mx.clip((v * v).sum(axis=-1, keepdims=True), a_min=1e-12, a_max=None))
        vn = mx.expand_dims(v_norm, axis=-2)  # [B, T, Hkv, 1, D]
        proj = (y_g * vn).sum(axis=-1, keepdims=True) * vn
        return (y_g - proj).reshape(B, T, H, D)

    def __call__(self, x: mx.array, v_embed: mx.array | None = None) -> mx.array:
        bsz, seqlen, dim = x.shape
        q = self.c_q(x).reshape(bsz, seqlen, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
        k = self.c_k(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
        v = self.c_v(x).reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
        if v_embed is not None:
            v = v + v_embed.reshape(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
        q = rms_norm(q).astype(COMPUTE_DTYPE)
        k = rms_norm(k).astype(COMPUTE_DTYPE)
        rd = self.rope_dims
        if rd < self.head_dim:
            q_rope, q_pass = q[..., :rd], q[..., rd:]
            k_rope, k_pass = k[..., :rd], k[..., rd:]
            q = mx.concatenate([self.rope(q_rope), q_pass], axis=-1)
            k = mx.concatenate([self.rope(k_rope), k_pass], axis=-1)
        else:
            q = self.rope(q)
            k = self.rope(k)
        q = q * self.q_gain.astype(q.dtype)[None, :, None, None]
        if self.window_size > 0:
            # Causal + sliding window mask: attend only to [i-window+1, i]
            ii = mx.arange(seqlen)[:, None]
            jj = mx.arange(seqlen)[None, :]
            mask = mx.where((jj > ii) | (ii - jj >= self.window_size),
                            mx.array(-1e9, dtype=q.dtype), mx.array(0.0, dtype=q.dtype))
            y = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask)
        else:
            y = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask="causal")
        y = y.transpose(0, 2, 1, 3)  # [B, T, H, D]
        if self.use_xsa:
            y = self._xsa(y, v.transpose(0, 2, 1, 3))
        y = y.reshape(bsz, seqlen, dim)
        return self.proj(y)


class MLP(nn.Module):
    # Baseline MLP uses relu^2 instead of GELU/SiLU. It is cheap and works well in this setup.
    def __init__(self, dim: int, mlp_mult: int):
        super().__init__()
        hidden = dim * mlp_mult
        self.fc = CastedLinear(dim, hidden)
        self.proj = CastedLinear(hidden, dim)

    def __call__(self, x: mx.array) -> mx.array:
        x = nn.leaky_relu(self.fc(x), negative_slope=0.5)
        return self.proj(x * x)


class ValueEmbedding(nn.Module):
    def __init__(self, vocab_size: int, ve_dim: int, kv_dim: int):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, ve_dim)
        self.embed.weight = mx.zeros((vocab_size, ve_dim), dtype=mx.float32)
        self.proj = nn.Linear(ve_dim, kv_dim, bias=False)
        self.proj.weight = mx.zeros((kv_dim, ve_dim), dtype=mx.float32)
        self.scale = mx.array(0.1, dtype=mx.float32)

    def __call__(self, token_ids: mx.array) -> mx.array:
        return self.proj(self.embed(token_ids.astype(mx.int32))) * self.scale


class SmearGate(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.gate = mx.zeros((dim,), dtype=mx.float32)

    def __call__(self, x: mx.array) -> mx.array:
        g = mx.sigmoid(self.gate.astype(x.dtype))[None, None, :]
        x_prev = mx.concatenate([mx.zeros_like(x[:, :1]), x[:, :-1]], axis=1)
        return (1.0 - g) * x + g * x_prev


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, rope_dims: int = 0,
                 layer_idx: int = 0, ln_scale: bool = False):
        super().__init__()
        self.attn_norm = RMSNormNoWeight()
        self.mlp_norm = RMSNormNoWeight()
        self.attn = CausalSelfAttention(dim, num_heads, num_kv_heads, rope_base, qk_gain_init, rope_dims=rope_dims)
        self.mlp = MLP(dim, mlp_mult)
        self.attn_scale = mx.ones((dim,), dtype=mx.float32)
        self.mlp_scale = mx.ones((dim,), dtype=mx.float32)
        self.resid_mix = mx.array(np.stack((np.ones((dim,), dtype=np.float32), np.zeros((dim,), dtype=np.float32))))
        self._ln_scale_factor = 1.0 / math.sqrt(layer_idx + 1) if ln_scale else 1.0

    def __call__(self, x: mx.array, x0: mx.array, v_embed: mx.array | None = None) -> mx.array:
        mix = self.resid_mix.astype(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), v_embed=v_embed)
        x = x + s * self.attn_scale.astype(x.dtype)[None, None, :] * attn_out
        x = x + s * self.mlp_scale.astype(x.dtype)[None, None, :] * self.mlp(self.mlp_norm(x))
        return x


class BigramHashEmbedding(nn.Module):
    """Learnable bigram-hash embedding bias. Zero-init so it starts as identity."""
    def __init__(self, hash_size: int, proj_dim: int, model_dim: int):
        super().__init__()
        self.hash_size = hash_size
        self.embed = nn.Embedding(hash_size, proj_dim)
        self.embed.weight = mx.zeros((hash_size, proj_dim), dtype=mx.float32)
        self.proj = nn.Linear(proj_dim, model_dim, bias=False)
        self.proj.weight = mx.zeros((model_dim, proj_dim), dtype=mx.float32)
        self.scale = mx.array(0.05, dtype=mx.float32)

    def __call__(self, token_ids: mx.array) -> mx.array:
        t = token_ids.astype(mx.int32)
        # Padding bucket for first token in each sequence
        pad = mx.full(t.shape[:-1] + (1,), self.hash_size - 1, dtype=mx.int32)
        # Bigram hash for positions 1..T-1
        h_rest = mx.bitwise_xor(36313 * t[..., 1:], 27191 * t[..., :-1]) % (self.hash_size - 1)
        h = mx.concatenate([pad, h_rest], axis=-1)
        out = self.proj(self.embed(h))
        return out * self.scale.astype(out.dtype)


class GPT(nn.Module):
    def __init__(self, vocab_size: int, num_layers: int, dim: int, num_heads: int, num_kv_heads: int, mlp_mult: int,
                 logit_chunk_tokens: int, logit_softcap: float, rope_base: float, tied_embed_init_std: float,
                 qk_gain_init: float, rope_dims: int = 0, z_loss_weight: float = 0.0,
                 bigram_hash_size: int = 4096, ln_scale: bool = False,
                 mtp_num_heads: int = 0, mtp_loss_weight: float = 0.15,
                 window_size: int = 0, window_attn_layers: str = ""):
        super().__init__()
        if logit_softcap <= 0.0:
            raise ValueError(f"logit_softcap must be positive, got {logit_softcap}")
        self.logit_chunk_tokens = logit_chunk_tokens
        self.logit_softcap = logit_softcap
        self.z_loss_weight = z_loss_weight
        self.mtp_num_heads = mtp_num_heads
        self.mtp_loss_weight = mtp_loss_weight
        self.tok_emb = nn.Embedding(vocab_size, dim)
        self.bigram_hash = BigramHashEmbedding(hash_size=bigram_hash_size, proj_dim=128, model_dim=dim)
        self.smear_gate = SmearGate(dim)
        self.vocab_bias = mx.zeros((vocab_size,), dtype=mx.float32)
        kv_dim = (dim // num_heads) * num_kv_heads
        self.ve_deep_0 = ValueEmbedding(vocab_size, 128, kv_dim)
        self.ve_deep_1 = ValueEmbedding(vocab_size, 128, kv_dim)
        # VE layers: inject token identity into V at layers num_layers-2, num_layers-1
        self._ve_layer_0 = num_layers - 2
        self._ve_layer_1 = num_layers - 1
        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 = mx.ones((self.num_skip_weights, dim), dtype=mx.float32)
        self.blocks = [
            Block(dim, num_heads, num_kv_heads, mlp_mult, rope_base, qk_gain_init,
                  rope_dims=rope_dims, layer_idx=i, ln_scale=ln_scale)
            for i in range(num_layers)
        ]
        # MTP auxiliary heads (training-only, discarded before quantization)
        self.mtp_heads = [CastedLinear(dim, vocab_size) for _ in range(mtp_num_heads)]
        self.final_norm = RMSNormNoWeight()
        # XSA on ALL layers (was last 4 only)
        for block in self.blocks:
            block.attn.use_xsa = True
        # Window attention: assign sliding window to specified layers
        if window_size > 0 and window_attn_layers:
            win_layers = set(int(x) for x in window_attn_layers.split(",") if x.strip())
            for i, block in enumerate(self.blocks):
                if i in win_layers:
                    block.attn.window_size = window_size

        for b in self.blocks:
            b.attn.proj.weight = mx.zeros_like(b.attn.proj.weight)
            b.mlp.proj.weight = mx.zeros_like(b.mlp.proj.weight)
        self.tok_emb.weight = (
            mx.random.normal(self.tok_emb.weight.shape, dtype=mx.float32) * tied_embed_init_std
        ).astype(COMPUTE_DTYPE)

    def softcap(self, logits: mx.array) -> mx.array:
        c = self.logit_softcap
        return c * mx.tanh(logits / c)

    def __call__(self, input_ids: mx.array) -> mx.array:
        x = self.tok_emb(input_ids).astype(COMPUTE_DTYPE)
        x = x + self.bigram_hash(input_ids).astype(COMPUTE_DTYPE)
        x = rms_norm(x)
        x = self.smear_gate(x)
        x0 = x
        skips: list[mx.array] = []
        ve_cache: dict[int, mx.array] = {}
        vl0, vl1 = self._ve_layer_0, self._ve_layer_1
        for i in range(self.num_encoder_layers):
            ve = self.ve_deep_0(input_ids).astype(COMPUTE_DTYPE) if i == vl0 else (
                 self.ve_deep_1(input_ids).astype(COMPUTE_DTYPE) if i == vl1 else None)
            if ve is not None:
                ve_cache[i] = ve
            x = self.blocks[i](x, x0, v_embed=ve)
            skips.append(x)
        for i in range(self.num_decoder_layers):
            if skips:
                x = x + self.skip_weights[i].astype(x.dtype)[None, None, :] * skips.pop()
            bi = self.num_encoder_layers + i
            ve = ve_cache.get(bi) or (self.ve_deep_0(input_ids).astype(COMPUTE_DTYPE) if bi == vl0 else (
                 self.ve_deep_1(input_ids).astype(COMPUTE_DTYPE) if bi == vl1 else None))
            x = self.blocks[bi](x, x0, v_embed=ve)
        return self.final_norm(x)

    def _logits(self, x: mx.array) -> mx.array:
        logits_proj = x @ self.tok_emb.weight.astype(x.dtype).T + self.vocab_bias.astype(x.dtype)
        return self.softcap(logits_proj)

    def loss(self, input_ids: mx.array, target_ids: mx.array, training: bool = True) -> mx.array:
        x_3d = self(input_ids)
        x = x_3d.reshape(-1, self.tok_emb.weight.shape[1])
        y = target_ids.reshape(-1)
        logits = self._logits(x)
        ce = nn.losses.cross_entropy(logits.astype(mx.float32), y, reduction="mean")
        if self.z_loss_weight > 0:
            lse = mx.logsumexp(logits.astype(mx.float32), axis=-1)
            ce = ce + self.z_loss_weight * (lse * lse).mean()
        # MTP: auxiliary heads predict tokens k+1 steps ahead (training only)
        if training and self.mtp_num_heads > 0 and len(self.mtp_heads) > 0:
            seqlen = input_ids.shape[-1]
            dim = x_3d.shape[-1]
            mtp_sum = mx.array(0.0)
            for k, mtp_head in enumerate(self.mtp_heads):
                vt = seqlen - (k + 1)
                if vt <= 0:
                    continue
                ml = mtp_head(x_3d[:, :vt].reshape(-1, dim))
                ml = self.softcap(ml)
                mtp_sum = mtp_sum + nn.losses.cross_entropy(
                    ml.astype(mx.float32), target_ids[:, k+1:].reshape(-1), reduction="mean")
            ce = ce + self.mtp_loss_weight * mtp_sum / max(self.mtp_num_heads, 1)
        return ce

    def forward_logits(self, input_ids: mx.array) -> mx.array:
        return self._logits(self(input_ids))

# ==============================================================================
# OPTIMIZERS (MUON + ADAM SPLIT)
# ==============================================================================
class Muon:
    # Muon applies SGD-momentum to matrix gradients, then orthogonalizes the result before the
    # parameter update.
    def __init__(self, keys: list[str], params: dict[str, mx.array], args: Hyperparameters,
                 weight_decay: float = 0.04):
        self.keys = keys
        self.args = args
        self.weight_decay = weight_decay
        self.buffers = {k: mx.zeros_like(params[k]) for k in keys}

    def step(self, params: dict[str, mx.array], grads: dict[str, mx.array], step: int, lr_mul: float) -> dict[str, mx.array]:
        if self.args.muon_momentum_warmup_steps:
            t = min(step / self.args.muon_momentum_warmup_steps, 1.0)
            momentum = (1.0 - t) * self.args.muon_momentum_warmup_start + t * self.args.muon_momentum
        else:
            momentum = self.args.muon_momentum
        lr = self.args.matrix_lr * lr_mul
        out: dict[str, mx.array] = {}
        for k in self.keys:
            p = params[k]
            g = grads[k]
            buf = momentum * self.buffers[k] + g
            self.buffers[k] = buf
            g_eff = g + momentum * buf
            g_ortho = zeropower_newtonschulz5(g_eff, self.args.muon_backend_steps)
            scale = math.sqrt(max(1.0, float(p.shape[0]) / float(p.shape[1])))
            p_wd = p * (1.0 - lr * self.weight_decay) if self.weight_decay > 0 else p
            out[k] = p_wd - lr * (g_ortho * scale).astype(p.dtype)
        return out


class SplitOptimizers:
    # - embeddings: Adam with the tied-embedding LR
    # - block matrices (2D): Muon
    # - block scalars + skip weights: Adam
    # This preserves the high-level optimization behavior even though MLX internals differ.
    def __init__(self, model: GPT, args: Hyperparameters):
        self.args = args
        params = dict(tree_flatten(model.parameters()))
        self.embed_key = "tok_emb.weight"
        self.matrix_keys = [
            k
            for k, p in params.items()
            if k.startswith("blocks.") and p.ndim == 2 and not any(pattern in k for pattern in CONTROL_TENSOR_NAME_PATTERNS)
        ]
        self.scalar_keys = [
            k
            for k, p in params.items()
            if k == "skip_weights" or (k.startswith("blocks.") and (p.ndim < 2 or any(pattern in k for pattern in CONTROL_TENSOR_NAME_PATTERNS)))
        ]

        self.muon = Muon(self.matrix_keys, params, args)
        self.adam_embed = optim.Adam(
            learning_rate=args.tied_embed_lr,
            betas=[args.beta1, args.beta2],
            eps=args.adam_eps,
            bias_correction=True,
        )
        self.adam_scalar = optim.Adam(
            learning_rate=args.scalar_lr,
            betas=[args.beta1, args.beta2],
            eps=args.adam_eps,
            bias_correction=True,
        )

    def step(self, model: GPT, grads_tree: dict, step: int, lr_mul: float) -> None:
        params = dict(tree_flatten(model.parameters()))
        grads = dict(tree_flatten(grads_tree))
        updated = dict(params)

        updated.update(self.muon.step(params, grads, step=step, lr_mul=lr_mul))

        self.adam_embed.learning_rate = self.args.tied_embed_lr * lr_mul
        updated.update(
            self.adam_embed.apply_gradients(
                {self.embed_key: grads[self.embed_key]},
                {self.embed_key: params[self.embed_key]},
            )
        )

        self.adam_scalar.learning_rate = self.args.scalar_lr * lr_mul
        scalar_grads = {k: grads[k] for k in self.scalar_keys}
        scalar_params = {k: params[k] for k in self.scalar_keys}
        updated.update(self.adam_scalar.apply_gradients(scalar_grads, scalar_params))

        model.update(tree_unflatten(list(updated.items())))

# ==============================================================================
# QUANTIZATION (INT8 + ZLIB)
# ==============================================================================
# - per-row int8 for 2D float tensors
# - per-tensor int8 for other float tensors
# - fp16 passthrough for small float tensors
# - exact passthrough for non-floats

MX_DTYPE_FROM_NAME = {
    "float32": mx.float32,
    "float16": mx.float16,
    "bfloat16": mx.bfloat16,
}

INT8_KEEP_FLOAT_MAX_NUMEL = 65_536
INT8_KEEP_FLOAT_STORE_DTYPE = np.float16
INT8_PER_ROW_SCALE_DTYPE = np.float16
INT8_CLIP_PERCENTILE = 99.99984
INT8_CLIP_Q = INT8_CLIP_PERCENTILE / 100.0


def _np_float32(arr: mx.array) -> np.ndarray:
    return np.array(arr.astype(mx.float32), dtype=np.float32, copy=False)


def keep_float_array(name: str, arr: mx.array, passthrough_orig_dtypes: dict[str, str]) -> np.ndarray:
    if any(pattern in name for pattern in INT8_KEEP_FLOAT_FP32_NAME_PATTERNS):
        return np.ascontiguousarray(_np_float32(arr))
    if arr.dtype in {mx.float32, mx.bfloat16}:
        passthrough_orig_dtypes[name] = str(arr.dtype).split(".")[-1]
        return np.ascontiguousarray(np.array(arr.astype(mx.float16), dtype=INT8_KEEP_FLOAT_STORE_DTYPE, copy=False))
    return np.ascontiguousarray(np.array(arr, copy=True))


def quantize_float_array(arr: mx.array) -> tuple[np.ndarray, np.ndarray]:
    f32 = _np_float32(arr)
    if f32.ndim == 2:
        # Matrices get one scale per row, which usually tracks output-channel
        # ranges much better than a single tensor-wide scale.
        clip_abs = np.quantile(np.abs(f32), INT8_CLIP_Q, axis=1) if f32.size else np.empty((f32.shape[0],), dtype=np.float32)
        clipped = np.clip(f32, -clip_abs[:, None], clip_abs[:, None])
        scale = np.maximum(clip_abs / 127.0, 1.0 / 127.0).astype(np.float32, copy=False)
        q = np.clip(np.round(clipped / scale[:, None]), -127, 127).astype(np.int8, copy=False)
        return np.ascontiguousarray(q), np.ascontiguousarray(scale.astype(INT8_PER_ROW_SCALE_DTYPE, copy=False))

    # Vectors / scalars use a simpler per-tensor scale.
    clip_abs = float(np.quantile(np.abs(f32).reshape(-1), INT8_CLIP_Q)) if f32.size else 0.0
    scale = np.array(clip_abs / 127.0 if clip_abs > 0.0 else 1.0, dtype=np.float32)
    q = np.clip(np.round(np.clip(f32, -clip_abs, clip_abs) / scale), -127, 127).astype(np.int8, copy=False)
    return np.ascontiguousarray(q), scale


def quantize_state_dict_int8(flat_state: dict[str, mx.array]) -> tuple[dict[str, object], dict[str, int]]:
    quantized: dict[str, np.ndarray] = {}
    scales: dict[str, np.ndarray] = {}
    dtypes: dict[str, str] = {}
    passthrough: dict[str, np.ndarray] = {}
    passthrough_orig_dtypes: dict[str, str] = {}
    qmeta: dict[str, dict[str, object]] = {}
    stats = dict.fromkeys(
        ("param_count", "num_tensors", "num_float_tensors", "num_nonfloat_tensors", "baseline_tensor_bytes", "int8_payload_bytes"),
        0,
    )
    for name, arr in flat_state.items():
        stats["param_count"] += int(arr.size)
        stats["num_tensors"] += 1
        stats["baseline_tensor_bytes"] += int(arr.nbytes)
        if not mx.issubdtype(arr.dtype, mx.floating):
            stats["num_nonfloat_tensors"] += 1
            passthrough[name] = np.ascontiguousarray(np.array(arr))
            stats["int8_payload_bytes"] += int(passthrough[name].nbytes)
            continue

        # Small float tensors are cheap enough to keep directly. We still downcast
        # fp32/bf16 passthrough tensors to fp16 so metadata does not dominate size.
        if int(arr.size) <= INT8_KEEP_FLOAT_MAX_NUMEL:
            kept = keep_float_array(name, arr, passthrough_orig_dtypes)
            passthrough[name] = kept
            stats["int8_payload_bytes"] += int(kept.nbytes)
            continue

        stats["num_float_tensors"] += 1
        q, s = quantize_float_array(arr)
        if s.ndim > 0:
            qmeta[name] = {"scheme": "per_row", "axis": 0}
        quantized[name] = q
        scales[name] = s
        dtypes[name] = str(arr.dtype).split(".")[-1]
        stats["int8_payload_bytes"] += int(q.nbytes + s.nbytes)
    obj: dict[str, object] = {
        "__quant_format__": "int8_clean_per_row_v1",
        "quantized": quantized,
        "scales": scales,
        "dtypes": dtypes,
        "passthrough": passthrough,
    }
    if qmeta:
        obj["qmeta"] = qmeta
    if passthrough_orig_dtypes:
        obj["passthrough_orig_dtypes"] = passthrough_orig_dtypes
    return obj, stats


def dequantize_state_dict_int8(quant_obj: dict[str, object]) -> dict[str, mx.array]:
    out: dict[str, mx.array] = {}
    qmeta = quant_obj.get("qmeta", {})
    passthrough_orig_dtypes = quant_obj.get("passthrough_orig_dtypes", {})
    for name, q in quant_obj["quantized"].items():
        q_np = np.asarray(q, dtype=np.int8)
        dtype_name = quant_obj["dtypes"][name]
        scale = np.asarray(quant_obj["scales"][name], dtype=np.float32)
        if qmeta.get(name, {}).get("scheme") == "per_row" or scale.ndim > 0:
            # Broadcast the saved row scale back across trailing dimensions.
            out_arr = q_np.astype(np.float32) * scale.reshape((q_np.shape[0],) + (1,) * (q_np.ndim - 1))
        else:
            out_arr = q_np.astype(np.float32) * float(scale)
        out[name] = mx.array(out_arr, dtype=MX_DTYPE_FROM_NAME[dtype_name])
    for name, arr in quant_obj["passthrough"].items():
        # Restore small tensors, undoing the temporary fp16 storage cast if needed.
        out_arr = np.array(arr, copy=True)
        orig_dtype = passthrough_orig_dtypes.get(name)
        if isinstance(orig_dtype, str):
            out[name] = mx.array(out_arr, dtype=MX_DTYPE_FROM_NAME[orig_dtype])
        else:
            out[name] = mx.array(out_arr)
    return out


# GPTQ-lite Int6: clip_range=31, best-of-5 percentile search per row.

def gptq_lite_quantize_array(arr: mx.array, clip_range: int = 31) -> tuple[np.ndarray, np.ndarray]:
    f32 = _np_float32(arr)
    if f32.ndim == 2:
        best_q, best_s, best_err = None, None, float("inf")
        for pct in [0.9990, 0.9995, 0.9999, 0.99999, 1.0]:
            rc = np.quantile(np.abs(f32), pct, axis=1) if pct < 1.0 else np.abs(f32).max(axis=1)
            s = np.maximum(rc / clip_range, 1.0 / clip_range).astype(np.float16)
            q = np.clip(np.round(f32 / s[:, None].astype(np.float32)), -clip_range, clip_range).astype(np.int8)
            err = float(np.mean((f32 - q.astype(np.float32) * s[:, None].astype(np.float32)) ** 2))
            if err < best_err:
                best_q, best_s, best_err = q, s, err
        return np.ascontiguousarray(best_q), np.ascontiguousarray(best_s)
    amax = float(np.abs(f32).max())
    s = np.float16(max(amax / clip_range, 1.0 / clip_range))
    q = np.clip(np.round(f32 / float(s)), -clip_range, clip_range).astype(np.int8)
    return np.ascontiguousarray(q), s


def gptq_lite_quantize_state_dict(flat_state: dict[str, mx.array]) -> dict[str, object]:
    quantized: dict[str, np.ndarray] = {}
    scales: dict[str, np.ndarray] = {}
    dtypes: dict[str, str] = {}
    passthrough: dict[str, np.ndarray] = {}
    pod: dict[str, str] = {}
    for name, arr in flat_state.items():
        if not mx.issubdtype(arr.dtype, mx.floating) or int(arr.size) <= INT8_KEEP_FLOAT_MAX_NUMEL:
            passthrough[name] = keep_float_array(name, arr, pod) if mx.issubdtype(arr.dtype, mx.floating) else np.ascontiguousarray(np.array(arr))
            continue
        q, s = gptq_lite_quantize_array(arr)
        quantized[name] = q
        scales[name] = s
        dtypes[name] = str(arr.dtype).split(".")[-1]
    obj: dict = {"__quant_format__": "gptq_lite_int6_v1", "quantized": quantized,
                 "scales": scales, "dtypes": dtypes, "passthrough": passthrough}
    if pod:
        obj["passthrough_orig_dtypes"] = pod
    return obj


def dequantize_gptq_lite_state_dict(quant_obj: dict) -> dict[str, mx.array]:
    out: dict[str, mx.array] = {}
    pod = quant_obj.get("passthrough_orig_dtypes", {})
    for name, q in quant_obj["quantized"].items():
        q_np = np.asarray(q, dtype=np.int8)
        s = np.asarray(quant_obj["scales"][name], dtype=np.float32)
        if s.ndim > 0:
            out_arr = q_np.astype(np.float32) * s.reshape((q_np.shape[0],) + (1,) * (q_np.ndim - 1))
        else:
            out_arr = q_np.astype(np.float32) * float(s)
        out[name] = mx.array(out_arr, dtype=MX_DTYPE_FROM_NAME[quant_obj["dtypes"][name]])
    for name, arr in quant_obj["passthrough"].items():
        out_arr = np.array(arr, copy=True)
        orig_dtype = pod.get(name)
        out[name] = mx.array(out_arr, dtype=MX_DTYPE_FROM_NAME[orig_dtype]) if orig_dtype else mx.array(out_arr)
    return out


def byte_shuffle(data: bytes, element_size: int = 2) -> bytes:
    """Transpose bytes so same-position bytes are contiguous (improves compression)."""
    n = len(data) // element_size
    remainder = len(data) % element_size
    result = bytearray(len(data))
    for b in range(element_size):
        for i in range(n):
            result[b * n + i] = data[i * element_size + b]
    if remainder:
        result[element_size * n:] = data[element_size * n:]
    return bytes(result)


def byte_unshuffle(data: bytes, element_size: int = 2) -> bytes:
    """Inverse of byte_shuffle."""
    n = len(data) // element_size
    remainder = len(data) % element_size
    result = bytearray(len(data))
    for b in range(element_size):
        for i in range(n):
            result[i * element_size + b] = data[b * n + i]
    if remainder:
        result[element_size * n:] = data[element_size * n:]
    return bytes(result)


def compress_best(data: bytes) -> tuple[bytes, str]:
    """Try zstd-22 and brotli+byte-shuffle, return the smallest."""
    zstd_blob = _zstd_mlx.ZstdCompressor(level=22).compress(data)
    best, method = zstd_blob, "zstd"
    if _brotli is not None:
        for es in (2, 3):
            shuffled = byte_shuffle(data, element_size=es)
            br_blob = _brotli.compress(shuffled, quality=11)
            if len(br_blob) < len(best):
                best, method = br_blob, f"brotli_bs{es}"
    return best, method


def decompress_auto(blob: bytes, method: str) -> bytes:
    """Decompress using the method tag from compress_best."""
    if method == "zstd":
        return _zstd_mlx.ZstdDecompressor().decompress(blob)
    elif method.startswith("brotli_bs"):
        es = int(method.split("bs")[1])
        return byte_unshuffle(_brotli.decompress(blob), element_size=es)
    raise ValueError(f"Unknown compression method: {method}")


def build_sentencepiece_luts(
    sp: spm.SentencePieceProcessor, vocab_size: int
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    sp_vocab_size = int(sp.vocab_size())
    table_size = max(sp_vocab_size, vocab_size)
    base_bytes_lut = np.zeros((table_size,), dtype=np.int16)
    has_leading_space_lut = np.zeros((table_size,), dtype=np.bool_)
    is_boundary_token_lut = 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_lut[token_id] = False
        if sp.is_byte(token_id):
            base_bytes_lut[token_id] = 1
            continue
        piece = sp.id_to_piece(token_id)
        if piece.startswith("▁"):
            has_leading_space_lut[token_id] = True
            piece = piece[1:]
        base_bytes_lut[token_id] = len(piece.encode("utf-8"))
    return base_bytes_lut, has_leading_space_lut, is_boundary_token_lut


def validate_dataset_tokenizer_pair(data_path: str, tokenizer_path: str) -> tuple[str, int, int | None]:
    # The shard directory and tokenizer are coupled: val_bpb is only meaningful if we
    # decode bytes with the exact tokenizer that produced the shards. The manifest
    # lets the training script fail fast on accidental dataset/tokenizer mismatches.
    dataset_dir = Path(data_path).resolve()
    actual_train_files = len(list(dataset_dir.glob("fineweb_train_*.bin")))
    if len(dataset_dir.parents) < 2:
        return dataset_dir.name, actual_train_files, None
    manifest_path = dataset_dir.parents[1] / "manifest.json"
    if not manifest_path.is_file():
        return dataset_dir.name, actual_train_files, None

    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    dataset_entry = next((x for x in manifest.get("datasets", []) if x.get("name") == dataset_dir.name), None)
    if dataset_entry is None:
        return dataset_dir.name, actual_train_files, None

    tokenizer_name = dataset_entry.get("tokenizer_name")
    tokenizer_entry = (
        next((x for x in manifest.get("tokenizers", []) if x.get("name") == tokenizer_name), None)
        if tokenizer_name
        else None
    )
    expected_name = Path((tokenizer_entry or {}).get("model_path") or (tokenizer_entry or {}).get("path") or "").name
    if expected_name and Path(tokenizer_path).name != expected_name:
        raise ValueError(f"{dataset_dir.name} expects tokenizer {expected_name}, got {Path(tokenizer_path).name}")
    expected_train_files = (dataset_entry.get("stats") or {}).get("files_train")
    if expected_train_files is not None:
        expected_train_files = int(expected_train_files)
        if actual_train_files > expected_train_files:
            raise ValueError(
                f"{dataset_dir.name} has more train shards than expected: found {actual_train_files}, "
                f"manifest says {expected_train_files}"
            )
    return dataset_dir.name, actual_train_files, expected_train_files


def load_validation_tokens(pattern: str, seq_len: int) -> np.ndarray:
    files = [Path(p) for p in sorted(glob.glob(pattern))]
    if not files:
        raise FileNotFoundError(f"No files found for pattern: {pattern}")
    # The export pipeline writes the fixed first-50k-doc validation set to fineweb_val_*.
    tokens = np.ascontiguousarray(np.concatenate([load_data_shard(file) for file in files], axis=0))
    usable = ((tokens.size - 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 loss_and_grad_chunked(
    args: Hyperparameters,
    train_loader: TokenLoader,
    compiled_loss_and_grad,
) -> tuple[mx.array, dict]:
    chunk_sizes = token_chunks(args.microbatch_tokens, args.train_seq_len, args.mlx_max_microbatch_tokens)
    total_tokens = float(sum(chunk_sizes))
    loss_value = mx.array(0.0, dtype=mx.float32)
    grad_accum: dict[str, mx.array] | None = None
    for chunk_tokens in chunk_sizes:
        x, y = train_loader.next_batch(chunk_tokens, args.train_seq_len)
        loss, grads = compiled_loss_and_grad(x, y)
        scale = float(y.size) / total_tokens
        loss_value = loss_value + loss.astype(mx.float32) * scale
        grad_accum = accumulate_flat_grads(grad_accum, grads, scale)
        if args.mlx_eager_eval:
            mx.eval(loss_value, grad_accum)  # materialize each chunk to cap peak memory
    return loss_value, tree_unflatten(list(grad_accum.items()))


def slot_loss_for_chunk(model: GPT, x: mx.array, y: mx.array,
                        slot_steps: int = 8, slot_lr: float = 0.005) -> mx.array:
    """SLOT: optimize additive delta at last hidden layer per eval chunk.
    Gradients flow only through lm_head (backbone hidden states are stop_gradient'd).
    Returns the eval loss with the optimized delta applied."""
    d_model = model.tok_emb.weight.shape[1]
    # Compute hidden states once (frozen backbone)
    h = mx.stop_gradient(model(x))  # [B, T, d_model]

    # Optimize delta
    delta = mx.zeros((1, 1, d_model))
    # Manual AdamW state
    m = mx.zeros_like(delta)
    v_state = mx.zeros_like(delta)
    beta1, beta2, eps = 0.9, 0.999, 1e-8

    for step_i in range(slot_steps):
        # Forward through lm_head only
        logits = model._logits((h + delta).reshape(-1, d_model))
        loss = nn.losses.cross_entropy(logits.astype(mx.float32), y.reshape(-1), reduction="mean")
        # Compute gradient of loss w.r.t. delta
        grad_delta = mx.grad(
            lambda d: nn.losses.cross_entropy(
                model._logits((h + d).reshape(-1, d_model)).astype(mx.float32),
                y.reshape(-1), reduction="mean")
        )(delta)
        # AdamW update (no weight decay on delta)
        t = step_i + 1
        m = beta1 * m + (1 - beta1) * grad_delta
        v_state = beta2 * v_state + (1 - beta2) * grad_delta * grad_delta
        m_hat = m / (1 - beta1 ** t)
        v_hat = v_state / (1 - beta2 ** t)
        delta = delta - slot_lr * m_hat / (mx.sqrt(v_hat) + eps)
        mx.eval(delta, m, v_state)

    # Final eval with optimized delta
    logits = model._logits((h + delta).reshape(-1, d_model))
    final_loss = nn.losses.cross_entropy(logits.astype(mx.float32), y.reshape(-1), reduction="mean")
    mx.eval(final_loss)
    return final_loss


def eval_val(
    args: Hyperparameters,
    compiled_loss,
    val_tokens: np.ndarray,
    base_bytes_lut: np.ndarray,
    has_leading_space_lut: np.ndarray,
    is_boundary_token_lut: np.ndarray,
    log_fn: Callable[[str], None] | None = None,
    model: GPT | None = None,
) -> tuple[float, float]:
    # Validation computes two metrics:
    # - val_loss: token cross-entropy (natural log)
    # - val_bpb: tokenizer-agnostic compression metric used by the challenge
    use_slot = args.slot_enabled and model is not None
    val_batch_tokens = args.val_batch_size // args.grad_accum_steps
    if val_batch_tokens < args.train_seq_len:
        raise ValueError(
            "VAL_BATCH_SIZE must provide at least one sequence; "
            f"got VAL_BATCH_SIZE={args.val_batch_size}, GRAD_ACCUM_STEPS={args.grad_accum_steps}, "
            f"TRAIN_SEQ_LEN={args.train_seq_len}"
        )
    val_batch_seqs = val_batch_tokens // args.train_seq_len
    total_seqs = (val_tokens.size - 1) // args.train_seq_len
    total_batches = max((total_seqs + val_batch_seqs - 1) // val_batch_seqs, 1)
    total_loss_sum = 0.0
    total_tokens = 0.0
    total_bytes = 0.0
    for batch_idx, batch_seq_start in enumerate(range(0, total_seqs, val_batch_seqs), start=1):
        batch_seq_end = min(batch_seq_start + val_batch_seqs, total_seqs)
        raw_start = batch_seq_start * args.train_seq_len
        raw_end = batch_seq_end * args.train_seq_len + 1
        chunk = val_tokens[raw_start:raw_end]
        x_np = chunk[:-1].reshape(-1, args.train_seq_len)
        y_np = chunk[1:].reshape(-1, args.train_seq_len)
        x = mx.array(x_np, dtype=mx.int32)
        y = mx.array(y_np, dtype=mx.int32)
        chunk_token_count = float(y.size)
        if use_slot:
            batch_loss = slot_loss_for_chunk(model, x, y, args.slot_steps, args.slot_lr)
        else:
            batch_loss = compiled_loss(x, y).astype(mx.float32)
            mx.eval(batch_loss)
        total_loss_sum += float(batch_loss.item()) * chunk_token_count
        prev_ids = x_np.reshape(-1)
        tgt_ids = y_np.reshape(-1)
        bytes_np = base_bytes_lut[tgt_ids].astype(np.int16, copy=True)
        bytes_np += (
            has_leading_space_lut[tgt_ids] & ~is_boundary_token_lut[prev_ids]
        ).astype(np.int16, copy=False)
        total_tokens += chunk_token_count
        total_bytes += float(bytes_np.astype(np.float64).sum())
        if log_fn is not None and total_batches > 1 and (
            batch_idx == 1 or batch_idx == total_batches or batch_idx % 25 == 0
        ):
            log_fn(f"val_progress:{batch_idx}/{total_batches}")
    val_loss = total_loss_sum / total_tokens
    bits_per_token = val_loss / math.log(2.0)
    val_bpb = bits_per_token * (total_tokens / total_bytes)
    return val_loss, val_bpb

# -----------------------------
# TRAINING
# -----------------------------

def clip_grad_tree(grads_tree: dict, max_norm: float) -> dict:
    if max_norm <= 0:
        return grads_tree
    flat = dict(tree_flatten(grads_tree))
    total_sq = mx.array(0.0, dtype=mx.float32)
    for grad in flat.values():
        total_sq = total_sq + mx.sum(grad.astype(mx.float32) * grad.astype(mx.float32))
    mx.eval(total_sq)
    total_norm = math.sqrt(float(total_sq.item()))
    if total_norm <= max_norm:
        return grads_tree
    scale = max_norm / (total_norm + 1e-12)
    return tree_unflatten([(k, g * scale) for k, g in flat.items()])


def main() -> None:
    # ==============================================================================
    # TOKENIZER + VALIDATION METRIC SETUP
    # ==============================================================================
    args = Hyperparameters()
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    logfile = out_dir / f"{args.run_id}.txt"
    print(logfile)

    def log(msg: str, console: bool = True) -> None:
        if console:
            print(msg)
        with logfile.open("a", encoding="utf-8") as f:
            print(msg, file=f)

    code = Path(__file__).read_text(encoding="utf-8")
    log(code, console=False)
    log("=" * 100, console=False)
    log(f"Running Python {sys.version}", console=False)
    log(f"Running MLX {mx.__version__}", console=False)
    log("=" * 100, console=False)

    if not args.tie_embeddings:
        raise NotImplementedError("train_gpt_mlx.py only supports tied embeddings")
    if not args.tokenizer_path.endswith(".model"):
        raise ValueError(f"TOKENIZER_PATH must point to a 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_name, actual_train_files, expected_train_files = validate_dataset_tokenizer_pair(
        args.data_path,
        args.tokenizer_path,
    )
    val_tokens = load_validation_tokens(args.val_files, args.train_seq_len)

    base_bytes_lut, has_leading_space_lut, is_boundary_token_lut = build_sentencepiece_luts(
        sp, args.vocab_size
    )

    # ==============================================================================
    # TRAINING SETUP
    # ==============================================================================
    mx.random.seed(args.seed)

    train_loader = TokenLoader(args.train_files, log_fn=log, dataset_name=dataset_name)

    # ==============================================================================
    # MODEL + OPTIMIZER SETUP
    # ==============================================================================
    model = GPT(
        vocab_size=args.vocab_size,
        num_layers=args.num_layers,
        dim=args.model_dim,
        num_heads=args.num_heads,
        num_kv_heads=args.num_kv_heads,
        mlp_mult=args.mlp_mult,
        logit_chunk_tokens=args.logit_chunk_tokens,
        logit_softcap=args.logit_softcap,
        rope_base=args.rope_base,
        tied_embed_init_std=args.tied_embed_init_std,
        qk_gain_init=args.qk_gain_init,
        rope_dims=args.rope_dims,
        z_loss_weight=args.z_loss_weight,
        bigram_hash_size=args.bigram_hash_size,
        ln_scale=args.ln_scale,
        mtp_num_heads=args.mtp_num_heads,
        mtp_loss_weight=args.mtp_loss_weight,
        window_size=args.window_size,
        window_attn_layers=args.window_attn_layers,
    )
    opt = SplitOptimizers(model, args)

    # ==============================================================================
    # COMPILED TRAIN / EVAL FUNCTIONS (MLX)
    # ==============================================================================
    # The crucial MLX detail is capture scope: this model contains non-trainable arrays too (for example
    # inside RoPE modules), so compiling only against trainable parameters throws "uncaptured inputs".
    # Compiling the model-bound functions and capturing the full model state fixes that while still
    # returning gradients only for trainable parameters via nn.value_and_grad(...).
    compiled_loss = mx.compile(lambda x, y: model.loss(x, y, training=False), inputs=model.state, outputs=model.state)
    compiled_loss_and_grad = mx.compile(
        nn.value_and_grad(model, lambda x, y: model.loss(x, y)),
        inputs=model.state,
        outputs=model.state,
    )

    # Print config once so logs are self-describing.
    n_params = sum(int(np.prod(p.shape)) for _, p in tree_flatten(model.parameters()))
    log(f"run_id:{args.run_id}")
    log(f"mlx_version:{mx.__version__}")
    log(f"train_loader:shards pattern={args.train_files}")
    log(f"val_loader:shards pattern={args.val_files} tokens:{val_tokens.size - 1}")
    if expected_train_files is None:
        log(f"train_loader:dataset:{dataset_name} train_shards:{actual_train_files}")
    elif actual_train_files < expected_train_files:
        log(
            f"WARNING: train_loader:subset dataset:{dataset_name} "
            f"train_shards:{actual_train_files}/{expected_train_files} "
            f"new epochs will arrive sooner than the full dataset"
        )
    else:
        log(f"train_loader:dataset:{dataset_name} train_shards:{actual_train_files}/{expected_train_files}")
    log(f"tokenizer_path:{args.tokenizer_path}")
    log(
        f"model_params:{n_params} vocab_size:{args.vocab_size} layers:{args.num_layers} "
        f"dim:{args.model_dim} heads:{args.num_heads} kv_heads:{args.num_kv_heads} "
        f"seq_len:{args.train_seq_len} tie_embeddings:{args.tie_embeddings}"
    )
    log(
        f"iterations:{args.iterations} train_batch_tokens:{args.train_batch_tokens} grad_accum_steps:{args.grad_accum_steps} "
        f"microbatch_tokens:{args.microbatch_tokens} microbatch_batch_size:{args.microbatch_tokens // args.train_seq_len} "
        f"val_batch_size:{args.val_batch_size} "
        f"warmup_steps:{args.warmup_steps} max_wallclock_seconds:{args.max_wallclock_seconds:.3f}"
    )
    log(f"mlx_max_microbatch_tokens:{args.mlx_max_microbatch_tokens}")
    log(
        f"optimizer:muon+adam muon_matrix_params:{len(opt.matrix_keys)} scalar_params:{len(opt.scalar_keys)} "
        f"embed_lr:{args.tied_embed_lr} "
        f"matrix_lr:{args.matrix_lr} scalar_lr:{args.scalar_lr} "
        f"muon_momentum:{args.muon_momentum} muon_steps:{args.muon_backend_steps}"
    )
    log(f"val_bpb:enabled tokenizer_kind=sentencepiece tokenizer_path={args.tokenizer_path}")
    log(f"compute_dtype:{COMPUTE_DTYPE} compile:True")
    log(
        f"dtypes tok_emb:{model.tok_emb.weight.dtype} "
        f"linear_weight:{model.blocks[0].attn.c_q.weight.dtype} "
        f"skip_weights:{model.skip_weights.dtype}"
    )

    # ==============================================================================
    # TRAINING LOOP
    # ==============================================================================
    if args.warmup_steps > 0:
        # Warmup should only prime MLX compile/allocation paths. Updating parameters here forces us
        # to snapshot and restore model/optimizer state, which is expensive on unified-memory Macs.
        # Instead we run the real train shapes, force the loss/grads to materialize, and then reset
        # the loader so measured training still starts from the true init and token window.
        for warmup_step in range(args.warmup_steps):
            accum: dict[str, mx.array] | None = None
            warmup_loss = mx.array(0.0, dtype=mx.float32)
            grad_scale = 1.0 / args.grad_accum_steps
            for _ in range(args.grad_accum_steps):
                warmup_loss, grads = loss_and_grad_chunked(args, train_loader, compiled_loss_and_grad)
                accum = accumulate_flat_grads(accum, grads, grad_scale)
            mx.eval(warmup_loss, accum)
            mx.synchronize()
            if args.warmup_steps <= 20 or (warmup_step + 1) % 10 == 0 or warmup_step + 1 == args.warmup_steps:
                log(f"warmup_step:{warmup_step + 1}/{args.warmup_steps}")

        # Prime the standalone eval graph once too. It is compiled separately from value_and_grad.
        val_batch_tokens = args.val_batch_size // args.grad_accum_steps
        if val_batch_tokens < args.train_seq_len:
            raise ValueError(
                "VAL_BATCH_SIZE must provide at least one sequence; "
                f"got VAL_BATCH_SIZE={args.val_batch_size}, GRAD_ACCUM_STEPS={args.grad_accum_steps}, "
                f"TRAIN_SEQ_LEN={args.train_seq_len}"
            )
        warm_val_seqs = min(val_batch_tokens // args.train_seq_len, (val_tokens.size - 1) // args.train_seq_len)
        warm_chunk = val_tokens[: warm_val_seqs * args.train_seq_len + 1]
        x_val = mx.array(warm_chunk[:-1].reshape(-1, args.train_seq_len), dtype=mx.int32)
        y_val = mx.array(warm_chunk[1:].reshape(-1, args.train_seq_len), dtype=mx.int32)
        warm_val_loss = compiled_loss(x_val, y_val)
        mx.eval(warm_val_loss)
        mx.synchronize()

        train_loader = TokenLoader(args.train_files, log_fn=log, dataset_name=dataset_name)

    # EMA weight averaging (float params only — skip integer/bool buffers)
    EMA_DECAY = args.ema_decay
    ema_params = {k: v.astype(mx.float32) for k, v in tree_flatten(model.parameters())
                  if mx.issubdtype(v.dtype, mx.floating)}

    train_time_ms = 0.0
    max_wallclock_ms = 1000.0 * args.max_wallclock_seconds if args.max_wallclock_seconds > 0 else None
    stop_after_step: int | None = None
    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)
        if last_step or (args.val_loss_every > 0 and step % args.val_loss_every == 0):
            train_time_ms += 1000.0 * (time.perf_counter() - t0)
            # Validation always scans the same fixed full validation split.
            val_loss, val_bpb = eval_val(
                args,
                compiled_loss,
                val_tokens,
                base_bytes_lut,
                has_leading_space_lut,
                is_boundary_token_lut,
                log_fn=log,
            )
            if step % 25 == 0 or last_step:
                log(
                    f"step:{step}/{args.iterations} val_loss:{val_loss:.4f} val_bpb:{val_bpb:.4f} "
                    f"train_time:{train_time_ms:.0f}ms step_avg:{train_time_ms / max(step, 1):.2f}ms"
                )
            t0 = time.perf_counter()
        if last_step:
            if stop_after_step is not None and step < args.iterations:
                log(f"stopping_early: wallclock_cap train_time:{train_time_ms:.0f}ms step:{step}/{args.iterations}")
            break

        lr_mul = args.lr_mul(step, train_time_ms + 1000.0 * (time.perf_counter() - t0))
        if args.qat_start_fraction > 0 and lr_mul < args.qat_start_fraction and not CastedLinear._qat_enabled:
            CastedLinear._qat_enabled = True
            log(f"late_qat:enabled step:{step} lr_mul:{lr_mul:.4f}")
        step_t0 = time.perf_counter()

        accum: dict[str, mx.array] | None = None
        train_loss = mx.array(0.0, dtype=mx.float32)
        grad_scale = 1.0 / args.grad_accum_steps
        for _ in range(args.grad_accum_steps):
            loss, grads = loss_and_grad_chunked(args, train_loader, compiled_loss_and_grad)
            accum = accumulate_flat_grads(accum, grads, grad_scale)
            train_loss = train_loss + loss.astype(mx.float32) * grad_scale
            if args.mlx_eager_eval:
                mx.eval(train_loss, accum)  # materialize each microbatch to cap peak memory

        grads = tree_unflatten(list(accum.items()))
        grads = clip_grad_tree(grads, args.grad_clip_norm)
        train_loss_value = float(train_loss.item())
        opt.step(model, grads, step=step, lr_mul=lr_mul)

        # EMA update (float params only) — must eval to prevent graph buildup
        for k, v in tree_flatten(model.parameters()):
            if k in ema_params:
                ema_params[k] = ema_params[k] * EMA_DECAY + v.astype(mx.float32) * (1.0 - EMA_DECAY)
        mx.eval(ema_params)
        mx.synchronize()

        step_ms = 1000.0 * (time.perf_counter() - step_t0)
        approx_train_time_ms = train_time_ms + 1000.0 * (time.perf_counter() - t0)
        tok_s = args.train_batch_tokens / (step_ms / 1000.0)
        step += 1
        if args.train_log_every > 0 and (step <= 10 or step % args.train_log_every == 0 or stop_after_step is not None):
            log(
                f"step:{step}/{args.iterations} train_loss:{train_loss_value:.4f} "
                f"train_time:{approx_train_time_ms:.0f}ms step_avg:{approx_train_time_ms / step:.2f}ms tok_s:{tok_s:.0f}"
            )
        if max_wallclock_ms is not None and stop_after_step is None and approx_train_time_ms >= max_wallclock_ms:
            stop_after_step = step

    # ==============================================================================
    # FINAL SERIALIZATION + QUANTIZED ROUNDTRIP EVAL
    # ==============================================================================
    # We always write a raw artifact and a quantized artifact, then validate the
    # quantized roundtrip directly by loading the dequantized tensors back into the
    # model and running one final validation pass.
    # Load EMA weights for final artifact
    model_params_now = dict(tree_flatten(model.parameters()))
    ema_typed = {k: ema_params[k].astype(model_params_now[k].dtype) for k in ema_params}
    model.update(tree_unflatten(list(ema_typed.items())))
    mx.eval(model.state)

    # Strip MTP heads before artifact (training-only, not needed for inference)
    model.mtp_heads = []
    model.mtp_num_heads = 0

    out_path = out_dir / f"{args.run_id}_mlx_model.npz"
    flat_state = {k: v for k, v in tree_flatten(model.state)}
    mx.savez(str(out_path), **flat_state)
    log(f"saved_model:{out_path} bytes:{out_path.stat().st_size}")

    quant_obj, quant_stats = quantize_state_dict_int8(flat_state)
    quant_raw = pickle.dumps(quant_obj, protocol=pickle.HIGHEST_PROTOCOL)
    quant_blob = _zstd_mlx.ZstdCompressor(level=22).compress(quant_raw)
    quant_serialized_bytes = len(quant_raw)
    quant_path = out_dir / f"{args.run_id}_mlx_model.int8.ptz"
    with quant_path.open("wb") as f:
        f.write(quant_blob)
    quant_file_bytes = quant_path.stat().st_size
    ratio = quant_stats["baseline_tensor_bytes"] / max(quant_stats["int8_payload_bytes"], 1)
    log(
        f"serialized_model_int8_zlib:{quant_file_bytes} bytes "
        f"(payload:{quant_stats['int8_payload_bytes']} raw_pickle:{quant_serialized_bytes} payload_ratio:{ratio:.2f}x)"
    )

    with quant_path.open("rb") as f:
        quant_blob_disk = f.read()
    try:
        quant_flat = dequantize_state_dict_int8(pickle.loads(_zstd_mlx.ZstdDecompressor().decompress(quant_blob_disk)))
    except Exception:
        quant_flat = dequantize_state_dict_int8(pickle.loads(zlib.decompress(quant_blob_disk)))
    model.update(tree_unflatten(list(quant_flat.items())))
    mx.eval(model.state)
    uncompiled_loss = lambda x, y: model.loss(x, y, training=False)
    q_t0 = time.perf_counter()
    q_val_loss, q_val_bpb = eval_val(
        args,
        uncompiled_loss,
        val_tokens,
        base_bytes_lut,
        has_leading_space_lut,
        is_boundary_token_lut,
        log_fn=log,
        model=model,
    )
    q_eval_ms = 1000.0 * (time.perf_counter() - q_t0)
    log(f"final_int8_zlib_roundtrip val_loss:{q_val_loss:.4f} val_bpb:{q_val_bpb:.4f} eval_time:{q_eval_ms:.0f}ms")
    log(f"final_int8_zlib_roundtrip_exact val_loss:{q_val_loss:.8f} val_bpb:{q_val_bpb:.8f}")

    # GPTQ-lite int6: smaller artifact for 16MB submission limit
    gptq_obj = gptq_lite_quantize_state_dict(flat_state)
    gptq_raw = pickle.dumps(gptq_obj, protocol=pickle.HIGHEST_PROTOCOL)
    # Try both zstd and brotli+byte-shuffle, pick the smallest
    gptq_blob, gptq_method = compress_best(gptq_raw)
    gptq_zstd_blob = _zstd_mlx.ZstdCompressor(level=22).compress(gptq_raw)
    gptq_path = out_dir / f"{args.run_id}_mlx_model.int6.ptz"
    with gptq_path.open("wb") as f:
        f.write(gptq_zstd_blob)  # always write zstd for roundtrip eval compatibility
    code_bytes = len(code.encode("utf-8"))
    gptq_file_bytes = gptq_path.stat().st_size
    log(f"GPTQ-lite int6+zstd:{gptq_file_bytes} bytes total_submission:{gptq_file_bytes + code_bytes}")
    log(f"GPTQ-lite int6+{gptq_method}:{len(gptq_blob)} bytes total:{len(gptq_blob) + code_bytes} (best compression)")

    with gptq_path.open("rb") as f:
        gptq_blob_disk = f.read()
    gptq_flat = dequantize_gptq_lite_state_dict(pickle.loads(_zstd_mlx.ZstdDecompressor().decompress(gptq_blob_disk)))
    model.update(tree_unflatten(list(gptq_flat.items())))
    mx.eval(model.state)
    gq_t0 = time.perf_counter()
    gq_val_loss, gq_val_bpb = eval_val(
        args, uncompiled_loss, val_tokens,
        base_bytes_lut, has_leading_space_lut, is_boundary_token_lut, log_fn=log,
        model=model,
    )
    gq_ms = 1000.0 * (time.perf_counter() - gq_t0)
    log(f"final_int6_gptq_roundtrip val_loss:{gq_val_loss:.4f} val_bpb:{gq_val_bpb:.4f} eval_time:{gq_ms:.0f}ms")
    log(f"final_int6_gptq_roundtrip_exact val_loss:{gq_val_loss:.8f} val_bpb:{gq_val_bpb:.8f}")


if __name__ == "__main__":
    main()