train.py

#!/usr/bin/env python3
"""
train.py: masked-diffusion chat SFT for MetaDiffusion-600M.

Usage:
    python train.py --init-checkpoint init/metadiffusion-600M-instruct.pt \
        --data-dir data --output-dir checkpoints --max-steps 30000
    python train.py ... --keep-knowledge --keep-mult 0.05
"""

import argparse
import datetime
import glob
import heapq
import json
import logging
import math
import os
import re
import signal
import sys
import time
from dataclasses import asdict

import numpy as np
import torch
from torch.optim.lr_scheduler import LambdaLR

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)


def format_bytes(b):
    for unit in ["B", "KB", "MB", "GB", "TB"]:
        if b < 1024:
            return f"{b:.1f} {unit}"
        b /= 1024
    return f"{b:.1f} PB"


def format_eta(seconds):
    seconds = max(0, int(seconds))
    h, rem = divmod(seconds, 3600)
    m, s = divmod(rem, 60)
    if h > 0:
        return f"{h}h{m:02d}m"
    if m > 0:
        return f"{m}m{s:02d}s"
    return f"{s}s"


def detect_max_batch_size(model, seq_len, device, keep_free_fraction=0.1,
                          vocab=151677, reserve_bytes=0):
    """Forward+backward probe to find the largest batch that fits.

    reserve_bytes: memory the optimizer states need during real training
    (2 x n_params x 4 for fp32 Adam, x 1 for 8-bit Adam). Without this the
    probe over-predicts and the first real step OOMs."""
    if not torch.cuda.is_available():
        return 8
    # mem_get_info() returns (free, total). Base the budget on ACTUAL free
    # memory: other processes may hold part of the card.
    free_mem, total_mem = torch.cuda.mem_get_info()
    logger.info(f"GPU: {format_bytes(total_mem)} total, {format_bytes(free_mem)} free")
    mem_limit = free_mem - int(total_mem * keep_free_fraction) - reserve_bytes
    logger.info(f"VRAM budget: {format_bytes(mem_limit)} "
                f"(based on {format_bytes(free_mem)} free, "
                f"reserving {format_bytes(reserve_bytes)} for optimizer + accum grads)")
    model = model.to(device)
    model.train()

    last_working = 1
    for bs in [1, 2, 4, 8, 16, 24, 32, 48]:
        torch.cuda.empty_cache()
        try:
            ids = torch.randint(0, vocab, (bs, seq_len), device=device)
            labels = ids.clone()
            mask = torch.rand(bs, seq_len, device=device) < 0.5
            t = torch.rand(bs, device=device)
            logits = model(ids, t)
            loss, n = model.compute_loss(logits, labels, mask)
            if n > 0:
                (loss / 4).backward()
            torch.cuda.synchronize()
            peak = torch.cuda.max_memory_allocated()
            logger.info(f"  batch={bs:>2d}: peak {format_bytes(peak)} "
                        f"({peak/total_mem*100:.0f}%)")
            if peak >= mem_limit:
                break
            last_working = bs
        except RuntimeError as e:
            if "out of memory" in str(e).lower():
                logger.info(f"  batch={bs:>2d}: OOM")
                break
            raise
        finally:
            model.zero_grad(set_to_none=True)
            torch.cuda.empty_cache()
    logger.info(f"Detected max batch_size={last_working}")
    return last_working


def get_step_from_filename(filename):
    basename = os.path.basename(filename)
    m = re.match(r"step_(\d+)(?:_\w+)?\.pt$", basename)
    return int(m.group(1)) if m else None

def load_best_val_steps(stats_path, max_n):
    """Return the steps with the N lowest val losses from stats.jsonl.

    Falls back to train loss when val never ran (patience=0)."""
    if not os.path.exists(stats_path):
        return set()
    val_entries, loss_entries = [], []
    with open(stats_path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                e = json.loads(line)
            except json.JSONDecodeError:
                continue
            if "step" not in e:
                continue
            if e.get("val_loss") is not None:
                val_entries.append((float(e["val_loss"]), int(e["step"])))
            elif "loss" in e:
                loss_entries.append((float(e["loss"]), int(e["step"])))
    pool = val_entries if val_entries else loss_entries
    return {s for _, s in heapq.nsmallest(max_n, pool)}


def cleanup_checkpoints(output_dir, keep_last_n, keep_best_n, stats_path):
    """Keep the N latest regular checkpoints + the N best-by-val-loss
    checkpoints; delete everything else to save storage."""
    # sort by step NUMBER, not filename: lexicographic order breaks once
    # steps hit 5 digits (step_30000.pt < step_4000.pt alphabetically)
    all_ckpts = sorted(glob.glob(os.path.join(output_dir, "step_*.pt")),
                       key=lambda c: (get_step_from_filename(c) or -1, c))
    if len(all_ckpts) <= keep_last_n + keep_best_n:
        return

    last_steps = {get_step_from_filename(c) for c in all_ckpts[-keep_last_n:]}
    best_steps = load_best_val_steps(stats_path, keep_best_n)

    kept = set()
    # latest N: regular files only (no _valbest suffix)
    for s in last_steps:
        for c in all_ckpts:
            if "valbest" not in c and get_step_from_filename(c) == s:
                kept.add(c)
    # best N by val: prefer the regular file at that step, else its valbest
    for s in best_steps:
        cands = [c for c in all_ckpts if get_step_from_filename(c) == s]
        if not cands:
            continue
        regular = [c for c in cands if "valbest" not in c]
        kept.add(regular[0] if regular else cands[0])

    deleted = 0
    for c in all_ckpts:
        if c not in kept:
            try:
                os.remove(c)
                deleted += 1
            except OSError:
                pass
    if deleted:
        logger.info(f"Cleaned {deleted} checkpoints "
                    f"(kept {len(kept)}: {keep_last_n} latest + {keep_best_n} best-val)")


def cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps, min_lr_ratio=0.1):
    def lr_lambda(step):
        if step < warmup_steps:
            return float(step) / float(max(1, warmup_steps))
        progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps))
        return max(min_lr_ratio, 0.5 * (1.0 + math.cos(math.pi * progress)))
    return LambdaLR(optimizer, lr_lambda)


def last_eos_mask(clean, eos_token_id, resp_bool):
    """True only at the rightmost <|im_end|> in the response region per row.

    Historical multi-turn terminators stay unforced; the final terminator
    still gets a gradient every step under --eos-mask-always."""
    is_eos = (clean == eos_token_id) & resp_bool
    rev = torch.flip(is_eos, [1]).int()
    last = torch.flip(rev.cumsum(1) == 1, [1]) & is_eos
    return last


def no_decay_name(name, param):
    return param.ndim < 2 or "norm" in name or name.endswith(".bias")


def main():
    p = argparse.ArgumentParser(description="MetaDiffusion-600M chat SFT")
    p.add_argument("--init-checkpoint", required=True)
    p.add_argument("--data-dir", default="data")
    p.add_argument("--output-dir", default="checkpoints")
    p.add_argument("--max-steps", type=int, default=30000)
    p.add_argument("--seq-len", type=int, default=512)
    p.add_argument("--batch-size", type=int, default=0, help="0 = auto-detect")
    p.add_argument("--grad-accum-steps", type=int, default=2)
    p.add_argument("--lr", type=float, default=5e-5)
    p.add_argument("--min-lr-ratio", type=float, default=0.1)
    p.add_argument("--warmup-steps", type=int, default=200)
    p.add_argument("--weight-decay", type=float, default=0.01)
    p.add_argument("--transferred-lr-mult", type=float, default=0.33)
    p.add_argument("--max-grad-norm", type=float, default=1.0)
    p.add_argument("--optimizer", default="auto",
                   choices=["auto", "adamw", "adamw8bit"],
                   help="auto = torchao/bitsandbytes 8-bit Adam if installed, "
                        "else AdamW. adamw8bit forces 8-bit (warns if missing).")
    p.add_argument("--dtype", default="bfloat16", choices=["float32", "bfloat16", "float16"])
    p.add_argument("--mask-all", action="store_true", help="mask the whole sequence, not just responses")
    p.add_argument("--mask-ratio-min", type=float, default=0.0)
    p.add_argument("--mask-ratio-max", type=float, default=1.0)
    p.add_argument("--curriculum", action="store_true",
                   help="Ramp the mask ratio from --mask-ratio-min to 1.0 over the "
                        "run. Val stays at fixed --val-t.")
    p.add_argument("--curriculum-early-stop-gate",
                   action=argparse.BooleanOptionalAction, default=True,
                   help="While the curriculum ramp has not yet covered --val-t, "
                        "val is out-of-domain (extrapolation): compute + log it, "
                        "but do NOT update best-val / patience / valbest. Without "
                        "this, early stopping can fire mid-ramp on OOD noise and "
                        "kill the run before the model learns high-mask denoising. "
                        "--no-curriculum-early-stop-gate restores the old behavior.")
    p.add_argument("--eos-weight", type=float, default=1.0,
                   help="Loss multiplier on the <|im_end|> token when it is a "
                        "masked target (teach termination; try 10, hammer 25)")
    p.add_argument("--eos-token-id", type=int, default=151645,
                   help="<|im_end|> token id (Qwen3: 151645)")
    p.add_argument("--eos-mask-always", action="store_true",
                   help="Force-mask the last <|im_end|> terminator in every window, "
                        "regardless of t: the terminator gets a gradient on every "
                        "step and the curriculum can never starve it")
    p.add_argument("--keep-knowledge", action="store_true", help="RND1-style LR split")
    p.add_argument("--keep-mult", type=float, default=0.05, help="LR mult for MLP/norm/embed when --keep-knowledge")
    p.add_argument("--save-every", type=int, default=2000)
    p.add_argument("--log-every", type=int, default=50)
    p.add_argument("--val-every", type=int, default=500,
                   help="Val loss check interval (needs ids_val.bin/resp_val.bin)")
    p.add_argument("--patience", type=int, default=0,
                   help="Early stop after N val checks without improvement (0 = off)")
    p.add_argument("--min-delta", type=float, default=1e-4,
                   help="Min val loss improvement to count as improvement")
    p.add_argument("--val-batches", type=int, default=50,
                   help="Batches averaged per val check")
    p.add_argument("--val-t", type=float, default=0.5,
                   help="Fixed mask ratio for val checks (stable early stopping; "
                        "t~U(0,1) makes val swing +-0.5 and patience fires on noise)")
    p.add_argument("--keep-last-n", type=int, default=3,
                   help="Keep the N latest regular checkpoints")
    p.add_argument("--keep-best-n", type=int, default=3,
                   help="Keep the N lowest-val-loss checkpoints (train-loss fallback)")
    p.add_argument("--resume-from", default=None)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--device", default="cuda:0")
    args = p.parse_args()

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[args.dtype]

    # --- Model ---
    from model import MetaDiffusionConfig, MetaDiffusionLM
    logger.info(f"Loading init checkpoint: {args.init_checkpoint}")
    ckpt = torch.load(args.init_checkpoint, map_location="cpu", weights_only=False)
    config = MetaDiffusionConfig(
        **{k: v for k, v in ckpt["config"].items() if k in MetaDiffusionConfig.__dataclass_fields__}
    )
    model = MetaDiffusionLM(config)
    model.load_state_dict(ckpt["model_state_dict"], strict=True)
    model = model.to(device=device, dtype=dtype)
    logger.info(f"Model: {config.num_hidden_layers}L x {config.hidden_size}W, "
                f"vocab={config.mask_vocab_size}, params={sum(p.numel() for p in model.parameters())/1e6:.1f}M, {args.dtype}")

    # --- Data ---
    data_dir = args.data_dir
    ids_arr = np.memmap(os.path.join(data_dir, "ids.bin"), dtype=np.uint32, mode="r")
    resp_arr = np.memmap(os.path.join(data_dir, "resp.bin"), dtype=np.uint8, mode="r")
    assert ids_arr.shape == resp_arr.shape, "ids.bin / resp.bin length mismatch"
    n_tokens = len(ids_arr)
    logger.info(f"Data: {n_tokens:,} tokens, seq_len={args.seq_len}")

    meta_path = os.path.join(data_dir, "meta.json")
    data_meta = {}
    if os.path.exists(meta_path):
        with open(meta_path) as f:
            data_meta = json.load(f)
    if not data_meta.get("val_held_out") and data_meta.get("n_val_samples", 0) > 0:
        skip = int(data_meta["n_val_samples"]) * int(data_meta.get("seq_len", args.seq_len))
        if skip > 0 and skip < n_tokens:
            ids_arr = ids_arr[skip:]
            resp_arr = resp_arr[skip:]
            n_tokens = len(ids_arr)
            logger.info(f"Legacy leaked val: skipped first {skip:,} tokens of "
                        f"ids.bin ({data_meta['n_val_samples']} samples). "
                        f"Train tokens now {n_tokens:,}")

    # --- Validation split (for early stopping) ---
    val_ids_arr = val_resp_arr = None
    val_n = 0
    if args.patience > 0:
        v_ids = os.path.join(data_dir, "ids_val.bin")
        v_resp = os.path.join(data_dir, "resp_val.bin")
        if os.path.exists(v_ids) and os.path.exists(v_resp):
            val_ids_arr = np.memmap(v_ids, dtype=np.uint32, mode="r")
            val_resp_arr = np.memmap(v_resp, dtype=np.uint8, mode="r")
            val_n = len(val_ids_arr)
            logger.info(f"Val data: {val_n:,} tokens (early stopping active, "
                        f"patience={args.patience}, val_every={args.val_every})")
        else:
            logger.warning("--patience set but no ids_val.bin/resp_val.bin found; "
                           "re-run prepare_data.py (writes a val split by default). "
                           "Early stopping disabled.")
            args.patience = 0

    def sample_windows(batch_size, ids_src, resp_src, n_src, rng=None):
        n_windows = n_src // args.seq_len
        if n_windows < 1:
            raise RuntimeError(
                f"Need at least {args.seq_len} tokens, have {n_src}")
        if rng is None:
            idx = np.random.randint(0, n_windows, size=batch_size)
        else:
            idx = rng.randint(0, n_windows, size=batch_size)
        starts = idx * args.seq_len
        ids = np.stack([ids_src[s:s + args.seq_len] for s in starts])
        resp = np.stack([resp_src[s:s + args.seq_len] for s in starts])
        return torch.from_numpy(ids).long(), torch.from_numpy(resp).bool()

    def sample_batch(batch_size, ids_arr_local, resp_arr_local, n_tokens_local,
                     max_ratio=None, rng=None):
        """Aligned windows; mask within the response region."""
        clean, resp_bool = sample_windows(
            batch_size, ids_arr_local, resp_arr_local, n_tokens_local, rng=rng)
        max_r = args.mask_ratio_max if max_ratio is None else max_ratio
        if rng is None:
            t = torch.rand(batch_size) * (max_r - args.mask_ratio_min) + args.mask_ratio_min
            mask_u = torch.rand(batch_size, args.seq_len)
        else:
            t = torch.from_numpy(
                rng.rand(batch_size) * (max_r - args.mask_ratio_min) + args.mask_ratio_min
            ).float()
            mask_u = torch.from_numpy(rng.rand(batch_size, args.seq_len)).float()
        mask_positions = mask_u < t[:, None]
        if not args.mask_all:
            mask_positions = mask_positions & resp_bool
        if args.eos_mask_always:
            mask_positions = mask_positions | last_eos_mask(
                clean, args.eos_token_id, resp_bool)
        input_ids = clean.clone()
        input_ids[mask_positions] = config.mask_token_id
        return input_ids, clean, mask_positions, t

    # --- Optimizer selection (before batch detection: it reserves memory) ---
    OptimizerCls = None
    if args.optimizer in ("auto", "adamw8bit"):
        for mod_name, cls_name, label in (
                ("torchao.optim", "AdamW8bit", "torchao"),
                ("bitsandbytes.optim", "AdamW8bit", "bitsandbytes")):
            try:
                mod = __import__(mod_name, fromlist=[cls_name])
                OptimizerCls = getattr(mod, cls_name)
                logger.info(f"Optimizer: 8-bit Adam ({label})")
                break
            except ImportError:
                continue
    if OptimizerCls is None:
        from torch.optim import AdamW as OptimizerCls
        if args.optimizer == "adamw8bit":
            logger.warning("adamw8bit requested but neither torchao nor "
                           "bitsandbytes is installed; using AdamW "
                           "(pip install torchao)")

    # --- Auto batch ---
    if args.batch_size <= 0 and torch.cuda.is_available():
        state_bytes = 1 if OptimizerCls.__name__ == "AdamW8bit" else 4
        n_params = sum(p.numel() for p in model.parameters())
        reserve_bytes = 2 * n_params * state_bytes  # optimizer states
        # with grad accumulation, the (accum-1) earlier micro-batch grads are
        # still alive when the last micro-batch's backward peaks
        reserve_bytes += (args.grad_accum_steps - 1) * n_params * 2  # bf16 grads
        logger.info(f"Optimizer/grad reserve: {format_bytes(reserve_bytes)} "
                    f"(states {state_bytes} B/param + "
                    f"{args.grad_accum_steps - 1} extra bf16 grads)")
        logger.info("Auto-detecting max batch size...")
        args.batch_size = detect_max_batch_size(
            model, args.seq_len, device,
            keep_free_fraction=0.1, vocab=config.mask_vocab_size,
            reserve_bytes=reserve_bytes)
    if args.batch_size <= 0:
        args.batch_size = 8
    logger.info(f"batch_size={args.batch_size} x grad_accum={args.grad_accum_steps} "
                f"= effective {args.batch_size * args.grad_accum_steps}")

    # --- Parameter groups ---
    # Only the diffusion-new modules get full LR. embed_tokens / lm_head are
    # almost entirely transferred AR rows (MASK + rainbow are 8 of 151677).
    new_keys = ("timestep_emb", "timestep_modulation")
    attn_key = "self_attn"
    buckets = {("new", True): [], ("new", False): [],
               ("attn", True): [], ("attn", False): [],
               ("non_attn", True): [], ("non_attn", False): []}
    for name, param in model.named_parameters():
        if any(k in name for k in new_keys):
            kind = "new"
        elif attn_key in name:
            kind = "attn"
        else:
            kind = "non_attn"
        buckets[(kind, not no_decay_name(name, param))].append(param)

    def make_group(kind, decay, lr, label):
        params = buckets[(kind, decay)]
        if not params:
            return None
        return {"params": params, "lr": lr,
                "weight_decay": args.weight_decay if decay else 0.0,
                "name": label}

    if args.keep_knowledge:
        raw = [
            make_group("new", True, args.lr, "new"),
            make_group("new", False, args.lr, "new_nodecay"),
            make_group("attn", True, args.lr, "attention"),
            make_group("attn", False, args.lr, "attention_nodecay"),
            make_group("non_attn", True, args.lr * args.keep_mult, "mlp_norm_embed"),
            make_group("non_attn", False, args.lr * args.keep_mult, "mlp_norm_embed_nodecay"),
        ]
    else:
        tr_lr = args.lr * args.transferred_lr_mult
        raw = [
            make_group("new", True, args.lr, "new"),
            make_group("new", False, args.lr, "new_nodecay"),
            make_group("attn", True, tr_lr, "transferred"),
            make_group("attn", False, tr_lr, "transferred_nodecay"),
            make_group("non_attn", True, tr_lr, "transferred"),
            make_group("non_attn", False, tr_lr, "transferred_nodecay"),
        ]
    param_groups = [g for g in raw if g is not None]
    for g in param_groups:
        n = sum(p.numel() for p in g["params"])
        logger.info(f"  group {g['name']}: {n:,} params, lr={g['lr']:.2e}, "
                    f"wd={g['weight_decay']}")

    optimizer = OptimizerCls(param_groups)
    scheduler = cosine_schedule_with_warmup(optimizer, args.warmup_steps, args.max_steps, args.min_lr_ratio)

    start_step = 0
    if args.resume_from:
        ckpt_r = torch.load(args.resume_from, map_location="cpu", weights_only=False)
        model.load_state_dict(ckpt_r["model_state_dict"])
        start_step = int(ckpt_r.get("step", get_step_from_filename(args.resume_from) or 0))
        if "optimizer_state" in ckpt_r:
            try:
                optimizer.load_state_dict(ckpt_r["optimizer_state"])
                scheduler.load_state_dict(ckpt_r.get("scheduler_state", {}))
                logger.info("Restored optimizer + scheduler state (seamless continuation)")
            except Exception as e:
                logger.warning(f"Could not restore optimizer state ({e}); "
                               "rebuilding fresh. (Changed hyperparams / param group "
                               "layout between segments?)")
                optimizer = OptimizerCls(param_groups)
                scheduler = cosine_schedule_with_warmup(
                    optimizer, args.warmup_steps, args.max_steps, args.min_lr_ratio)
        else:
            logger.warning("Checkpoint has no optimizer state (older run); "
                           "rebuilding fresh.")
        logger.info(f"Resumed from {args.resume_from} (continuing from step {start_step})")

    os.makedirs(args.output_dir, exist_ok=True)
    stats_path = os.path.join(args.output_dir, "stats.jsonl")
    stats_file = open(stats_path, "a")

    # --- Train loop ---
    @torch.no_grad()
    def compute_val_loss():
        """Masked CE over --val-batches of the held-out split.

        Fixed mask ratio (--val-t) + deterministic aligned windows AND
        seeded masks: torch.rand made the estimate depend on the training
        RNG and patience fired on noise."""
        model.eval()
        total_ce, total_n = 0.0, 0
        rng = np.random.RandomState(args.seed)
        for _ in range(args.val_batches):
            clean, resp_bool = sample_windows(
                args.batch_size, val_ids_arr, val_resp_arr, val_n, rng=rng)
            t = torch.full((args.batch_size,), args.val_t)
            mask_u = torch.from_numpy(rng.rand(args.batch_size, args.seq_len)).float()
            mask_positions = mask_u < t[:, None]
            if not args.mask_all:
                mask_positions = mask_positions & resp_bool
            if args.eos_mask_always:
                mask_positions = mask_positions | last_eos_mask(
                    clean, args.eos_token_id, resp_bool)
            input_ids = clean.clone()
            input_ids[mask_positions] = config.mask_token_id
            input_ids = input_ids.to(device)
            clean = clean.to(device)
            mask_positions = mask_positions.to(device)
            t = t.to(device)
            logits = model(input_ids, t)
            ce, n = model.compute_loss(logits, clean, mask_positions,
                                       pad_token_id=config.pad_token_id,
                                       eos_token_id=args.eos_token_id,
                                       eos_weight=args.eos_weight)
            total_ce += ce.item() * n
            total_n += n
        model.train()
        return total_ce / max(1, total_n)

    def save_ckpt(step, path):
        torch.save({
            "config": asdict(config),
            "model_state_dict": model.state_dict(),
            "optimizer_state": optimizer.state_dict(),
            "scheduler_state": scheduler.state_dict(),
            "step": step,
            "metadata": {
                "init_checkpoint": args.init_checkpoint,
                "keep_knowledge": args.keep_knowledge,
                "dtype": args.dtype,
            },
        }, path)
        logger.info(f"Saved {path}")

    model.train()
    try:
        scaler = torch.amp.GradScaler(
            "cuda", enabled=(args.dtype == "float16" and torch.cuda.is_available()))
    except AttributeError:  # torch < 2.3
        scaler = torch.cuda.amp.GradScaler(
            enabled=(args.dtype == "float16" and torch.cuda.is_available()))
    eff_bs = args.batch_size * args.grad_accum_steps
    best_val, best_step, no_improve = float("inf"), 0, 0

    # Ctrl+C handling: first press finishes the current step and saves it,
    # second press force-quits without saving.
    interrupted = False

    def _handle_sigint(sig, frame):
        nonlocal interrupted
        if interrupted:
            print("\nSecond Ctrl+C: force quitting without save.", flush=True)
            os._exit(130)
        interrupted = True
        print("\nCtrl+C received: finishing current step, then saving...",
              flush=True)

    signal.signal(signal.SIGINT, _handle_sigint)

    if interrupted:
        logger.info("Ctrl+C before training started; nothing to save.")
        stats_file.close()
        sys.exit(130)

    if args.curriculum:
        ramp_floor = max(args.mask_ratio_min, 0.1)
        logger.info(f"Curriculum: mask ratio ramps {ramp_floor:.2f} -> 1.00 "
                    f"over the run (val stays at fixed t={args.val_t})")
        if args.curriculum_early_stop_gate and args.patience > 0:
            logger.info(f"Curriculum early-stop gate ON: val at t={args.val_t} is "
                        f"OOD until the ramp covers it; early stopping / valbest "
                        f"inactive until then (--no-curriculum-early-stop-gate to "
                        f"disable)")
    if args.eos_weight != 1.0:
        logger.info(f"EOS weighting: <|im_end|> (id {args.eos_token_id}) "
                    f"loss x{args.eos_weight}")
    if args.eos_mask_always:
        logger.info("EOS mask-always: last <|im_end|> in each window is a "
                    "masked target (curriculum cannot starve the terminator)")
    logger.info(f"Training {args.max_steps:,} steps... (Ctrl+C saves current step)")
    t0 = time.time()
    step = start_step
    for step in range(start_step + 1, args.max_steps + 1):
        optimizer.zero_grad(set_to_none=True)
        total_ce = 0.0
        total_n = 0
        val_loss = None
        cur_max_ratio = None
        if args.curriculum:
            # easy-to-hard: mask ratio ramps from a 0.1 floor to 1.0 across the
            # run. Starting at 0.0 leaves whole micro-batches unmasked (no grad
            # path); the floor keeps t ~ U(0, 0.1) at step 1 (~100 masked
            # tokens per 2048-token batch). Segment-relative: any resume
            # segment covers the full ramp, so a shortened segment still trains
            # the high-mask regime (where im_end lives) end to end.
            seg_steps = max(1, args.max_steps - start_step)
            progress = (step - start_step) / seg_steps
            ramp_floor = max(args.mask_ratio_min, 0.1)
            cur_max_ratio = ramp_floor + (1.0 - ramp_floor) * progress
        did_backward = False
        for _ in range(args.grad_accum_steps):
            input_ids, clean, mask_positions, t = sample_batch(
                args.batch_size, ids_arr, resp_arr, n_tokens,
                max_ratio=cur_max_ratio)
            input_ids = input_ids.to(device)
            clean = clean.to(device)
            mask_positions = mask_positions.to(device)
            t = t.to(device)

            logits = model(input_ids, t)
            ce, n = model.compute_loss(logits, clean, mask_positions, pad_token_id=config.pad_token_id,
                                       eos_token_id=args.eos_token_id,
                                       eos_weight=args.eos_weight)
            if n == 0:
                continue  # no masked tokens this micro-batch (early curriculum);
                          # compute_loss returns a constant, not a grad path
            loss = ce / args.grad_accum_steps
            total_ce += ce.item() * n  # weighted by masked-token count: true per-token CE
            total_n += n
            scaler.scale(loss).backward()
            did_backward = True

        if did_backward:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()

        if interrupted:
            save_ckpt(step, os.path.join(args.output_dir, f"step_{step}.pt"))
            logger.info(f"Ctrl+C: saved current step -> step_{step}.pt")
            break

        if step % args.log_every == 0:
            elapsed = time.time() - t0
            # segment-relative step count: cumulative `step` over segment
            # elapsed would print a fake 100K+ tok/s after --resume-from
            tok_per_s = eff_bs * args.seq_len * (step - start_step) / max(1.0, elapsed)
            vram = torch.cuda.max_memory_allocated() / 1e9 if torch.cuda.is_available() else 0
            steps_done = max(1, step - start_step)
            eta_s = (args.max_steps - step) * (elapsed / steps_done)
            finish = datetime.datetime.now() + datetime.timedelta(seconds=eta_s)
            eta_str = (f"ETA {format_eta(eta_s)} (el {format_eta(elapsed)}, "
                       f"done {finish.strftime('%I:%M %p').lstrip('0')})")
            val_loss = None
            if args.patience > 0 and step % args.val_every == 0:
                val_loss = compute_val_loss()
                logger.info(f"step {step:>6d}/{args.max_steps} | loss={total_ce/max(1,total_n):.4f}"
                            f" | val={val_loss:.4f} | lr={scheduler.get_last_lr()[0]:.2e} | "
                            f"{tok_per_s:.0f} tok/s | vram={vram:.1f}GB | {eta_str}")
            else:
                logger.info(f"step {step:>6d}/{args.max_steps} | loss={total_ce/max(1,total_n):.4f}"
                            f" | lr={scheduler.get_last_lr()[0]:.2e} | "
                            f"{tok_per_s:.0f} tok/s | vram={vram:.1f}GB | {eta_str}")
            stats_file.write(json.dumps({
                "step": step, "loss": total_ce / max(1, total_n),
                "lr": float(scheduler.get_last_lr()[0]), "tok_per_s": tok_per_s,
                "vram_gb": vram, "val_loss": val_loss,
            }) + "\n")
            stats_file.flush()

        if step % args.save_every == 0:
            save_ckpt(step, os.path.join(args.output_dir, f"step_{step}.pt"))
            cleanup_checkpoints(args.output_dir, args.keep_last_n,
                                args.keep_best_n, stats_path)

        if args.patience > 0 and step % args.val_every == 0:
            if val_loss is None:
                val_loss = compute_val_loss()
            # Curriculum gate: before the ramp covers --val-t the val estimate
            # is extrapolation, not generalization. Keep the number for the
            # curve but do not let it move best_val / patience / valbest.
            gated = (args.curriculum_early_stop_gate and args.curriculum
                     and cur_max_ratio is not None and cur_max_ratio < args.val_t)
            if gated:
                logger.info(f"  [val] {val_loss:.4f} gated (t_max {cur_max_ratio:.3f} "
                            f"< val-t {args.val_t}); early stopping inactive")
            elif val_loss < best_val - args.min_delta:
                best_val, best_step, no_improve = val_loss, step, 0
                save_ckpt(step, os.path.join(args.output_dir, f"step_{step}_valbest.pt"))
                logger.info(f"  [val] new best {val_loss:.4f} -> step_{step}_valbest.pt")
            else:
                no_improve += 1
                logger.info(f"  [val] no improvement ({no_improve}/{args.patience}), "
                            f"best={best_val:.4f} @ step {best_step}")
                if no_improve >= args.patience:
                    logger.info(f"Early stop at step {step}; best val {best_val:.4f} "
                                f"at step {best_step} (step_{best_step}_valbest.pt)")
                    save_ckpt(step, os.path.join(args.output_dir, f"step_{step}.pt"))
                    break

    stats_file.close()
    if interrupted:
        logger.info(f"Stopped early (Ctrl+C) at step {step}. "
                    f"Resume with --resume-from {args.output_dir}/step_{step}.pt")
    else:
        logger.info(f"Done. Final checkpoint: {args.output_dir}/step_{step}.pt")


if __name__ == "__main__":
    main()