upload 05_train_v5.py

05_train_v5.py

@@ -0,0 +1,594 @@
+"""
+V5 Eğitim — 200M model, multi-stage curriculum (SmolLM2 tarzı).
+Hedef GPU: RTX PRO 6000 Blackwell (96GB VRAM).
+
+Mimari (model_v5.py):
+  - 18 layer, 14 head, 896 embd, 32K vocab, 1024 context
+  - RoPE + RMSNorm + SwiGLU + QK-norm + soft-cap + tied embeddings
+
+Stack:
+  - Muon (2D weights) + AdamW (1D + embedding)
+  - bf16 + TF32 + cudnn.benchmark + torch.compile
+  - Async prefetcher per stage
+  - Multi-stage data loader (weighted sampling, progresif)
+
+Curriculum (SmolLM2 stil) — Stage1=web(3-4B), Stage2=medium(9B), Stage3=premium(3B):
+  Faz 1  [0%   - 55%]  : Stage1 %25, Stage2 %65, Stage3 %10 (medium bulk + web)
+  Faz 2  [55%  - 85%]  : Stage1 %15, Stage2 %55, Stage3 %30 (premium ısınma)
+  Faz 3  [85%  - 100%] : Stage1 %5 , Stage2 %25, Stage3 %70 (PREMIUM annealing)
+
+Replay: Geçmiş stage'leri tamamen kesmiyoruz → catastrophic forgetting önlenir.
+
+Kullanim:
+    python 05_train_v5.py                         # bastan
+    python 05_train_v5.py --resume                # latest_ckpt
+    python 05_train_v5.py --compile               # torch.compile
+    python 05_train_v5.py --max-time 480          # 8 saatlik oturum
+
+Önceden: data/v5_stage1.bin, v5_stage2.bin, v5_stage3.bin, v5_val.bin hazır olmalı.
+"""
+
+import argparse
+import math
+import os
+import signal
+import sys
+import time
+from contextlib import nullcontext
+from pathlib import Path
+
+import numpy as np
+import torch
+from tokenizers import Tokenizer
+
+from model_v5 import GPTV5, GPTConfigV5
+from muon import Muon
+
+# ============================================================
+# Konfigurasyon — V5 (RTX PRO 6000 Blackwell 96GB için)
+# ============================================================
+DATA_DIR = Path(__file__).parent / "data"
+OUT_DIR = Path(__file__).parent / "runs" / "tr-200m-v5"
+
+MODEL_CONFIG = dict(
+    block_size=2048,
+    vocab_size=32000,
+    n_layer=18,
+    n_head=14,
+    n_embd=896,
+    dropout=0.0,
+    rope_theta=10000.0,
+    logit_softcap=30.0,
+)
+
+# RTX PRO 6000 Blackwell 96GB VRAM
+# T=2048 için activation memory T=1024'ün ~2x'i (Flash-Attn ile O(B*H*T))
+# bs=48, T=2048 → activations ~60-70GB, weights+opt ~5GB → ~70GB toplam (rahat)
+# 48 = 16*3 tensor core dostu
+BATCH_SIZE = 32
+GRAD_ACCUM_STEPS = 16          # etkin batch = 528, token/step ≈ 1.08M
+MAX_STEPS = 20_000             # ~21.6B token training (Modern overtraining, 108x Chinchilla)
+LOG_INTERVAL = 10
+EVAL_INTERVAL = 400
+EVAL_ITERS = 80
+SAVE_INTERVAL = 1000
+SAMPLE_INTERVAL = 2000
+
+# LR — 200M, etkin batch ~528 için
+MUON_LR = 0.022
+ADAM_LR = 3.5e-4
+MIN_LR_RATIO = 0.1
+WARMUP_STEPS = 1000            # 20K step için %5 warmup
+LR_DECAY_STEPS = 20_000
+
+# Optimizer
+WEIGHT_DECAY = 0.1            # 200M model için biraz weight decay yararlı
+ADAM_BETA1 = 0.9
+ADAM_BETA2 = 0.95
+MUON_MOMENTUM = 0.95
+GRAD_CLIP = 1.0
+
+# Curriculum faz sınırları (oran cinsinden)
+# Stage1 = WEB     (oscar, mc4, forum, fineweb_hq)  ~3-4B token
+# Stage2 = MEDIUM  (bellaturca, cosmos, culturax, havadis, cosmopedia) ~9B token (BULK)
+# Stage3 = PREMIUM (wiki, wikisource, tezler, akademik, finepdfs, ozenli) ~3B token
+PHASE1_END = 0.55             # 0-55%   : bulk learning (medium dominant)
+PHASE2_END = 0.85             # 55-85%  : premium ısınır
+                              # 85-100% : PREMIUM annealing
+
+# Faz başına karışım oranları [stage1=web, stage2=medium, stage3=premium]
+PHASE_MIX = {
+    1: (0.25, 0.65, 0.10),    # medium bulk + web, premium dokun
+    2: (0.15, 0.55, 0.30),    # premium ısın
+    3: (0.05, 0.25, 0.70),    # PREMIUM annealing — son finishing
+}
+# ============================================================
+
+LATEST_CKPT = OUT_DIR / "latest_ckpt.pt"
+BEST_CKPT   = OUT_DIR / "best_ckpt.pt"
+LOG_FILE    = OUT_DIR / "train.log"
+
+
+def get_lr_factor(step: int) -> float:
+    if step < WARMUP_STEPS:
+        return (step + 1) / (WARMUP_STEPS + 1)
+    if step > LR_DECAY_STEPS:
+        return MIN_LR_RATIO
+    decay_ratio = (step - WARMUP_STEPS) / (LR_DECAY_STEPS - WARMUP_STEPS)
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+    return MIN_LR_RATIO + coeff * (1.0 - MIN_LR_RATIO)
+
+
+def get_phase(step: int) -> int:
+    """Mevcut step'e göre curriculum faz (1/2/3)."""
+    p = step / max(MAX_STEPS, 1)
+    if p < PHASE1_END:
+        return 1
+    if p < PHASE2_END:
+        return 2
+    return 3
+
+
+def log(msg: str):
+    print(msg, flush=True)
+    try:
+        with open(LOG_FILE, "a", encoding="utf-8") as f:
+            f.write(msg + "\n")
+    except Exception:
+        pass
+
+
+# =====================================================================
+# Data — per-stage loader
+# =====================================================================
+class StageLoader:
+    def __init__(self, bin_path: Path, block_size: int, batch_size: int,
+                 device: torch.device, pin: bool = True, name: str = ""):
+        self.data = np.memmap(bin_path, dtype=np.uint16, mode="r")
+        self.block_size = block_size
+        self.batch_size = batch_size
+        self.device = device
+        self.pin = pin and device.type == "cuda"
+        self.name = name or bin_path.stem
+        n_tok = len(self.data)
+        log(f"  {bin_path.name}: {n_tok:,} token  (~{n_tok*2/1e9:.2f} GB)")
+        self.n_tokens = n_tok
+
+    def get_batch(self):
+        bs, T = self.batch_size, self.block_size
+        ix = np.random.randint(0, len(self.data) - T - 1, size=bs)
+        x_np = np.empty((bs, T), dtype=np.int64)
+        y_np = np.empty((bs, T), dtype=np.int64)
+        for k, i in enumerate(ix):
+            x_np[k] = self.data[i:i+T]
+            y_np[k] = self.data[i+1:i+1+T]
+        x = torch.from_numpy(x_np)
+        y = torch.from_numpy(y_np)
+        if self.pin:
+            x = x.pin_memory()
+            y = y.pin_memory()
+        x = x.to(self.device, non_blocking=True)
+        y = y.to(self.device, non_blocking=True)
+        return x, y
+
+
+class MultiStageLoader:
+    """Curriculum-aware sampler — fazlara göre stage karışımı değişir."""
+    def __init__(self, stage_loaders, rng=None):
+        # stage_loaders: [s1, s2, s3]
+        self.loaders = stage_loaders
+        self.rng = rng or np.random.default_rng()
+
+    def get_batch(self, phase: int):
+        mix = PHASE_MIX[phase]
+        # Tek bir stage seçimi (batch içi karışım değil — daha temiz gradient)
+        idx = self.rng.choice(len(self.loaders), p=mix)
+        return self.loaders[idx].get_batch(), idx
+
+
+class AsyncMultiStagePrefetcher:
+    """Faz bilgisi geçilen prefetch kuyruğu. Her get() çağrısında mevcut phase
+    kullanılır — geriden gelen batch'ler hâlâ önceki phase'in karışımındaysa
+    sorun değil (geçişler yumuşaktır)."""
+    def __init__(self, multi_loader: MultiStageLoader, phase_fn, queue_size=4):
+        import threading, queue
+        self.ml = multi_loader
+        self.phase_fn = phase_fn
+        self.q = queue.Queue(maxsize=queue_size)
+        self._stop = threading.Event()
+        self.thread = threading.Thread(target=self._produce, daemon=True)
+        self.thread.start()
+
+    def _produce(self):
+        while not self._stop.is_set():
+            try:
+                ph = self.phase_fn()
+                self.q.put(self.ml.get_batch(ph))
+            except Exception:
+                self._stop.set()
+                break
+
+    def get_batch(self):
+        return self.q.get()
+
+    def close(self):
+        self._stop.set()
+
+
+# =====================================================================
+# Eval / Sample
+# =====================================================================
+@torch.no_grad()
+def estimate_loss(model, val_loader, train_loaders, ctx, eval_iters: int):
+    """Val + her stage için train loss."""
+    out = {}
+    model.eval()
+
+    # Val
+    losses = torch.zeros(eval_iters)
+    for k in range(eval_iters):
+        x, y = val_loader.get_batch()
+        with ctx:
+            _, loss = model(x, y)
+        losses[k] = loss.item()
+    out["val"] = losses.mean().item()
+
+    # Her stage'den birkaç iter
+    n_small = max(eval_iters // 4, 8)
+    for i, ld in enumerate(train_loaders, start=1):
+        losses = torch.zeros(n_small)
+        for k in range(n_small):
+            x, y = ld.get_batch()
+            with ctx:
+                _, loss = model(x, y)
+            losses[k] = loss.item()
+        out[f"stage{i}"] = losses.mean().item()
+
+    model.train()
+    return out
+
+
+@torch.no_grad()
+def sample_text(model, tokenizer, device, ctx,
+                prompt: str = "Türkiye", max_new_tokens: int = 100,
+                temperature: float = 0.8, top_k: int = 50,
+                repetition_penalty: float = 1.15):
+    model.eval()
+    ids = tokenizer.encode(prompt).ids
+    x = torch.tensor([ids], dtype=torch.long, device=device)
+    real_model = model._orig_mod if hasattr(model, "_orig_mod") else model
+    with ctx:
+        out = real_model.generate(
+            x, max_new_tokens=max_new_tokens,
+            temperature=temperature, top_k=top_k,
+            repetition_penalty=repetition_penalty,
+        )
+    text = tokenizer.decode(out[0].tolist())
+    model.train()
+    return text
+
+
+# =====================================================================
+# Checkpointing
+# =====================================================================
+def atomic_save(state: dict, path: Path):
+    tmp = path.with_suffix(path.suffix + ".tmp")
+    torch.save(state, tmp)
+    if path.exists():
+        path.unlink()
+    tmp.rename(path)
+
+
+def build_state(model, opt_muon, opt_adam, scaler, step, best_val):
+    real_model = model._orig_mod if hasattr(model, "_orig_mod") else model
+    return {
+        "model": real_model.state_dict(),
+        "opt_muon": opt_muon.state_dict(),
+        "opt_adam": opt_adam.state_dict(),
+        "scaler": scaler.state_dict(),
+        "step": step,
+        "best_val": best_val,
+        "config": MODEL_CONFIG,
+        "version": "v5",
+    }
+
+
+# =====================================================================
+# Optimizer setup
+# =====================================================================
+def create_optimizers(model, device):
+    muon_params = []
+    adam_params = []
+
+    for name, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        if p.ndim < 2:
+            adam_params.append(p)
+        elif "wte" in name or "lm_head" in name:
+            adam_params.append(p)
+        else:
+            muon_params.append(p)
+
+    seen = set()
+    adam_params_unique = []
+    for p in adam_params:
+        if id(p) not in seen:
+            seen.add(id(p))
+            adam_params_unique.append(p)
+
+    n_muon = sum(p.numel() for p in muon_params)
+    n_adam = sum(p.numel() for p in adam_params_unique)
+    log(f"  Muon params:  {n_muon/1e6:.2f}M ({len(muon_params)} tensor)")
+    log(f"  AdamW params: {n_adam/1e6:.2f}M ({len(adam_params_unique)} tensor)")
+
+    opt_muon = Muon(
+        muon_params,
+        lr=MUON_LR,
+        momentum=MUON_MOMENTUM,
+        nesterov=True,
+        ns_steps=5,
+    )
+    opt_adam = torch.optim.AdamW(
+        adam_params_unique,
+        lr=ADAM_LR,
+        betas=(ADAM_BETA1, ADAM_BETA2),
+        weight_decay=WEIGHT_DECAY,
+        fused=(device.type == "cuda"),
+    )
+    return opt_muon, opt_adam
+
+
+# =====================================================================
+# Main
+# =====================================================================
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--resume", action="store_true")
+    parser.add_argument("--resume-best", action="store_true")
+    parser.add_argument("--compile", action="store_true")
+    parser.add_argument("--max-time", type=int, default=0,
+                        help="Maksimum süre (dakika)")
+    parser.add_argument("--max-steps", type=int, default=None)
+    parser.add_argument("--batch", type=int, default=None,
+                        help="Override BATCH_SIZE")
+    parser.add_argument("--grad-accum", type=int, default=None)
+    args = parser.parse_args()
+
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+
+    global MAX_STEPS, LR_DECAY_STEPS, BATCH_SIZE, GRAD_ACCUM_STEPS
+    if args.max_steps:
+        MAX_STEPS = args.max_steps
+        LR_DECAY_STEPS = args.max_steps
+    if args.batch:
+        BATCH_SIZE = args.batch
+    if args.grad_accum:
+        GRAD_ACCUM_STEPS = args.grad_accum
+
+    # Cihaz
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if device.type == "cpu":
+        log("UYARI: CUDA yok.")
+    else:
+        log(f"GPU: {torch.cuda.get_device_name(0)}")
+        log(f"CUDA: {torch.version.cuda}, PyTorch: {torch.__version__}")
+        torch.set_float32_matmul_precision("high")
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+        torch.backends.cudnn.benchmark = True
+        # Blackwell: Flash Attention v2/v3 backend zorla (sdpa içinden)
+        try:
+            torch.backends.cuda.enable_flash_sdp(True)
+            torch.backends.cuda.enable_mem_efficient_sdp(True)
+            torch.backends.cuda.enable_math_sdp(False)
+        except Exception:
+            pass
+        # Daha agresif allocator (büyük bs için fragmentasyon azalır)
+        os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF",
+                              "expandable_segments:True")
+        log("Perf: TF32 ON, cudnn.benchmark ON, Flash-SDPA ON")
+        # VRAM rapor
+        vram_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
+        log(f"VRAM: {vram_gb:.1f} GB")
+
+    use_bf16 = device.type == "cuda" and torch.cuda.is_bf16_supported()
+    dtype = torch.bfloat16 if use_bf16 else torch.float16
+    log(f"Mixed precision: {dtype}")
+    ctx = (nullcontext() if device.type == "cpu"
+           else torch.amp.autocast(device_type="cuda", dtype=dtype))
+
+    # Data — 3 stage + val
+    log("\nData yukleniyor...")
+    bs, T = BATCH_SIZE, MODEL_CONFIG["block_size"]
+    stage1 = StageLoader(DATA_DIR / "v5_stage1.bin", T, bs, device, name="stage1")
+    stage2 = StageLoader(DATA_DIR / "v5_stage2.bin", T, bs, device, name="stage2")
+    stage3 = StageLoader(DATA_DIR / "v5_stage3.bin", T, bs, device, name="stage3")
+    val_loader = StageLoader(DATA_DIR / "v5_val.bin", T, bs, device, name="val")
+
+    total_tokens = stage1.n_tokens + stage2.n_tokens + stage3.n_tokens
+    log(f"  Toplam: {total_tokens/1e9:.2f}B token")
+
+    multi = MultiStageLoader([stage1, stage2, stage3])
+
+    # Prefetcher — phase fonksiyonunu bir mutable ref ile vereceğiz
+    step_ref = {"step": 0}
+    def cur_phase():
+        return get_phase(step_ref["step"])
+
+    prefetch = AsyncMultiStagePrefetcher(multi, cur_phase, queue_size=4)
+    log("  Async multi-stage prefetcher aktif")
+
+    tokenizer = Tokenizer.from_file(str(DATA_DIR / "tokenizer-tr-v5.json"))
+
+    # Model
+    log("\nModel V5 olusturuluyor...")
+    cfg = GPTConfigV5(**MODEL_CONFIG)
+    model = GPTV5(cfg).to(device)
+    n_params = model.num_params()
+    log(f"  Toplam: {n_params/1e6:.2f}M param")
+    log(f"  Mimari: RoPE + RMSNorm + SwiGLU + QK-norm + soft-cap + tied emb")
+    log(f"  L={cfg.n_layer}, H={cfg.n_head}, d={cfg.n_embd}, T={cfg.block_size}")
+
+    # Optimizers
+    opt_muon, opt_adam = create_optimizers(model, device)
+    log(f"  Muon  LR: {MUON_LR}, Momentum: {MUON_MOMENTUM}")
+    log(f"  AdamW LR: {ADAM_LR}, WD: {WEIGHT_DECAY}")
+
+    scaler = torch.amp.GradScaler("cuda", enabled=(dtype == torch.float16))
+
+    # Resume
+    start_step = 0
+    best_val = float("inf")
+    resume_path = None
+    if args.resume_best and BEST_CKPT.exists():
+        resume_path = BEST_CKPT
+    elif args.resume and LATEST_CKPT.exists():
+        resume_path = LATEST_CKPT
+
+    if resume_path:
+        log(f"\nResume: {resume_path}")
+        ckpt = torch.load(resume_path, map_location=device, weights_only=False)
+        if ckpt.get("version") != "v5":
+            log("UYARI: V5 olmayan checkpoint!")
+        model.load_state_dict(ckpt["model"])
+        opt_muon.load_state_dict(ckpt["opt_muon"])
+        opt_adam.load_state_dict(ckpt["opt_adam"])
+        if "scaler" in ckpt:
+            scaler.load_state_dict(ckpt["scaler"])
+        start_step = ckpt["step"] + 1
+        best_val = ckpt.get("best_val", float("inf"))
+        log(f"  step={start_step}, best_val={best_val:.4f}")
+
+    # Compile
+    if args.compile:
+        log("torch.compile baslatiliyor...")
+        torch._dynamo.config.suppress_errors = False
+        os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR",
+                              str(OUT_DIR / "_inductor_cache"))
+        model = torch.compile(model, mode="default", dynamic=False)
+
+    # Sinyal
+    interrupt_flag = {"stop": False}
+    def signal_handler(sig, frame):
+        if interrupt_flag["stop"]:
+            log("\n[!] Ikinci Ctrl+C, cikiyor.")
+            sys.exit(1)
+        interrupt_flag["stop"] = True
+        log("\n[!] Ctrl+C alindi, kaydedilip cikilacak.")
+    signal.signal(signal.SIGINT, signal_handler)
+
+    total_tokens_per_step = BATCH_SIZE * GRAD_ACCUM_STEPS * MODEL_CONFIG["block_size"]
+    log(f"\nEgitim basliyor:")
+    log(f"  Step araligi: {start_step} → {MAX_STEPS}")
+    log(f"  Etkin batch:  {BATCH_SIZE * GRAD_ACCUM_STEPS}")
+    log(f"  Token/step:   {total_tokens_per_step:,}")
+    log(f"  Toplam token: {MAX_STEPS * total_tokens_per_step / 1e9:.1f}B")
+    log(f"  Curriculum: P1[0-{int(PHASE1_END*100)}%]  "
+        f"P2[{int(PHASE1_END*100)}-{int(PHASE2_END*100)}%]  "
+        f"P3[{int(PHASE2_END*100)}-100%]")
+
+    t_start = time.time()
+    step_t0 = time.time()
+    step = start_step
+    last_phase = -1
+    stage_hits = [0, 0, 0]
+
+    try:
+        while step < MAX_STEPS:
+            step_ref["step"] = step
+            phase = get_phase(step)
+            if phase != last_phase:
+                mix = PHASE_MIX[phase]
+                log(f"\n>>> FAZ {phase} basliyor (step {step}): "
+                    f"stage1={mix[0]:.0%}, stage2={mix[1]:.0%}, stage3={mix[2]:.0%}")
+                last_phase = phase
+
+            # LR
+            lr_factor = get_lr_factor(step)
+            muon_lr = MUON_LR * lr_factor
+            adam_lr = ADAM_LR * lr_factor
+            for pg in opt_muon.param_groups:
+                pg["lr"] = muon_lr
+            for pg in opt_adam.param_groups:
+                pg["lr"] = adam_lr
+
+            # Grad accumulation
+            opt_muon.zero_grad(set_to_none=True)
+            opt_adam.zero_grad(set_to_none=True)
+            loss_accum = 0.0
+            for _ in range(GRAD_ACCUM_STEPS):
+                (x, y), stage_idx = prefetch.get_batch()
+                stage_hits[stage_idx] += 1
+                with ctx:
+                    _, loss = model(x, y)
+                    loss = loss / GRAD_ACCUM_STEPS
+                scaler.scale(loss).backward()
+                loss_accum += loss.item()
+
+            scaler.unscale_(opt_adam)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
+            opt_muon.step()
+            scaler.step(opt_adam)
+            scaler.update()
+
+            # Log
+            if step % LOG_INTERVAL == 0:
+                dt = time.time() - step_t0
+                tps = (LOG_INTERVAL * total_tokens_per_step) / dt if step > start_step else 0
+                step_t0 = time.time()
+                elapsed_min = (time.time() - t_start) / 60
+                total_hits = sum(stage_hits) or 1
+                mix_str = "/".join(f"{h*100//total_hits}" for h in stage_hits)
+                log(f"step {step:>6} | P{phase} | loss {loss_accum:.4f} | "
+                    f"muon {muon_lr:.2e} adam {adam_lr:.2e} | "
+                    f"{tps/1e3:.0f}K tok/s | mix {mix_str} | {elapsed_min:.1f}m")
+                stage_hits = [0, 0, 0]
+
+            # Eval
+            if step > start_step and step % EVAL_INTERVAL == 0:
+                losses = estimate_loss(model, val_loader,
+                                       [stage1, stage2, stage3], ctx, EVAL_ITERS)
+                log(f"  >>> EVAL: val {losses['val']:.4f}  "
+                    f"s1 {losses['stage1']:.4f}  s2 {losses['stage2']:.4f}  "
+                    f"s3 {losses['stage3']:.4f}")
+                if losses["val"] < best_val:
+                    best_val = losses["val"]
+                    state = build_state(model, opt_muon, opt_adam, scaler, step, best_val)
+                    atomic_save(state, BEST_CKPT)
+                    log(f"  >>> BEST kaydedildi (val {best_val:.4f})")
+
+            # Save
+            if step > start_step and step % SAVE_INTERVAL == 0:
+                state = build_state(model, opt_muon, opt_adam, scaler, step, best_val)
+                atomic_save(state, LATEST_CKPT)
+
+            # Sample
+            if step > start_step and step % SAMPLE_INTERVAL == 0:
+                for prompt in ["Türkiye", "Yapay zeka", "Bu çalışmada"]:
+                    text = sample_text(model, tokenizer, device, ctx,
+                                       prompt=prompt, max_new_tokens=80)
+                    log(f"  [sample] {text!r}")
+
+            # Time
+            if args.max_time and (time.time() - t_start) / 60 >= args.max_time:
+                log(f"\n[time] {args.max_time} dakika doldu, kaydedilip cikiliyor.")
+                break
+
+            if interrupt_flag["stop"]:
+                break
+
+            step += 1
+
+    finally:
+        log("\nSon checkpoint yaziliyor...")
+        state = build_state(model, opt_muon, opt_adam, scaler, step, best_val)
+        atomic_save(state, LATEST_CKPT)
+        log(f"  latest_ckpt.pt → step {step}, best_val {best_val:.4f}")
+        prefetch.close()
+
+    log(f"\n[DONE] Step {step}/{MAX_STEPS}. Best val: {best_val:.4f}")
+
+
+if __name__ == "__main__":
+    main()