File size: 8,526 Bytes
51bb0b7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | """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()
|