Upload 05_train_v5_tpu.py with huggingface_hub

05_train_v5_tpu.py

@@ -0,0 +1,533 @@
+"""
+V5 Egitim — TPU v5e-8 (Kaggle veya GCP) versiyonu.
+
+Differences from 05_train_v5.py (CUDA):
+  - torch_xla device (xla_device)
+  - Native torch.optim.Muon (PyTorch 2.11+)
+  - xmp.spawn ile multi-chip data parallelism
+  - No Liger (CUDA-only, skip)
+  - No torch.compile (XLA already optimizes)
+  - bf16 native (TPU strength)
+  - Periodic HF push (Kaggle 9h session limit icin)
+
+Kullanim:
+  # Kaggle TPU v5e-8 notebook'ta:
+  python 05_train_v5_tpu.py --resume --hf-user musabc \
+      --hf-push-every 500 --max-time 510  # 8.5 saat (Kaggle 9h limit)
+
+NOT:
+  - Single-host v5e-8: 8 chip, xmp.spawn ile 8 process
+  - Checkpoint format CUDA ile uyumlu — resume calisir
+  - Muon state da otomatik portluyor (state_dict same)
+"""
+
+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
+import torch.nn.functional as F
+
+# TPU imports
+import torch_xla.core.xla_model as xm
+import torch_xla.distributed.parallel_loader as pl
+import torch_xla.distributed.xla_multiprocessing as xmp
+import torch_xla.runtime as xr
+
+# Liger'i kapat (CUDA-only)
+os.environ["NANOGPT_NO_LIGER"] = "1"
+
+# Model import (NANOGPT_NO_LIGER sonrasi)
+from model_v5 import GPTV5, GPTConfigV5
+
+# Native Muon (PyTorch 2.11+)
+try:
+    from torch.optim import Muon as TorchMuon
+    HAS_NATIVE_MUON = True
+except ImportError:
+    HAS_NATIVE_MUON = False
+    print("! torch.optim.Muon yok. PyTorch 2.11+ gerek")
+    print("  Fallback: muon.py kullanilacak (yavas olabilir XLA'da)")
+    from muon import Muon as TorchMuon
+
+
+# ============================================================
+# Konfigurasyon — TPU v5e-8 (8 chip, 16GB HBM/chip)
+# ============================================================
+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,
+)
+
+# Per-chip batch — 8 chip ile global etkin batch 520'ye yakin
+# 16GB per chip: bs=8, T=2048, activations ~10GB rahatca sigar
+BATCH_SIZE_PER_CHIP = 8
+GRAD_ACCUM_STEPS = 8           # 8 chip × bs 8 × accum 8 = 512 etkin batch
+MAX_STEPS = 20_000
+LOG_INTERVAL = 10
+EVAL_INTERVAL = 400
+EVAL_ITERS = 60
+SAVE_INTERVAL = 500            # daha sik kayit (Kaggle session limiti icin)
+HF_PUSH_INTERVAL = 500          # her 500 step HF'e push
+SAMPLE_INTERVAL = 2000
+
+MUON_LR = 0.022
+ADAM_LR = 3.5e-4
+MIN_LR_RATIO = 0.1
+WARMUP_STEPS = 1000
+LR_DECAY_STEPS = 20_000
+
+WEIGHT_DECAY = 0.1
+ADAM_BETA1 = 0.9
+ADAM_BETA2 = 0.95
+MUON_MOMENTUM = 0.95
+GRAD_CLIP = 1.0
+
+PHASE1_END = 0.55
+PHASE2_END = 0.85
+PHASE_MIX = {
+    1: (0.25, 0.65, 0.10),
+    2: (0.15, 0.55, 0.30),
+    3: (0.05, 0.25, 0.70),
+}
+# ============================================================
+
+LATEST_CKPT = OUT_DIR / "latest_ckpt.pt"
+BEST_CKPT = OUT_DIR / "best_ckpt.pt"
+LOG_FILE = OUT_DIR / "train_tpu.log"
+
+
+def get_lr_factor(step):
+    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):
+    p = step / max(MAX_STEPS, 1)
+    if p < PHASE1_END:
+        return 1
+    if p < PHASE2_END:
+        return 2
+    return 3
+
+
+def log_rank0(msg, rank=0):
+    if rank != 0:
+        return
+    print(msg, flush=True)
+    try:
+        with open(LOG_FILE, "a", encoding="utf-8") as f:
+            f.write(msg + "\n")
+    except Exception:
+        pass
+
+
+# =====================================================================
+# Data
+# =====================================================================
+class StageLoader:
+    def __init__(self, bin_path, block_size, batch_size, rank=0, world_size=1):
+        self.data = np.memmap(bin_path, dtype=np.uint16, mode="r")
+        self.block_size = block_size
+        self.batch_size = batch_size
+        self.rank = rank
+        self.world_size = world_size
+        self.n_tokens = len(self.data)
+        # Her rank farkli sample alsin diye seed offset
+        self.rng = np.random.default_rng(seed=42 + rank)
+        if rank == 0:
+            print(f"  {bin_path.name}: {self.n_tokens:,} token "
+                  f"(~{self.n_tokens*2/1e9:.2f} GB)")
+
+    def get_batch(self, device):
+        bs, T = self.batch_size, self.block_size
+        ix = self.rng.integers(0, self.n_tokens - 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).to(device)
+        y = torch.from_numpy(y_np).to(device)
+        return x, y
+
+
+class MultiStageLoader:
+    def __init__(self, stage_loaders, rng=None):
+        self.loaders = stage_loaders
+        self.rng = rng or np.random.default_rng()
+
+    def get_batch(self, phase, device):
+        mix = PHASE_MIX[phase]
+        idx = self.rng.choice(len(self.loaders), p=mix)
+        x, y = self.loaders[idx].get_batch(device)
+        return (x, y), idx
+
+
+# =====================================================================
+# Eval
+# =====================================================================
+@torch.no_grad()
+def estimate_loss(model, val_loader, train_loaders, device, eval_iters):
+    out = {}
+    model.eval()
+    # Val
+    losses = []
+    for _ in range(eval_iters):
+        x, y = val_loader.get_batch(device)
+        _, loss = model(x, y)
+        losses.append(loss.item())
+    out["val"] = sum(losses) / len(losses)
+    # Per-stage (kucuk sample)
+    n_small = max(eval_iters // 4, 8)
+    for i, ld in enumerate(train_loaders, start=1):
+        losses = []
+        for _ in range(n_small):
+            x, y = ld.get_batch(device)
+            _, loss = model(x, y)
+            losses.append(loss.item())
+        out[f"stage{i}"] = sum(losses) / len(losses)
+    model.train()
+    return out
+
+
+# =====================================================================
+# Checkpointing
+# =====================================================================
+def atomic_save(state, path):
+    tmp = path.with_suffix(path.suffix + ".tmp")
+    # CPU'ya kopyala (XLA tensor'lari)
+    cpu_state = {}
+    for k, v in state.items():
+        if isinstance(v, dict):
+            cpu_state[k] = {kk: vv.cpu() if torch.is_tensor(vv) else vv
+                            for kk, vv in v.items()}
+        elif torch.is_tensor(v):
+            cpu_state[k] = v.cpu()
+        else:
+            cpu_state[k] = v
+    torch.save(cpu_state, tmp)
+    if path.exists():
+        path.unlink()
+    tmp.rename(path)
+
+
+def build_state(model, opt_muon, opt_adam, step, best_val):
+    return {
+        "model": model.state_dict(),
+        "opt_muon": opt_muon.state_dict(),
+        "opt_adam": opt_adam.state_dict(),
+        "step": step,
+        "best_val": best_val,
+        "config": MODEL_CONFIG,
+        "version": "v5",
+    }
+
+
+def push_to_hf(ckpt_path, hf_user, hf_token, repo_suffix="ckpts"):
+    """Kaggle session bitiminde HF'e push."""
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi(token=hf_token)
+        repo_id = f"{hf_user}/nanogpt-tr-v5-{repo_suffix}"
+        api.upload_file(
+            path_or_fileobj=str(ckpt_path),
+            path_in_repo=ckpt_path.name,
+            repo_id=repo_id,
+            repo_type="model",
+            commit_message=f"upload {ckpt_path.name}",
+        )
+        return True
+    except Exception as e:
+        print(f"  ! HF push hatasi: {e}")
+        return False
+
+
+# =====================================================================
+# Optimizer
+# =====================================================================
+def create_optimizers(model, rank=0):
+    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)
+
+    if rank == 0:
+        n_muon = sum(p.numel() for p in muon_params)
+        n_adam = sum(p.numel() for p in adam_params_unique)
+        log_rank0(f"  Muon params:  {n_muon/1e6:.2f}M ({len(muon_params)} tensor)")
+        log_rank0(f"  AdamW params: {n_adam/1e6:.2f}M ({len(adam_params_unique)} tensor)")
+        log_rank0(f"  Muon backend: {'NATIVE torch.optim.Muon' if HAS_NATIVE_MUON else 'custom muon.py'}")
+
+    # Native Muon (PyTorch 2.11+)
+    opt_muon = TorchMuon(
+        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,
+    )
+    return opt_muon, opt_adam
+
+
+# =====================================================================
+# Main training (per-process, multi-chip)
+# =====================================================================
+def train_fn(rank, args):
+    device = xm.xla_device()
+    world_size = xr.world_size()
+
+    if rank == 0:
+        OUT_DIR.mkdir(parents=True, exist_ok=True)
+        log_rank0(f"TPU world size: {world_size}")
+        log_rank0(f"Device: {device}")
+        log_rank0(f"Per-chip batch: {BATCH_SIZE_PER_CHIP}")
+        log_rank0(f"Global effective batch: "
+                  f"{BATCH_SIZE_PER_CHIP * world_size * GRAD_ACCUM_STEPS}")
+
+    # Data
+    if rank == 0:
+        log_rank0("\nData yukleniyor...")
+    stage1 = StageLoader(DATA_DIR / "v5_stage1.bin",
+                          MODEL_CONFIG["block_size"], BATCH_SIZE_PER_CHIP,
+                          rank=rank, world_size=world_size)
+    stage2 = StageLoader(DATA_DIR / "v5_stage2.bin",
+                          MODEL_CONFIG["block_size"], BATCH_SIZE_PER_CHIP,
+                          rank=rank, world_size=world_size)
+    stage3 = StageLoader(DATA_DIR / "v5_stage3.bin",
+                          MODEL_CONFIG["block_size"], BATCH_SIZE_PER_CHIP,
+                          rank=rank, world_size=world_size)
+    val_loader = StageLoader(DATA_DIR / "v5_val.bin",
+                              MODEL_CONFIG["block_size"], BATCH_SIZE_PER_CHIP,
+                              rank=rank, world_size=world_size)
+    multi = MultiStageLoader([stage1, stage2, stage3])
+
+    # Model
+    if rank == 0:
+        log_rank0("\nModel V5 olusturuluyor...")
+    cfg = GPTConfigV5(**MODEL_CONFIG)
+    model = GPTV5(cfg).to(device)
+    if rank == 0:
+        n_params = model.num_params()
+        log_rank0(f"  Toplam: {n_params/1e6:.2f}M param")
+
+    opt_muon, opt_adam = create_optimizers(model, rank=rank)
+
+    # Resume
+    start_step = 0
+    best_val = float("inf")
+    resume_path = LATEST_CKPT if args.resume and LATEST_CKPT.exists() else None
+    if resume_path:
+        if rank == 0:
+            log_rank0(f"\nResume: {resume_path}")
+        ckpt = torch.load(resume_path, map_location="cpu", weights_only=False)
+        # State dict yukle, sonra device'a tasi
+        model.load_state_dict(ckpt["model"])
+        # XLA tensor'a tasinmis oldu (model.to(device) ile)
+        try:
+            opt_muon.load_state_dict(ckpt["opt_muon"])
+            opt_adam.load_state_dict(ckpt["opt_adam"])
+        except Exception as e:
+            if rank == 0:
+                log_rank0(f"  ! Optimizer state yuklenemedi: {e}")
+                log_rank0(f"    Sifirdan optimizer state ile devam edilecek")
+        start_step = ckpt["step"] + 1
+        best_val = ckpt.get("best_val", float("inf"))
+        if rank == 0:
+            log_rank0(f"  step={start_step}, best_val={best_val:.4f}")
+
+    # Sinyal
+    interrupt_flag = {"stop": False}
+    def signal_handler(sig, frame):
+        interrupt_flag["stop"] = True
+        if rank == 0:
+            log_rank0("\n[!] Ctrl+C, kaydedilip cikiliyor.")
+    signal.signal(signal.SIGINT, signal_handler)
+
+    # Egitim
+    if rank == 0:
+        token_per_step = BATCH_SIZE_PER_CHIP * world_size * GRAD_ACCUM_STEPS * MODEL_CONFIG["block_size"]
+        log_rank0(f"\nEgitim basliyor:")
+        log_rank0(f"  Step araligi: {start_step} -> {MAX_STEPS}")
+        log_rank0(f"  Token/step: {token_per_step:,}")
+
+    t_start = time.time()
+    step_t0 = time.time()
+    step = start_step
+    last_phase = -1
+    stage_hits = [0, 0, 0]
+    last_hf_push = 0
+
+    try:
+        while step < MAX_STEPS:
+            phase = get_phase(step)
+            if phase != last_phase and rank == 0:
+                mix = PHASE_MIX[phase]
+                log_rank0(f"\n>>> FAZ {phase} basliyor (step {step}): "
+                          f"s1={mix[0]:.0%}, s2={mix[1]:.0%}, s3={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
+
+            # Gradient 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 = multi.get_batch(phase, device)
+                stage_hits[stage_idx] += 1
+                with torch.autocast(device_type="xla", dtype=torch.bfloat16):
+                    _, loss = model(x, y)
+                    loss = loss / GRAD_ACCUM_STEPS
+                loss.backward()
+                loss_accum += loss.item()
+
+            # Grad clip + step (xm.optimizer_step ile cross-replica sync)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
+            xm.optimizer_step(opt_muon)
+            xm.optimizer_step(opt_adam)
+
+            # XLA flush — kritik!
+            xm.mark_step()
+
+            # Log
+            if step % LOG_INTERVAL == 0 and rank == 0:
+                dt = time.time() - step_t0
+                tps = (LOG_INTERVAL * BATCH_SIZE_PER_CHIP * world_size *
+                       GRAD_ACCUM_STEPS * MODEL_CONFIG["block_size"]) / 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_rank0(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 (rank 0)
+            if step > start_step and step % EVAL_INTERVAL == 0 and rank == 0:
+                losses = estimate_loss(model, val_loader,
+                                       [stage1, stage2, stage3], device, EVAL_ITERS)
+                log_rank0(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, step, best_val)
+                    atomic_save(state, BEST_CKPT)
+                    log_rank0(f"  >>> BEST kaydedildi (val {best_val:.4f})")
+
+            # Save (rank 0)
+            if step > start_step and step % SAVE_INTERVAL == 0 and rank == 0:
+                state = build_state(model, opt_muon, opt_adam, step, best_val)
+                atomic_save(state, LATEST_CKPT)
+
+            # HF push (Kaggle session limit icin)
+            if (args.hf_push_every and step > last_hf_push + args.hf_push_every
+                    and rank == 0 and args.hf_user and args.hf_token):
+                log_rank0(f"  >>> HF push step {step}...")
+                push_to_hf(LATEST_CKPT, args.hf_user, args.hf_token)
+                last_hf_push = step
+
+            # Time
+            if args.max_time and (time.time() - t_start) / 60 >= args.max_time:
+                if rank == 0:
+                    log_rank0(f"\n[time] {args.max_time} dakika doldu, "
+                              f"kaydedilip cikiliyor.")
+                break
+
+            if interrupt_flag["stop"]:
+                break
+
+            step += 1
+
+    finally:
+        if rank == 0:
+            log_rank0("\nSon checkpoint yaziliyor...")
+            state = build_state(model, opt_muon, opt_adam, step, best_val)
+            atomic_save(state, LATEST_CKPT)
+            log_rank0(f"  latest_ckpt.pt -> step {step}, best_val {best_val:.4f}")
+
+            # Final HF push
+            if args.hf_user and args.hf_token:
+                log_rank0("Final HF push...")
+                push_to_hf(LATEST_CKPT, args.hf_user, args.hf_token)
+                if BEST_CKPT.exists():
+                    push_to_hf(BEST_CKPT, args.hf_user, args.hf_token)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--resume", action="store_true")
+    parser.add_argument("--max-time", type=int, default=510,
+                        help="Maks dakika (Kaggle 9h limit icin 510)")
+    parser.add_argument("--max-steps", type=int, default=None)
+    parser.add_argument("--hf-user", type=str, default=None,
+                        help="HF push icin kullanici (kaydet+push)")
+    parser.add_argument("--hf-token", type=str,
+                        default=os.environ.get("HF_TOKEN"),
+                        help="HF token (env HF_TOKEN'dan da okur)")
+    parser.add_argument("--hf-push-every", type=int, default=500,
+                        help="Her N step'te HF'e push (0=kapali)")
+    args = parser.parse_args()
+
+    if args.max_steps:
+        global MAX_STEPS, LR_DECAY_STEPS
+        MAX_STEPS = args.max_steps
+        LR_DECAY_STEPS = args.max_steps
+
+    # Multi-chip spawn — TPU v5e-8 icin 8 process
+    # Kaggle TPU genelde tek-host, xmp.spawn yeter
+    xmp.spawn(train_fn, args=(args,))
+
+
+if __name__ == "__main__":
+    main()