| """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) |
| |
| 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() |
|
|
