training_v2/train/utils.py

"""Training utilities: optimizer setup, LR schedule, checkpointing, cloud backup."""

import json
import math
import os
import threading
import time
from pathlib import Path

import torch


def cosine_with_warmup(step, warmup, total, max_lr, min_lr_ratio=0.1):
    if step < warmup:
        return max_lr * (step + 1) / warmup
    progress = (step - warmup) / max(1, total - warmup)
    progress = min(1.0, progress)
    return min_lr_ratio * max_lr + 0.5 * (max_lr - min_lr_ratio * max_lr) * (1 + math.cos(math.pi * progress))


def wsd_scheduler(step, warmup, stable, total, max_lr, min_lr_ratio=0.1):
    """Warmup-Stable-Decay (WSD) learning rate schedule.

    Allows pausing in the stable phase for evaluation or corpus extension.
    Decay only happens in the final fraction — resuming mid-stable is safe.
    """
    if step < warmup:
        return max_lr * (step + 1) / warmup
    if step < warmup + stable:
        return max_lr
    decay_step = step - warmup - stable
    decay_total = max(1, total - warmup - stable)
    progress = min(1.0, decay_step / decay_total)
    return min_lr_ratio * max_lr + 0.5 * (max_lr - min_lr_ratio * max_lr) * (1 + math.cos(math.pi * progress))


def make_optimizer(model, lr, weight_decay=0.1, betas=(0.9, 0.95), fused=True):
    """AdamW with weight decay only on 2D weights (no decay on biases / norms / embeddings).

    Per Loshchilov & Hutter; same convention as nanoGPT.
    """
    decay, no_decay = [], []
    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if p.dim() >= 2 and "tok_emb" not in n:
            decay.append(p)
        else:
            no_decay.append(p)
    groups = [
        {"params": decay, "weight_decay": weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]
    extra = {}
    if fused and torch.cuda.is_available():
        try:
            return torch.optim.AdamW(groups, lr=lr, betas=betas, fused=True)
        except TypeError:
            pass
    return torch.optim.AdamW(groups, lr=lr, betas=betas, **extra)


def save_checkpoint(path, model, optimizer, scheduler_state, step, extra=None):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict() if optimizer is not None else None,
        "scheduler": scheduler_state,
        "step": step,
        "config": {k: getattr(model.cfg, k) for k in model.cfg.__dataclass_fields__},
        "extra": extra or {},
    }
    tmp = path.with_suffix(path.suffix + ".tmp")
    torch.save(payload, tmp)
    os.replace(tmp, path)


def load_checkpoint(path, model, optimizer=None, map_location="cpu"):
    payload = torch.load(path, map_location=map_location, weights_only=False)
    state = payload["model"]
    if any(k.startswith("_orig_mod.") for k in state):
        state = {k.replace("_orig_mod.", "", 1): v for k, v in state.items()}
    model.load_state_dict(state)
    if optimizer is not None and payload.get("optimizer"):
        optimizer.load_state_dict(payload["optimizer"])
    return payload.get("step", 0), payload.get("extra", {})


def count_tokens(loader_output_iter, n_steps, block_size, batch_size):
    """Approximate; effective tokens consumed per step."""
    return n_steps * block_size * batch_size


def log_jsonl(path, record):
    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(record, ensure_ascii=False) + "\n")


def _save_weights_only(src_path, dst_path):
    """Write a weights-only copy of a full checkpoint (drops optimizer state).

    Returns dst_path. Used to produce the lighter artifact uploaded to HuggingFace.
    """
    payload = torch.load(src_path, map_location="cpu", weights_only=False)
    slim = {
        "model": payload["model"],
        "scheduler": payload.get("scheduler"),
        "step": payload.get("step"),
        "config": payload.get("config"),
        "extra": payload.get("extra", {}),
    }
    dst_path = Path(dst_path)
    dst_path.parent.mkdir(parents=True, exist_ok=True)
    tmp = dst_path.with_suffix(dst_path.suffix + ".tmp")
    torch.save(slim, tmp)
    os.replace(tmp, dst_path)
    return dst_path


def _parse_gcs_uri(uri):
    """Split gs://bucket/prefix into (bucket, prefix). Trailing slash stripped."""
    if not uri.startswith("gs://"):
        raise ValueError(f"GCS uri must start with gs://, got {uri!r}")
    rest = uri[len("gs://"):]
    parts = rest.split("/", 1)
    bucket = parts[0]
    prefix = parts[1].rstrip("/") if len(parts) > 1 else ""
    return bucket, prefix


class CloudBackup:
    """Non-blocking checkpoint backup to GCS (primary) and HuggingFace (secondary).

    Each destination runs a single-slot worker: at most one upload in flight per
    destination. If a new trigger arrives while an upload is running, it is dropped
    (logged) rather than queued, so a slow network can never stall training or build
    an unbounded backlog. All failures are logged as warnings; training never stops.

    GCS receives the full checkpoint (model + optimizer). HuggingFace receives a
    weights-only copy. Successful uploads are recorded in {out_dir}/backup_state.json.

    Layout:
      GCS:  {gcs_prefix}/phase{phase}/last.pt          (overwritten every trigger)
            {gcs_prefix}/phase{phase}/step_{N:07d}.pt  (every `interval` steps)
      HF:   model_step_{N:07d}.pt  (weights only, every `interval` steps)
    """

    def __init__(self, out_dir, phase, gcs_uri=None, hf_repo=None,
                 interval=2000, hf_token=None, log=print):
        self.out_dir = Path(out_dir)
        self.phase = phase
        self.interval = max(1, int(interval))
        self.log = log
        self.state_path = self.out_dir / "backup_state.json"

        self._gcs_bucket = None
        self._gcs = None
        self._gcs_prefix = None
        self._hf_api = None
        self._hf_repo = None

        self._state_lock = threading.Lock()
        self._slots = {}      # dest -> threading.Lock (held while uploading)
        self._state = self._load_state()

        if gcs_uri:
            self._init_gcs(gcs_uri)
        if hf_repo:
            self._init_hf(hf_repo, hf_token)

    @property
    def enabled(self):
        return self._gcs is not None or self._hf_api is not None

    # ── destination setup ──────────────────────────────────────────────────────
    def _init_gcs(self, gcs_uri):
        try:
            from google.cloud import storage
            bucket_name, prefix = _parse_gcs_uri(gcs_uri)
            client = storage.Client()
            bucket = client.bucket(bucket_name)
            bucket.reload()  # forces auth + existence check
            self._gcs = client
            self._gcs_bucket = bucket
            self._gcs_prefix = prefix
            self._slots["gcs"] = threading.Lock()
            self.log(f"[backup] GCS enabled → gs://{bucket_name}/{prefix}")
        except Exception as e:
            self._gcs = None
            self.log(f"[backup][warn] GCS disabled (init failed): {e}")

    def _init_hf(self, hf_repo, hf_token):
        try:
            from huggingface_hub import HfApi
            token = hf_token or self._resolve_hf_token()
            if not token:
                self.log("[backup][warn] HF disabled: no token (HF_TOKEN / ~/tok2.txt)")
                return
            api = HfApi(token=token)
            api.create_repo(repo_id=hf_repo, repo_type="model",
                            private=True, exist_ok=True)
            self._hf_api = api
            self._hf_repo = hf_repo
            self._slots["hf"] = threading.Lock()
            self.log(f"[backup] HF enabled → {hf_repo} (private)")
        except Exception as e:
            self._hf_api = None
            self.log(f"[backup][warn] HF disabled (init failed): {e}")

    @staticmethod
    def _resolve_hf_token():
        tok = os.environ.get("HF_TOKEN")
        if tok:
            return tok.strip()
        tok_file = Path.home() / "tok2.txt"
        if tok_file.exists():
            t = tok_file.read_text(encoding="utf-8").strip()
            return t or None
        return None

    # ── state file ─────────────────────────────────────────────────────────────
    def _load_state(self):
        if self.state_path.exists():
            try:
                return json.loads(self.state_path.read_text(encoding="utf-8"))
            except Exception:
                pass
        return {"phase": self.phase, "gcs": {}, "hf": {}}

    def _record(self, dest, step, uri):
        with self._state_lock:
            self._state.setdefault(dest, {})[str(step)] = {
                "uri": uri, "ts": time.time(),
            }
            tmp = self.state_path.with_suffix(".json.tmp")
            tmp.write_text(json.dumps(self._state, indent=2), encoding="utf-8")
            os.replace(tmp, self.state_path)

    # ── public trigger ─────────────────────────────────────────────────────────
    def backup(self, ckpt_path, step, is_final=False):
        """Trigger backups for `ckpt_path` at `step`. Returns immediately.

        Uploads `step_{N}.pt` / `model_step_{N}.pt` only on interval boundaries
        (or when is_final). `last.pt` on GCS is refreshed on every trigger.
        """
        if not self.enabled:
            return
        ckpt_path = str(ckpt_path)
        keep_versioned = is_final or (step % self.interval == 0)

        if self._gcs is not None:
            self._dispatch("gcs", self._do_gcs, ckpt_path, step, keep_versioned)
        if self._hf_api is not None and keep_versioned:
            self._dispatch("hf", self._do_hf, ckpt_path, step, keep_versioned)

    def _dispatch(self, dest, fn, ckpt_path, step, keep_versioned):
        lock = self._slots[dest]
        if not lock.acquire(blocking=False):
            self.log(f"[backup][{dest}] busy, skipping step {step}")
            return

        def runner():
            try:
                fn(ckpt_path, step, keep_versioned)
            except Exception as e:
                self.log(f"[backup][{dest}][warn] step {step} failed: {e}")
            finally:
                lock.release()

        threading.Thread(target=runner, name=f"backup-{dest}", daemon=True).start()

    # ── workers ────────────────────────────────────────────────────────────────
    def _do_gcs(self, ckpt_path, step, keep_versioned):
        base = f"{self._gcs_prefix + '/' if self._gcs_prefix else ''}phase{self.phase}"
        last_blob = f"{base}/last.pt"
        self._gcs_bucket.blob(last_blob).upload_from_filename(
            ckpt_path, timeout=1800)
        if keep_versioned:
            ver_blob = f"{base}/step_{step:07d}.pt"
            self._gcs_bucket.blob(ver_blob).upload_from_filename(
                ckpt_path, timeout=1800)
            uri = f"gs://{self._gcs_bucket.name}/{ver_blob}"
        else:
            uri = f"gs://{self._gcs_bucket.name}/{last_blob}"
        self._record("gcs", step, uri)
        self.log(f"[backup][gcs] step {step} → {uri}")

    def _do_hf(self, ckpt_path, step, keep_versioned):
        slim_path = self.out_dir / f".hf_upload_step_{step:07d}.pt"
        try:
            _save_weights_only(ckpt_path, slim_path)
            dst = f"phase{self.phase}/model_step_{step:07d}.pt"
            self._hf_api.upload_file(
                path_or_fileobj=str(slim_path),
                path_in_repo=dst,
                repo_id=self._hf_repo,
                repo_type="model",
                commit_message=f"backup phase{self.phase} step {step}",
            )
            uri = f"hf://{self._hf_repo}/{dst}"
            self._record("hf", step, uri)
            self.log(f"[backup][hf] step {step} → {uri}")
        finally:
            try:
                slim_path.unlink(missing_ok=True)
            except OSError:
                pass

    def wait(self, timeout=None):
        """Block until in-flight uploads on every destination finish (best effort)."""
        for dest, lock in self._slots.items():
            acquired = lock.acquire(timeout=timeout if timeout is not None else -1)
            if acquired:
                lock.release()
            else:
                self.log(f"[backup][{dest}][warn] still uploading at shutdown")