insc/train/finetune.py

"""Fine-tune the Stoicheia torso on I.PHI inscriptions (diffusion objective, Ithaca-
oriented masking: contiguous spans of 1-10 chars dominate). Config-driven, DDP/torchrun.

Reuses Stoicheia modules wholesale (PYTHONPATH=$STOICHEIA_ROOT); this file only changes:
  - init from $INS_TORSO (weights only; fresh optimizer, fresh short schedule)
  - data mix: iphi shards + gold/silver replay (anti-forgetting)
  - noise mix: span-heavy, span lengths matched to the Ithaca eval (1-10)
  - quick PHI-val greedy-restore CER as the per-checkpoint eval

  torchrun --nproc_per_node=4 insc_train/finetune.py --config configs/finetune.json
"""
from __future__ import annotations

import argparse, json, math, os, sys, time
from pathlib import Path

import torch

from model.char_bert import CharBertConfig, CharBertEncoder, num_params
from train.collate import pack_batch
from train.data import DataConfig, MultiTierLoader, TierSpec
from train.loss import compute_loss
from train.noising import NoiseConfig
from train.train import ddp_setup, infinite_records, save_ckpt

INS_ROOT = Path(__file__).resolve().parents[1]


class MixDataset(torch.utils.data.IterableDataset):
    """Packed batches from an explicit DataConfig; each (rank, worker) gets a disjoint
    shard and its own RNG (same contract as Stoicheia's BatchDataset)."""
    def __init__(self, dcfg, ncfg, T, rows, seed, rank, world):
        super().__init__()
        self.dcfg, self.ncfg, self.T, self.rows = dcfg, ncfg, T, rows
        self.seed, self.rank, self.world = seed, rank, world

    def __iter__(self):
        import copy
        info = torch.utils.data.get_worker_info()
        wid = info.id if info else 0
        nw = info.num_workers if info else 1
        gshard = self.rank * nw + wid
        gtot = self.world * nw
        c = copy.deepcopy(self.dcfg)
        loader = MultiTierLoader(c, rank=gshard, world_size=gtot)
        g = torch.Generator().manual_seed(self.seed * 100003 + gshard)
        gen = infinite_records(loader)
        while True:
            yield pack_batch(gen, self.ncfg, self.T, self.rows, g)


def make_loader(dcfg, ncfg, T, rows, seed, rank, world, num_workers):
    return torch.utils.data.DataLoader(
        MixDataset(dcfg, ncfg, T, rows, seed, rank, world), batch_size=None,
        num_workers=num_workers, prefetch_factor=(2 if num_workers > 0 else None),
        persistent_workers=(num_workers > 0), pin_memory=True)


def lr_mult(step, total, warmup):
    if step < warmup:
        return step / max(warmup, 1)
    t = (step - warmup) / max(total - warmup, 1)
    return 0.02 + 0.98 * 0.5 * (1 + math.cos(math.pi * t))   # cosine to near-zero


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", required=True)
    a = ap.parse_args()
    cfg = json.loads(Path(a.config).read_text())
    cfg["out_dir"] = os.path.expandvars(cfg["out_dir"])
    rank, world, is_ddp = ddp_setup()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(cfg["seed"] + rank)
    torch.set_float32_matmul_precision("high")

    T = cfg["seq_len"]; rows = cfg["micro_batch"]; accum = cfg.get("grad_accum", 1)
    total = cfg["total_steps"]
    out = Path(cfg["out_dir"]); out.mkdir(parents=True, exist_ok=True)

    mcfg = CharBertConfig(attn_impl=cfg.get("attn", "flex"), d_model=cfg["d_model"],
                          n_heads=cfg["d_model"] // 64, depth=cfg["depth"],
                          char_window=cfg["char_window"], qk_norm=True)
    model = CharBertEncoder(mcfg).to(device)

    step0 = 0
    ckpt = out / "last.pt"
    if ckpt.exists():                                  # resume a crashed finetune
        sd = torch.load(ckpt, map_location=device)
        model.load_state_dict(sd["model"]); step0 = sd["step"]
        if rank == 0:
            print(f"resumed finetune from step {step0}")
    else:                                              # init from the pretrained torso
        torso = os.path.expandvars(cfg["torso"])
        sd = torch.load(torso, map_location=device)
        model.load_state_dict(sd["model"])
        if rank == 0:
            print(f"initialized from torso {torso} (pretrain step {sd.get('step')})")
    if step0 >= total:
        if rank == 0:
            print("finetune already complete")
        return
    if rank == 0:
        print(f"params={num_params(model)/1e6:.1f}M T={T} rows={rows}x{accum} world={world} "
              f"total={total}")
    if is_ddp:
        from torch.nn.parallel import DistributedDataParallel as DDP
        model = DDP(model, device_ids=[rank % torch.cuda.device_count()])
    fwd = torch.compile(model) if cfg.get("compile", True) and device.type == "cuda" else model

    decay_p = [p for p in model.parameters() if p.ndim >= 2]
    nodecay_p = [p for p in model.parameters() if p.ndim < 2]
    opt = torch.optim.AdamW([{"params": decay_p, "weight_decay": cfg.get("wd", 0.1)},
                             {"params": nodecay_p, "weight_decay": 0.0}],
                            lr=cfg["lr"], betas=(0.9, 0.95), fused=(device.type == "cuda"))
    if ckpt.exists() and "opt" in sd:
        try:
            opt.load_state_dict(sd["opt"])
        except Exception:
            pass

    # noise: Ithaca-oriented. Spans dominate, geometric length tuned to the 1-10 band.
    nc = cfg.get("noise", {})
    ncfg = NoiseConfig(w_span=nc.get("w_span", 0.45), w_word=nc.get("w_word", 0.15),
                       w_elastic=nc.get("w_elastic", 0.0), w_iid=nc.get("w_iid", 0.15),
                       w_halfword=nc.get("w_halfword", 0.15),
                       w_substitute=nc.get("w_substitute", 0.10),
                       span_mean=nc.get("span_mean", 4.0), span_max=nc.get("span_max", 12))

    ins = os.path.expandvars(cfg["iphi_shards"])
    gcb = os.path.expandvars(cfg["shards"])
    mix = cfg.get("mix", {"iphi": 1.0, "iphi_syn": 0.5, "gold": 0.15, "silver": 0.15})
    dcfg = DataConfig(tiers={
        "iphi":     TierSpec(ins, mix.get("iphi", 1.0), tier_filter="iphi"),
        "iphi_syn": TierSpec(ins, mix.get("iphi_syn", 0.5), tier_filter="iphi_syn"),
        "gold":     TierSpec(gcb, mix.get("gold", 0.15), tier_filter="pristine"),
        "silver":   TierSpec(gcb, mix.get("silver", 0.15), tier_filter="repaired"),
    }, window_chars=T, seed=cfg["seed"] + step0)

    it = iter(make_loader(dcfg, ncfg, T, rows, cfg["seed"] + step0, rank, world,
                          cfg.get("num_workers", 8)))

    model.train()
    t0 = time.time(); seen = 0
    metrics_f = out / "metrics.jsonl"
    import contextlib
    for step in range(step0, total):
        lr = cfg["lr"] * lr_mult(step, total, cfg.get("warmup", 200))
        for pg in opt.param_groups:
            pg["lr"] = lr
        opt.zero_grad(set_to_none=True)
        for micro in range(accum):
            batch = next(it)
            batch = {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in batch.items()}
            sync = (model.no_sync() if (is_ddp and micro < accum - 1)
                    else contextlib.nullcontext())
            with sync:
                with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
                    out_h = fwd(batch)
                    loss, logs = compute_loss(out_h, batch, lam=cfg.get("lam", 0.1))
                (loss / accum).backward()
            seen += int((batch["seg_id"] > 0).sum()) * world
        gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        if torch.isfinite(gnorm):
            opt.step()

        if rank == 0 and step % cfg.get("log_every", 20) == 0:
            rec = dict(step=step, lr=round(lr, 7), gnorm=round(float(gnorm), 3),
                       chps=round(seen / (time.time() - t0) / 1e6, 3), **logs)
            print("  " + " ".join(f"{k}={v}" for k, v in rec.items()), flush=True)
            with open(metrics_f, "a") as f:
                f.write(json.dumps(rec) + "\n")
        if rank == 0 and step > step0 and step % cfg.get("ckpt_every", 500) == 0:
            core = model.module if is_ddp else model
            save_ckpt(dict(model=core.state_dict(), opt=opt.state_dict(), step=step + 1,
                           cfg=cfg), ckpt)

    if rank == 0:
        core = model.module if is_ddp else model
        state = dict(model=core.state_dict(), opt=opt.state_dict(), step=total, cfg=cfg)
        save_ckpt(state, out / "final.pt")
        save_ckpt(state, ckpt)
        print("FINETUNE DONE")
    if is_ddp:
        import torch.distributed as dist
        dist.destroy_process_group()


if __name__ == "__main__":
    main()