train_300m.py

"""
LUNA 300M — Config-Driven Training with Validation & HF Upload
===============================================================
Same proven architecture as LUNA 100M, scaled to ~303M parameters.
Adds: text-generation validation every N steps + HF checkpoint upload.

Usage:
    python train_300m.py                                          # defaults
    python train_300m.py --config train_config_300m.yaml          # explicit
    python train_300m.py --data_path /workspace/data/litdata_pretrain_final
"""

import os
import gc
import sys
import math
import time
import json
import shutil
import argparse
import yaml
import psutil
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.amp import autocast, GradScaler
from pathlib import Path

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")


# ─── Model ────────────────────────────────────────────────────────────────────

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len=1024):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        t = torch.arange(max_seq_len).float()
        freqs = torch.einsum("i,j->ij", t, inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        self.register_buffer("cos_cached", emb.cos())
        self.register_buffer("sin_cached", emb.sin())

    def forward(self, seq_len):
        return self.cos_cached[:seq_len], self.sin_cached[:seq_len]


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


def apply_rotary(x, cos, sin):
    c = cos.unsqueeze(0).unsqueeze(0)
    s = sin.unsqueeze(0).unsqueeze(0)
    return x * c + rotate_half(x) * s


class CausalSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size, rotary_pct=0.25):
        super().__init__()
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.rot_dim = int(self.head_dim * rotary_pct)
        self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=True)
        self.c_proj = nn.Linear(n_embd, n_embd, bias=True)
        self.rotary = RotaryEmbedding(self.rot_dim, block_size)

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x).reshape(B, T, 3, self.n_head, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)
        cos, sin = self.rotary(T)
        q = torch.cat([apply_rotary(q[..., :self.rot_dim], cos, sin), q[..., self.rot_dim:]], dim=-1)
        k = torch.cat([apply_rotary(k[..., :self.rot_dim], cos, sin), k[..., self.rot_dim:]], dim=-1)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.c_proj(y.transpose(1, 2).contiguous().view(B, T, C))


class MLP(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.fc = nn.Linear(n_embd, 4 * n_embd, bias=True)
        self.gelu = nn.GELU()
        self.proj = nn.Linear(4 * n_embd, n_embd, bias=True)

    def forward(self, x):
        return self.proj(self.gelu(self.fc(x)))


class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, block_size)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = MLP(n_embd)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x


class LUNAModel(nn.Module):
    def __init__(self, vocab_size, block_size, n_layer, n_embd, n_head):
        super().__init__()
        self.wte = nn.Embedding(vocab_size, n_embd)
        self.blocks = nn.ModuleList([Block(n_embd, n_head, block_size) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        self.lm_head.weight = self.wte.weight  # tie
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            m.weight.data.normal_(mean=0.0, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                m.bias.data.zero_()

    def forward(self, idx, targets=None, return_logits=True):
        x = self.wte(idx)
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
            if not return_logits:
                logits = None
        return logits, loss

    @property
    def num_params(self):
        return sum(p.numel() for p in self.parameters()) - self.wte.weight.numel()


# ─── Dataset ──────────────────────────────────────────────────────────────────

class LitDataDataset(torch.utils.data.Dataset):
    def __init__(self, data_path: str, block_size: int = 1024):
        import struct, numpy as np
        self.block_size = block_size
        self.data_path = Path(data_path)
        with open(self.data_path / "index.json") as f:
            idx = json.load(f)
        self.chunks_meta = idx["chunks"]
        self._cum_blocks = []
        total = 0
        for c in self.chunks_meta:
            n = c["dim"] // (block_size + 1)
            total += n
            self._cum_blocks.append(total)
        self.total_blocks = total
        self._chunk_cache = {}

    def _load_chunk(self, chunk_idx: int):
        if chunk_idx in self._chunk_cache:
            return self._chunk_cache[chunk_idx]
        import struct, numpy as np
        meta = self.chunks_meta[chunk_idx]
        with open(self.data_path / meta["filename"], "rb") as f:
            raw = f.read()
        num_items = struct.unpack_from("<I", raw, 0)[0]
        header_bytes = (num_items + 2) * 4
        tokens = torch.from_numpy(np.frombuffer(raw[header_bytes:], dtype=np.int32).copy())
        if len(self._chunk_cache) >= 4:
            del self._chunk_cache[next(iter(self._chunk_cache))]
        self._chunk_cache[chunk_idx] = tokens
        return tokens

    def __len__(self):
        return self.total_blocks

    def __getitem__(self, idx):
        chunk_idx = 0
        for i, cum in enumerate(self._cum_blocks):
            if idx < cum:
                chunk_idx = i
                break
        prev = self._cum_blocks[chunk_idx - 1] if chunk_idx > 0 else 0
        tokens = self._load_chunk(chunk_idx)
        s = (idx - prev) * (self.block_size + 1)
        e = s + self.block_size + 1
        chunk = tokens[s:e]
        if len(chunk) < self.block_size + 1:
            pad = torch.zeros(self.block_size + 1, dtype=torch.int32)
            pad[:len(chunk)] = chunk
            chunk = pad
        chunk = chunk.long()
        return chunk[:self.block_size], chunk[1:self.block_size + 1]


# ─── Hardware Detection ────────────────────────────────────────────────────────

def probe_hardware():
    info = {
        "cpu_cores": os.cpu_count() or 4,
        "ram_gb": psutil.virtual_memory().total / 1024**3,
    }
    if torch.cuda.is_available():
        props = torch.cuda.get_device_properties(0)
        info.update({
            "device": "cuda",
            "gpu_name": props.name,
            "vram_gb": props.total_memory / 1024**3,
            "sm_major": props.major,
        })
        if props.major >= 8:
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True
            info["precision"] = "bf16"
            info["dtype"] = torch.bfloat16
        else:
            info["precision"] = "fp16"
            info["dtype"] = torch.float16
    else:
        info.update({
            "device": "cpu",
            "gpu_name": "CPU",
            "vram_gb": 0,
            "sm_major": 0,
            "precision": "fp32",
            "dtype": torch.float32,
        })
    return info


def probe_max_batch(model, device, dtype, seq_len, vocab_size, max_search=4096, grad_accum_sim=4):
    """Binary search for max micro_batch. Safety: x0.70."""
    tmp_opt = torch.optim.AdamW(model.parameters(), lr=1e-4)
    lo, hi, best = 1, max_search, 1
    while lo <= hi:
        mid = (lo + hi) // 2
        try:
            torch.cuda.empty_cache(); gc.collect()
            tmp_opt.zero_grad(set_to_none=True)
            for _ in range(grad_accum_sim):
                x = torch.randint(0, vocab_size, (mid, seq_len), device=device)
                t = torch.randint(0, vocab_size, (mid, seq_len), device=device)
                with autocast(device_type="cuda", dtype=dtype):
                    _, loss = model(x, t, return_logits=False)
                    loss = loss / grad_accum_sim
                loss.backward()
                del x, t, loss
            tmp_opt.step()
            tmp_opt.zero_grad(set_to_none=True)
            best = mid; lo = mid + 1
            torch.cuda.empty_cache()
        except torch.cuda.OutOfMemoryError:
            try: del x, t, loss
            except: pass
            torch.cuda.empty_cache()
            tmp_opt.zero_grad(set_to_none=True)
            hi = mid - 1
        except RuntimeError as e:
            err_str = str(e).lower()
            if ("out of memory" in err_str
                    or "no valid triton configs" in err_str
                    or "out of resource" in err_str
                    or "inductorerror" in type(e).__name__.lower()):
                try: del x, t, loss
                except: pass
                torch.cuda.empty_cache()
                tmp_opt.zero_grad(set_to_none=True)
                hi = mid - 1
            else:
                raise
        except Exception as e:
            err_str = str(e).lower()
            if ("out of resource" in err_str
                    or "no valid triton configs" in err_str):
                try: del x, t, loss
                except: pass
                torch.cuda.empty_cache()
                tmp_opt.zero_grad(set_to_none=True)
                hi = mid - 1
            else:
                raise
    del tmp_opt; torch.cuda.empty_cache(); gc.collect()
    safe = max(1, int(best * 0.70))
    print(f"  Probe found max_batch={best}, using {safe} (70% safety, tested with {grad_accum_sim} accum steps)")
    return safe


# ─── LR Schedule ──────────────────────────────────────────────────────────────

def cosine_lr(step, warmup, total, lr_max, lr_min):
    if step < warmup:
        return lr_max * (step + 1) / warmup
    p = (step - warmup) / max(1, total - warmup)
    return lr_min + 0.5 * (1 + math.cos(math.pi * p)) * (lr_max - lr_min)


# ─── Validation (Text Generation) ─────────────────────────────────────────────

VALIDATION_PROMPTS = [
    "The theory of general relativity, proposed by Albert Einstein, describes gravity as",
    "The process of photosynthesis in green plants converts sunlight into chemical energy by",
    "The French Revolution, which began in 1789, was primarily caused by",
    "In machine learning, neural networks learn to recognize patterns by",
    "The philosophical concept of free will suggests that humans can",
    "The law of supply and demand states that when the price of a good increases,",
    "In Shakespeare's Hamlet, the protagonist faces a moral dilemma about whether to",
    "Water molecules consist of two hydrogen atoms bonded to one oxygen atom, giving it the property of",
    "The Pythagorean theorem states that in a right triangle, the square of the hypotenuse equals",
    "The Amazon rainforest, often called the lungs of the Earth, spans across several countries and",
]


@torch.no_grad()
def run_validation(model, tokenizer, device, dtype, step, seq_len=1024, max_new=64):
    """Generate text from 10 diverse prompts to visually track quality progression."""
    model.eval()
    print(f"\n{'=' * 72}")
    print(f"  VALIDATION @ step {step}  ({len(VALIDATION_PROMPTS)} prompts, {max_new} tokens each)")
    print(f"{'=' * 72}")

    for i, prompt in enumerate(VALIDATION_PROMPTS):
        input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
        if input_ids.shape[1] > seq_len:
            input_ids = input_ids[:, -seq_len:]

        generated = input_ids.clone()
        for _ in range(max_new):
            ctx = generated[:, -seq_len:] if generated.shape[1] > seq_len else generated
            with autocast(device_type=device.type, dtype=dtype, enabled=(device.type == "cuda")):
                logits, _ = model(ctx)
            next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
            generated = torch.cat([generated, next_token], dim=1)
            if next_token.item() == tokenizer.eos_token_id:
                break

        output = tokenizer.decode(generated[0], skip_special_tokens=True)
        continuation = output[len(prompt):].strip()[:200]
        print(f"\n  [{i + 1:2d}] {prompt}")
        print(f"       -> {continuation}")

    print(f"\n{'=' * 72}\n")
    model.train()


# ─── HuggingFace Upload ───────────────────────────────────────────────────────

def upload_to_hf(out_dir: Path, cfg: dict):
    """Upload all checkpoints + config + tokenizer to HuggingFace."""
    hf_repo = cfg.get("hf_repo", "")
    if not hf_repo:
        print("  [UPLOAD] No hf_repo configured, skipping.")
        return

    token = os.environ.get("HF_TOKEN", "")
    if not token:
        print("  [UPLOAD] HF_TOKEN not set, skipping. Set it to upload.")
        return

    try:
        from huggingface_hub import HfApi
    except ImportError:
        print("  [UPLOAD] huggingface_hub not installed, skipping.")
        return

    print(f"\n{'=' * 72}")
    print(f"  UPLOADING TO HUGGINGFACE: {hf_repo}")
    print(f"{'=' * 72}")

    api = HfApi(token=token)
    api.create_repo(repo_id=hf_repo, repo_type="model", exist_ok=True)

    # Upload the full output directory (all checkpoints + final)
    print(f"  Uploading {out_dir} ...")
    api.upload_folder(
        repo_id=hf_repo,
        repo_type="model",
        folder_path=str(out_dir),
        path_in_repo="pretrained",
    )
    print(f"  Uploaded model checkpoints.")

    # Upload config
    config_name = "train_config_300m.yaml"
    if os.path.exists(config_name):
        api.upload_file(
            path_or_fileobj=config_name,
            path_in_repo=f"config/{config_name}",
            repo_id=hf_repo,
            repo_type="model",
        )
        print(f"  Uploaded {config_name}")

    # Upload tokenizer
    tok_dir = Path(cfg.get("tokenizer_dir", ""))
    if tok_dir.exists():
        for tf in tok_dir.iterdir():
            if tf.is_file():
                api.upload_file(
                    path_or_fileobj=str(tf),
                    path_in_repo=f"tokenizer/{tf.name}",
                    repo_id=hf_repo,
                    repo_type="model",
                )
        print(f"  Uploaded tokenizer.")

    print(f"  Done -> https://huggingface.co/{hf_repo}")
    print(f"{'=' * 72}\n")


# ─── Config Loading ───────────────────────────────────────────────────────────

def load_config(config_path: str) -> dict:
    with open(config_path, encoding="utf-8") as f:
        raw = yaml.safe_load(f)

    cfg = {
        "auto_config":    raw.get("auto_config", True),
        "data_path":      raw.get("data_path", "Base/data/litdata_pretrain_final"),
        "out_dir":        raw.get("out_dir", "out/pretrain/luna-300m"),
        "tokenizer_dir":  raw.get("tokenizer_dir", "Base/checkpoints/EleutherAI/pythia-160m"),
        "init_model_path": raw.get("init_model_path"),
        # model
        "vocab_size":     raw["model"]["vocab_size"],
        "seq_len":        raw["model"]["seq_len"],
        "n_layer":        raw["model"]["n_layer"],
        "n_embd":         raw["model"]["n_embd"],
        "n_head":         raw["model"]["n_head"],
        # train
        "max_tokens":     raw["train"]["max_tokens"],
        "lr_warmup_steps":raw["train"]["lr_warmup_steps"],
        "save_interval":  raw["train"]["save_interval"],
        "log_interval":   raw["train"]["log_interval"],
        "eval_interval":  raw["train"].get("eval_interval", 100),
        "max_norm":       raw["train"]["max_norm"],
        # optimizer
        "lr":             raw["optimizer"]["lr"],
        "min_lr":         raw["optimizer"]["min_lr"],
        "weight_decay":   raw["optimizer"]["weight_decay"],
        "betas":          tuple(raw["optimizer"]["betas"]),
        "eps":            raw["optimizer"]["eps"],
        # batch
        "global_batch":   raw["batch"]["global_batch"],
        "micro_batch":    raw["batch"]["micro_batch"],
        "grad_accum":     raw["batch"]["grad_accum"],
        # dataloader
        "num_workers":    raw["dataloader"]["num_workers"],
        "pin_memory":     raw["dataloader"]["pin_memory"],
        # hardware
        "precision":      raw["hardware"]["precision"],
        "compile":        raw["hardware"]["compile"],
        # upload
        "hf_repo":        raw.get("upload", {}).get("hf_repo", ""),
        "upload_checkpoints": raw.get("upload", {}).get("upload_checkpoints", True),
    }
    return cfg


def apply_cli_overrides(cfg: dict, cli_args: argparse.Namespace) -> dict:
    for key, val in vars(cli_args).items():
        if key == "config":
            continue
        if val is not None:
            cfg[key] = val
    return cfg


def resolve_auto(cfg: dict, hw: dict) -> dict:
    if not cfg["auto_config"]:
        print("  [CONFIG] auto_config=false -- using manual values as-is")
        cfg.update({"_hw": hw})
        return cfg

    print("  [CONFIG] auto_config=true -- tuning settings to this hardware")
    cfg["precision"] = hw["precision"]
    cfg["_dtype"] = hw["dtype"]

    auto_workers = hw["cpu_cores"] // 2
    max_by_ram = max(0, int(hw["ram_gb"] * 0.25 * 1024 / 268))
    cfg["num_workers"] = min(auto_workers, max_by_ram, hw["cpu_cores"])
    if cfg["num_workers"] == -1:
        cfg["num_workers"] = 0

    cfg["pin_memory"] = hw["ram_gb"] > 16 and hw["device"] == "cuda"
    cfg["_auto_warmup"] = True

    base_global = 120
    cfg["lr"] = cfg["lr"] * math.sqrt(cfg["global_batch"] / base_global)
    cfg["min_lr"] = cfg["min_lr"] * math.sqrt(cfg["global_batch"] / base_global)

    cfg["_hw"] = hw
    return cfg


# ─── Training ─────────────────────────────────────────────────────────────────

SEP = "=" * 72


def train(cfg: dict):
    hw = cfg["_hw"]
    device = torch.device(hw["device"])

    if device.type == "cuda":
        torch.cuda.empty_cache()
        gc.collect()
        free_gb = (torch.cuda.get_device_properties(0).total_memory
                   - torch.cuda.memory_allocated()) / 1024**3
        print(f"  GPU free before model load: {free_gb:.1f} GB")

    if cfg["auto_config"]:
        dtype = hw.get("dtype", torch.float32)
    else:
        dtype = {"bf16": torch.bfloat16, "fp16": torch.float16,
                 "fp32": torch.float32}.get(cfg["precision"], torch.float32)

    # Compute model size label
    total_p = (cfg["vocab_size"] * cfg["n_embd"]
               + cfg["n_layer"] * (
                   2 * cfg["n_embd"]                                    # ln1
                   + cfg["n_embd"] * 3 * cfg["n_embd"] + 3 * cfg["n_embd"]  # c_attn
                   + cfg["n_embd"] * cfg["n_embd"] + cfg["n_embd"]          # c_proj
                   + 2 * cfg["n_embd"]                                    # ln2
                   + cfg["n_embd"] * 4 * cfg["n_embd"] + 4 * cfg["n_embd"]  # mlp fc
                   + 4 * cfg["n_embd"] * cfg["n_embd"] + cfg["n_embd"]      # mlp proj
               )
               + 2 * cfg["n_embd"])  # ln_f
    size_label = f"{total_p / 1e6:.0f}M"

    print(SEP)
    print(f"  LUNA {size_label} - Pretraining")
    print(SEP)
    mode = "AUTO" if cfg["auto_config"] else "MANUAL"
    print(f"  Config mode  : {mode}")
    print(f"  GPU          : {hw['gpu_name']}  ({hw['vram_gb']:.1f} GB)")
    print(f"  RAM          : {hw['ram_gb']:.1f} GB   CPU: {hw['cpu_cores']} cores")
    print(f"  Precision    : {cfg['precision']}   dtype={dtype}")
    print(f"  Workers      : {cfg['num_workers']}   pin_memory={cfg['pin_memory']}")
    print(f"  Model arch   : layers={cfg['n_layer']} embd={cfg['n_embd']} heads={cfg['n_head']}")

    # ── Model ─────────────────────────────────────────────────────────────────
    print(f"\n  Building LUNA-{size_label}...")
    model = LUNAModel(
        vocab_size=cfg["vocab_size"],
        block_size=cfg["seq_len"],
        n_layer=cfg["n_layer"],
        n_embd=cfg["n_embd"],
        n_head=cfg["n_head"],
    ).to(device)

    actual_params = sum(p.numel() for p in model.parameters())
    unique_params = model.num_params
    print(f"  Parameters   : {actual_params:,} total  ({unique_params:,} unique, tied embeddings)")

    # ── Batch sizing (BEFORE torch.compile to avoid Triton kernel issues) ────
    if cfg["auto_config"] and device.type == "cuda":
        print(f"\n  Probing max micro_batch_size (VRAM search)...")
        max_mbs = probe_max_batch(
            model, device, dtype, cfg["seq_len"], cfg["vocab_size"]
        )
        model.apply(model._init_weights)
        torch.cuda.empty_cache(); gc.collect()
        grad_accum = max(1, math.ceil(cfg["global_batch"] / max_mbs))
        effective_batch = max_mbs * grad_accum
        print(f"  AUTO -> micro_batch={max_mbs}, grad_accum={grad_accum}, "
              f"effective_batch={effective_batch}")
    else:
        max_mbs = cfg["micro_batch"]
        grad_accum = cfg["grad_accum"]
        effective_batch = max_mbs * grad_accum
        print(f"\n  MANUAL -> micro_batch={max_mbs}, grad_accum={grad_accum}, "
              f"effective_batch={effective_batch}")

    tokens_per_step = effective_batch * cfg["seq_len"]
    print(f"  Tokens/step : {tokens_per_step:,}")

    # ── torch.compile (after batch probe to avoid Triton shared-memory errors)
    compiled_model = False
    if cfg.get("compile", True) and sys.platform == "linux" and device.type == "cuda":
        try:
            print("  Compiling model (torch.compile)...")
            model = torch.compile(model)
            compiled_model = True
            print("  torch.compile: enabled")
        except Exception as e:
            print(f"  torch.compile: failed ({e}), continuing without")
    else:
        reason = "non-linux" if sys.platform != "linux" else "disabled"
        print(f"  torch.compile: skipped ({reason})")

    # ── Optional init from existing weights ───────────────────────────────────
    init_model_path = cfg.get("init_model_path")
    if init_model_path:
        init_path = Path(init_model_path)
        if not init_path.exists():
            raise FileNotFoundError(f"Initial model checkpoint not found: {init_path}")
        print(f"\n  Init model   : {init_path}")
        init_ckpt = torch.load(init_path, map_location=device, weights_only=True)
        init_state = init_ckpt["model"] if isinstance(init_ckpt, dict) and "model" in init_ckpt else init_ckpt
        raw_model = model._orig_mod if hasattr(model, "_orig_mod") else model
        raw_model.load_state_dict(init_state, strict=True)
        print("  Loaded initial weights from checkpoint")

    # ── Dataset ───────────────────────────────────────────────────────────────
    print(f"\n  Dataset: {cfg['data_path']}")
    dataset = LitDataDataset(cfg["data_path"], block_size=cfg["seq_len"])
    print(f"  Blocks  : {len(dataset):,}   ({len(dataset) * cfg['seq_len']:,} tokens)")

    loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=max_mbs,
        shuffle=True,
        num_workers=cfg["num_workers"],
        pin_memory=cfg["pin_memory"],
        drop_last=True,
        prefetch_factor=4 if cfg["num_workers"] > 0 else None,
        persistent_workers=cfg["num_workers"] > 0,
    )

    # ── Tokenizer (for validation generation) ─────────────────────────────────
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(cfg["tokenizer_dir"])

    # ── Optimiser ─────────────────────────────────────────────────────────────
    try:
        optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=cfg["lr"], weight_decay=cfg["weight_decay"],
            betas=cfg["betas"], eps=cfg["eps"],
            fused=True,
        )
    except TypeError:
        optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=cfg["lr"], weight_decay=cfg["weight_decay"],
            betas=cfg["betas"], eps=cfg["eps"],
        )

    use_scaler = dtype == torch.float16
    scaler = GradScaler(enabled=use_scaler)

    # ── Schedule ──────────────────────────────────────────────────────────────
    total_steps = max(1, cfg["max_tokens"] // tokens_per_step)
    if cfg["auto_config"] and cfg.get("_auto_warmup"):
        warmup_steps = max(50, min(500, total_steps // 20))
    else:
        warmup_steps = min(cfg["lr_warmup_steps"], total_steps)

    out_dir = Path(cfg["out_dir"])
    out_dir.mkdir(parents=True, exist_ok=True)

    # Save model config for eval script
    model_config = {
        "vocab_size": cfg["vocab_size"],
        "seq_len": cfg["seq_len"],
        "n_layer": cfg["n_layer"],
        "n_embd": cfg["n_embd"],
        "n_head": cfg["n_head"],
    }
    with open(out_dir / "model_config.json", "w") as f:
        json.dump(model_config, f, indent=2)

    print(f"\n  max_tokens   : {cfg['max_tokens']:,}")
    print(f"  total_steps  : {total_steps:,}")
    print(f"  warmup_steps : {warmup_steps}")
    print(f"  lr           : {cfg['lr']:.2e} -> {cfg['min_lr']:.2e}")
    print(f"  save every   : {cfg['save_interval']} steps")
    print(f"  eval every   : {cfg['eval_interval']} steps (10 generation prompts)")
    print(f"  out_dir      : {out_dir}")
    print(SEP)

    # ── Resume ────────────────────────────────────────────────────────────────
    start_step = 0
    ckpt_path = out_dir / "latest.pt"
    if ckpt_path.exists():
        print(f"\n  Resuming from {ckpt_path}...")
        ckpt = torch.load(ckpt_path, map_location=device, weights_only=True)
        raw_model = model._orig_mod if hasattr(model, "_orig_mod") else model
        raw_model.load_state_dict(ckpt["model"])
        optimizer.load_state_dict(ckpt["optimizer"])
        start_step = ckpt["step"]
        print(f"  Resumed at step {start_step}")

    # ── Loop ──────────────────────────────────────────────────────────────────
    model.train()
    data_iter = iter(loader)

    def get_batch():
        nonlocal data_iter
        try:
            return next(data_iter)
        except StopIteration:
            data_iter = iter(loader)
            return next(data_iter)

    run_t0 = time.perf_counter()
    tokens_seen = start_step * tokens_per_step
    step = start_step
    eval_interval = cfg["eval_interval"]

    print(f"\n  Starting training (step {start_step} -> {total_steps})...")

    while step < total_steps:
        t0 = time.perf_counter()
        lr_now = cosine_lr(step, warmup_steps, total_steps, cfg["lr"], cfg["min_lr"])
        for pg in optimizer.param_groups:
            pg["lr"] = lr_now

        optimizer.zero_grad(set_to_none=True)
        total_loss = 0.0

        for _ in range(grad_accum):
            x, t = get_batch()
            x = x.to(device, non_blocking=True)
            t = t.to(device, non_blocking=True)
            with autocast(device_type=device.type, dtype=dtype, enabled=(device.type == "cuda")):
                _, loss = model(x, t, return_logits=False)
                loss = loss / grad_accum
            scaler.scale(loss).backward()
            total_loss += loss.item()

        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["max_norm"])
        scaler.step(optimizer)
        scaler.update()

        if device.type == "cuda":
            torch.cuda.synchronize()

        dt = time.perf_counter() - t0
        step += 1
        tokens_seen += tokens_per_step

        if step % cfg["log_interval"] == 0 or step <= 2:
            tps = tokens_per_step / dt
            steps_left = total_steps - step
            eta_h = steps_left * dt / 3600
            vram = torch.cuda.max_memory_allocated() / 1024**3 if device.type == "cuda" else 0
            print(f"  step {step:6d}/{total_steps} | loss {total_loss:.4f} | "
                  f"lr {lr_now:.2e} | {tps:,.0f} tok/s | VRAM {vram:.1f}GB | ETA {eta_h:.1f}h")

        # ── Validation generation ─────────────────────────────────────────────
        if step % eval_interval == 0:
            raw_model = model._orig_mod if hasattr(model, "_orig_mod") else model
            run_validation(raw_model, tokenizer, device, dtype, step,
                           seq_len=cfg["seq_len"], max_new=64)

        # ── Checkpoint save ───────────────────────────────────────────────────
        if step % cfg["save_interval"] == 0 or step == total_steps:
            raw_model = model._orig_mod if hasattr(model, "_orig_mod") else model
            step_dir = out_dir / f"step-{step:08d}"
            step_dir.mkdir(parents=True, exist_ok=True)
            torch.save(raw_model.state_dict(), step_dir / "lit_model.pth")
            with open(step_dir / "model_config.json", "w") as f:
                json.dump(model_config, f, indent=2)
            torch.save({"step": step, "model": raw_model.state_dict(),
                        "optimizer": optimizer.state_dict(),
                        "tokens_seen": tokens_seen},
                       out_dir / "latest.pt")
            print(f"  Saved -> {step_dir}")

    # ── Final ─────────────────────────────────────────────────────────────────
    final_dir = out_dir / "final"
    final_dir.mkdir(parents=True, exist_ok=True)
    raw_model = model._orig_mod if hasattr(model, "_orig_mod") else model
    torch.save(raw_model.state_dict(), final_dir / "lit_model.pth")
    with open(final_dir / "model_config.json", "w") as f:
        json.dump(model_config, f, indent=2)

    tok_src = Path(cfg["tokenizer_dir"])
    if tok_src.exists():
        shutil.copytree(tok_src, final_dir / "tokenizer", dirs_exist_ok=True)

    total_h = (time.perf_counter() - run_t0) / 3600
    print(SEP)
    print(f"  Training Done! {total_h:.2f} h -> {final_dir}")
    print(SEP)

    # ── Final validation ──────────────────────────────────────────────────────
    run_validation(raw_model, tokenizer, device, dtype, step,
                   seq_len=cfg["seq_len"], max_new=128)

    # ── Upload to HuggingFace ─────────────────────────────────────────────────
    if cfg.get("upload_checkpoints", True):
        upload_to_hf(out_dir, cfg)


# ─── Entry point ──────────────────────────────────────────────────────────────

def parse_args():
    p = argparse.ArgumentParser(description="LUNA 300M Trainer")
    p.add_argument("--config", type=str, default="train_config_300m.yaml")
    p.add_argument("--data_path",   type=str,   default=None)
    p.add_argument("--out_dir",     type=str,   default=None)
    p.add_argument("--init_model_path", type=str, default=None)
    p.add_argument("--max_tokens",  type=int,   default=None)
    p.add_argument("--micro_batch", type=int,   default=None)
    p.add_argument("--global_batch",type=int,   default=None)
    p.add_argument("--lr",          type=float, default=None)
    p.add_argument("--num_workers", type=int,   default=None)
    p.add_argument("--save_interval",type=int,  default=None)
    p.add_argument("--log_interval",type=int,   default=None)
    p.add_argument("--eval_interval",type=int,  default=None)
    p.add_argument("--hf_repo",    type=str,   default=None)
    p.add_argument("--auto_config", type=lambda x: x.lower() in ("1","true","yes"),
                   default=None)
    return p.parse_args()


if __name__ == "__main__":
    args = parse_args()
    cfg = load_config(args.config)
    cfg = apply_cli_overrides(cfg, args)
    hw = probe_hardware()
    cfg = resolve_auto(cfg, hw)
    train(cfg)