05_train_v5.py

"""
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(max-autotune)
  - Liger Kernel (Triton fused RMSNorm / SwiGLU / chunked CE)
  - Flash-SDPA + cuDNN SDPA (Blackwell aware)
  - Async prefetcher per stage
  - Multi-stage data loader (weighted sampling, progresif)

Kurulum (Lightning AI'da):
    pip install liger-kernel

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
#   Default: H100 80GB (Thunder Compute) — bs=40, ga=13, etkin batch 520
#   RTX PRO 6000 96GB için:        --batch 48 --grad-accum 11
#   A100 80GB (eski) için:         --batch 36 --grad-accum 14
# ============================================================
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,
)

# H100 80GB için: bs=40, T=2048 → activations ~55-60GB, weights+opt ~5GB
# Liger Kernel ile ~10GB tasarruf → toplam ~55GB (rahat marj)
# 40 = 8*5 tensor core dostu
BATCH_SIZE = 40
GRAD_ACCUM_STEPS = 13          # etkin batch = 520, token/step ≈ 1.06M
MAX_STEPS = 20_000             # ~21.3B token training (Modern overtraining, 106x 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("--compile-mode", type=str, default="max-autotune",
                        choices=["default", "reduce-overhead",
                                 "max-autotune", "max-autotune-no-cudagraphs"],
                        help="torch.compile mode (default: max-autotune)")
    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 + cuDNN SDPA backend
        try:
            torch.backends.cuda.enable_flash_sdp(True)
            torch.backends.cuda.enable_mem_efficient_sdp(True)
            torch.backends.cuda.enable_math_sdp(False)
            # cuDNN 9.5+ Blackwell-aware attention (PyTorch 2.5+)
            if hasattr(torch.backends.cuda, "enable_cudnn_sdp"):
                torch.backends.cuda.enable_cudnn_sdp(True)
        except Exception:
            pass
        # Allocator — fragmentasyon + tuning
        os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF",
                              "expandable_segments:True,max_split_size_mb:512")
        # Inductor autotuning genisletilmis
        try:
            import torch._inductor.config as ind_cfg
            ind_cfg.coordinate_descent_tuning = True
            ind_cfg.triton.unique_kernel_names = True
            ind_cfg.fx_graph_cache = True
        except Exception:
            pass
        log("Perf: TF32 ON, cudnn.benchmark ON, Flash+cuDNN SDPA ON")
        # Liger durumu
        try:
            from model_v5 import HAS_LIGER
            log(f"Liger Kernel: {'ON (Triton fused ops)' if HAS_LIGER else 'OFF'}")
        except Exception:
            pass
        # 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 — max-autotune (Blackwell icin Triton variant arama)
    if args.compile:
        compile_mode = args.compile_mode
        log(f"torch.compile baslatiliyor (mode={compile_mode})...")
        log("  (ilk 200-500 step yavas — kernel autotuning, sabir)")
        torch._dynamo.config.suppress_errors = False
        os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR",
                              str(OUT_DIR / "_inductor_cache"))
        model = torch.compile(model, mode=compile_mode, 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()