"""Stage B: fine-tune a Stage-A checkpoint with ground-truth pinball + teacher-ensemble pinball (TimesFM-2.5 + Chronos-2). python -m modular_ts.distill --init runs/run2_8k/ckpt.pt --teach teach/synth,teach/gep --lam 0.5 --minutes 45 --out runs/run3 Batches: with prob --p_distill a teacher batch (V=1, T=C+F windows, min-context split s=C/32; synthetic rows: GT pinball + teacher pinball; GiftEvalPretrain rows: teacher pinball ONLY, their future is filled with the teacher-ensemble median just for the normalisation frame and never used as a target), else a fresh synthetic Stage-A batch (GT pinball only) to avoid forgetting. Teacher ensemble = mean of the two teachers' quantiles (Vincentisation) at 0.1..0.9. """ from __future__ import annotations import argparse import glob import json import math import os import time import numpy as np import torch from modular_ts.model import ModularTS, count_params, make_config from modular_ts.synth import SynthStream from modular_ts.train import export, lr_at def load_teach(dirs, dev, max_per_dir=0): ctx, fut, q, src = [], [], [], [] for si, d in enumerate(dirs): files = sorted(glob.glob(os.path.join(d, "shard_*.npz"))) n = 0 for f in files: z = np.load(f) qe = 0.5 * (z["q_tfm"].astype(np.float32) + z["q_cr2"].astype(np.float32)) qe = np.sort(qe, -1) ctx.append(z["ctx"]); fut.append(z["fut"]); q.append(qe); src.append(np.full(len(qe), si)) n += len(qe) if max_per_dir and n >= max_per_dir: break print(f"[teach] {d}: {n} windows from {len(files)} shards", flush=True) c = torch.from_numpy(np.concatenate(ctx)).to(dev) f = torch.from_numpy(np.concatenate(fut)).to(dev) qq = torch.from_numpy(np.concatenate(q)).to(dev) s = torch.from_numpy(np.concatenate(src)).to(dev) return c, f, qq, s def main(): ap = argparse.ArgumentParser() ap.add_argument("--init", required=True, help="Stage-A ckpt.pt (model weights only are used)") ap.add_argument("--teach", required=True, help="comma list of teacher shard dirs; dirs containing 'gep' are teacher-only") ap.add_argument("--lam", type=float, default=0.5) ap.add_argument("--p_distill", type=float, default=0.6) ap.add_argument("--lr", type=float, default=3e-4) ap.add_argument("--warmup", type=int, default=200) ap.add_argument("--minutes", type=float, default=45) ap.add_argument("--tokens_per_batch", type=int, default=262144) ap.add_argument("--ctx", type=int, default=0, help="model context_length for export (0 = from init ckpt)") ap.add_argument("--compile", action="store_true") ap.add_argument("--log_every", type=int, default=100) ap.add_argument("--seed", type=int, default=7) ap.add_argument("--out", required=True) args = ap.parse_args() os.makedirs(args.out, exist_ok=True) dev = torch.device("cuda") torch.manual_seed(args.seed) torch.backends.cuda.matmul.allow_tf32 = True st = torch.load(args.init, map_location=dev) ctx_len = args.ctx or st["args"]["ctx"] cfg = make_config(st["args"].get("size", "2.5M"), context_length=ctx_len) model = ModularTS(cfg).to(dev) model.load_state_dict(st["model"]) print("params", count_params(model), "init", args.init, "step", st["step"], flush=True) dirs = [d for d in args.teach.split(",") if d] tc, tf, tq, tsrc = load_teach(dirs, dev) teacher_only = torch.tensor(["gep" in d for d in dirs], device=dev)[tsrc] C, F, H = tc.shape[1], tf.shape[1], tq.shape[1] T = C + F P = cfg.patch_size Bd = max(1, args.tokens_per_batch // (T // P)) nT = tc.shape[0] print(f"[teach] total {nT} windows (teacher-only {int(teacher_only.sum())}), C={C} F={F} H={H}, rows/batch {Bd}", flush=True) # hold out 2% of synthetic teacher windows for a GT/KD validation number g = torch.Generator(device=dev); g.manual_seed(args.seed) perm = torch.randperm(nT, generator=g, device=dev) syn_idx = perm[~teacher_only[perm]] val_idx = syn_idx[:min(4096, max(1, len(syn_idx) // 50))] trmask = torch.ones(nT, dtype=torch.bool, device=dev); trmask[val_idx] = False tr_idx = torch.nonzero(trmask).squeeze(1) def distill_batch(idx): c, f, q, to = tc[idx], tf[idx], tq[idx], teacher_only[idx] med = q[:, :, 4] fill = torch.full_like(f, float("nan")) fill[:, :H] = med f = torch.where(to[:, None], fill, f) vals = torch.cat([c, f], 1) obs = torch.isfinite(vals) vals = torch.where(obs, vals, torch.zeros_like(vals)) teach = torch.zeros(len(idx), T, q.shape[-1], device=dev) teach[:, C:C + H] = q tobs = torch.zeros(len(idx), T, dtype=torch.bool, device=dev) tobs[:, C:C + H] = True R = len(idx) return dict(values=vals[:, None], observed=obs[:, None], roles=torch.zeros(R, 1, dtype=torch.long, device=dev), s=torch.full((R,), C // P, dtype=torch.long, device=dev), teacher=teach, teacher_obs=tobs, gt_row=(~to).float()) synth = SynthStream(4096, args.tokens_per_batch, dev, seed=args.seed + 99) decay, no_decay = [], [] for n, p in model.named_parameters(): (decay if p.dim() >= 2 else no_decay).append(p) opt = torch.optim.AdamW([{"params": decay, "weight_decay": 0.05}, {"params": no_decay, "weight_decay": 0.0}], lr=args.lr, betas=(0.9, 0.95), eps=1e-8, fused=True) fwd = torch.compile(model.forward_train, dynamic=False) if args.compile else model.forward_train amp = torch.autocast(device_type="cuda", dtype=torch.bfloat16) lam = args.lam def evaluate(): model.eval() out = {"val_gt": 0.0, "val_kd": 0.0} with torch.no_grad(), amp: for i in range(0, len(val_idx), Bd): b = distill_batch(val_idx[i:i + Bd]) l, s_ = model.forward_train(b["values"], b["observed"], b["roles"], b["s"], teacher=b["teacher"], teacher_obs=b["teacher_obs"], gt_row=b["gt_row"], lam=lam) out["val_gt"] += float(s_["pinball"]); out["val_kd"] += float(s_["kd"]) n = math.ceil(len(val_idx) / Bd) model.train() return {k: v / n for k, v in out.items()} log = open(os.path.join(args.out, "log.jsonl"), "a") rng = np.random.default_rng(args.seed) t0 = time.time(); step = 0 rec0 = {"step": 0, **evaluate()} print(json.dumps(rec0), flush=True); log.write(json.dumps(rec0) + "\n") model.train() while (time.time() - t0) / 60 < args.minutes: frac = (time.time() - t0) / 60 / args.minutes lr = lr_at(step, max(step + 1, int(step / max(frac, 1e-6))) if step > args.warmup else 10**8, args.warmup, args.lr) for gp in opt.param_groups: gp["lr"] = lr is_d = rng.random() < args.p_distill with amp: if is_d: idx = tr_idx[torch.randint(0, len(tr_idx), (Bd,), device=dev)] b = distill_batch(idx) loss, s_ = fwd(b["values"], b["observed"], b["roles"], b["s"], False, b["teacher"], b["teacher_obs"], b["gt_row"], lam) else: b = synth.next_batch() loss, s_ = fwd(b["values"], b["observed"], b["roles"], b["s"]) if not torch.isfinite(loss): opt.zero_grad(set_to_none=True); step += 1 continue loss.backward() gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step(); opt.zero_grad(set_to_none=True) step += 1 if step % args.log_every == 0: rec = {"step": step, "lr": lr, "loss": float(loss), "kind": "distill" if is_d else "synth", "kd": float(s_.get("kd", float("nan"))), "gnorm": float(gn), "min": (time.time() - t0) / 60} if step % (args.log_every * 10) == 0: rec.update(evaluate()) print(json.dumps(rec), flush=True); log.write(json.dumps(rec) + "\n"); log.flush() torch.save({"model": model.state_dict(), "step": step, "args": {**st["args"], "distill": vars(args)}}, os.path.join(args.out, "ckpt.pt")) export(model, os.path.join(args.out, "export")) fin = {"steps": step, **evaluate(), "minutes": (time.time() - t0) / 60, "lam": lam} print("FINAL", json.dumps(fin), flush=True) json.dump(fin, open(os.path.join(args.out, "final.json"), "w")) if __name__ == "__main__": main()