| """Train Modular-TS from scratch with DHF + pinball loss. |
| |
| python -m modular_ts.train --size 2.5M --ctx 4096 --tokens_per_batch 65536 --steps 20000 \ |
| --lr 2e-3 --warmup 1000 --out runs/smoke --max_minutes 35 [--compile] [--data synthetic] |
| |
| Logs JSONL to <out>/log.jsonl (step, loss, val, lr, enc tokens/s, grad norm), |
| checkpoints to <out>/ckpt.pt and an inference export (config.json + model.safetensors, |
| loadable by timestral.TimestralForecaster) to <out>/export/. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import os |
| import time |
|
|
| import torch |
|
|
| from modular_ts.model import ModularTS, count_params, make_config |
| from modular_ts.synth import SynthStream |
|
|
|
|
| def export(model: ModularTS, out_dir: str) -> None: |
| from safetensors.torch import save_file |
| os.makedirs(out_dir, exist_ok=True) |
| sd = {k: v.detach().float().cpu().contiguous() for k, v in model.state_dict().items()} |
| save_file(sd, os.path.join(out_dir, "model.safetensors")) |
| model.config.save(os.path.join(out_dir, "config.json")) |
|
|
|
|
| def lr_at(step, total, warmup, peak, final_frac=0.05): |
| if step < warmup: |
| return peak * (step + 1) / warmup |
| prog = min(1.0, (step - warmup) / max(1, total - warmup)) |
| return peak * (final_frac + (1 - final_frac) * 0.5 * (1 + math.cos(math.pi * prog))) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--size", default="2.5M") |
| ap.add_argument("--ctx", type=int, default=4096) |
| ap.add_argument("--tokens_per_batch", type=int, default=65536, help="encoder patch tokens per step") |
| ap.add_argument("--steps", type=int, default=20000, help="schedule length (cosine horizon)") |
| ap.add_argument("--max_minutes", type=float, default=0, help="stop early after this wall time (0 = off)") |
| ap.add_argument("--lr", type=float, default=2e-3) |
| ap.add_argument("--warmup", type=int, default=1000) |
| ap.add_argument("--wd", type=float, default=0.05) |
| ap.add_argument("--clip", type=float, default=1.0) |
| ap.add_argument("--data", default="synthetic", choices=["synthetic", "synthetic+gifteval"]) |
| ap.add_argument("--real_frac", type=float, default=0.5) |
| ap.add_argument("--gifteval_subset", default="", help="comma list of GiftEvalPretrain sub-datasets (default all kept)") |
| ap.add_argument("--gep_selection", default="", help="json {dataset: [shard files]} -> fast GPU-resident GiftEvalPretrain stream") |
| ap.add_argument("--synth", default="v1", choices=["v1", "v2"], help="v2 = modular_ts.synth2 prior (v3)") |
| ap.add_argument("--p_new", type=float, default=0.55, help="synth v2: share of new generator families") |
| ap.add_argument("--real_v2", action="store_true", help="v3 real stream: resolution sampling + TSMixup + ds_temp 0.3") |
| ap.add_argument("--min_lr_frac", type=float, default=0.05) |
| ap.add_argument("--time_schedule", action="store_true", help="cosine progress = elapsed / max_minutes (smoke runs)") |
| ap.add_argument("--compile", action="store_true") |
| ap.add_argument("--bf16", action="store_true", default=True) |
| ap.add_argument("--no_bf16", dest="bf16", action="store_false") |
| ap.add_argument("--log_every", type=int, default=25) |
| ap.add_argument("--val_every", type=int, default=250) |
| ap.add_argument("--ckpt_every", type=int, default=1000) |
| ap.add_argument("--out", default="runs/dev") |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--resume", action="store_true") |
| ap.add_argument("--resume_from", default="", help="ckpt to resume model+opt from (default <out>/ckpt.pt)") |
| ap.add_argument("--reset_schedule", action="store_true", help="on resume, restart warmup+cosine at the resumed step") |
| ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| args = ap.parse_args() |
|
|
| os.makedirs(args.out, exist_ok=True) |
| torch.manual_seed(args.seed) |
| dev = torch.device(args.device) |
| if dev.type == "cuda": |
| torch.backends.cuda.matmul.allow_tf32 = True |
| torch.backends.cudnn.allow_tf32 = True |
|
|
| cfg = make_config(args.size, context_length=args.ctx) |
| model = ModularTS(cfg).to(dev) |
| pc = count_params(model) |
| print("params", pc, flush=True) |
|
|
| if args.synth == "v2": |
| from modular_ts.synth2 import SynthStream2 |
| synth = SynthStream2(args.ctx, args.tokens_per_batch, dev, seed=args.seed, p_new=args.p_new) |
| else: |
| synth = SynthStream(args.ctx, args.tokens_per_batch, dev, seed=args.seed) |
| stream = synth |
| manifest = {"data": args.data, "synth_version": args.synth, "real_v2": args.real_v2, "synthetic": "modular_ts.synth (KernelSynth-style GP + TSI + AR + intermittent/bursty + multivariatizers + covariate tasks)"} |
| if args.data == "synthetic+gifteval": |
| from modular_ts.realdata import GiftEvalPretrainStream, MixedStream |
| sub = [s for s in args.gifteval_subset.split(",") if s] or None |
| if args.gep_selection: |
| from modular_ts.realdata import FastGiftEvalPretrain |
| if args.real_v2: |
| from modular_ts.realdata2 import FastGiftEvalPretrain2 as FastGiftEvalPretrain |
| real = FastGiftEvalPretrain(args.ctx, args.tokens_per_batch, dev, json.load(open(args.gep_selection)), seed=args.seed) |
| else: |
| real = GiftEvalPretrainStream(args.ctx, args.tokens_per_batch, dev, seed=args.seed, datasets_subset=sub) |
| stream = MixedStream(synth, real, args.real_frac, seed=args.seed) |
| manifest.update(real.manifest()) |
| json.dump(manifest, open(os.path.join(args.out, "data_manifest.json"), "w"), indent=1) |
|
|
| |
| vgen = SynthStream(args.ctx, args.tokens_per_batch // 2, dev, seed=10_000 + args.seed, chunk=1) |
| val_batches = [vgen.next_batch(V) for V in (1, 4, 16)] |
|
|
| 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": args.wd}, {"params": no_decay, "weight_decay": 0.0}], |
| lr=args.lr, betas=(0.9, 0.95), eps=1e-8, fused=dev.type == "cuda") |
| step = 0 |
| ck = os.path.join(args.out, "ckpt.pt") |
| rck = args.resume_from or ck |
| step0 = 0 |
| prev_elapsed = 0.0 |
| if args.resume and os.path.exists(rck): |
| st = torch.load(rck, map_location=dev) |
| model.load_state_dict(st["model"]) |
| opt.load_state_dict(st["opt"]) |
| step = st["step"] |
| if args.reset_schedule: |
| step0 = step |
| elif args.time_schedule and "elapsed_min" in st: |
| prev_elapsed = st["elapsed_min"] |
| print("resumed at", step, "from", rck, flush=True) |
|
|
| fwd = model.forward_train |
| if args.compile: |
| try: |
| fwd = torch.compile(model.forward_train, dynamic=False) |
| except Exception as e: |
| print("compile failed, eager:", e, flush=True) |
|
|
| amp = torch.autocast(device_type=dev.type, dtype=torch.bfloat16, enabled=args.bf16 and dev.type == "cuda") |
|
|
| def evaluate(): |
| model.eval() |
| tot = 0.0 |
| with torch.no_grad(), amp: |
| for b in val_batches: |
| l, _ = model.forward_train(b["values"], b["observed"], b["roles"], b["s"]) |
| tot += l.item() |
| model.train() |
| return tot / len(val_batches) |
|
|
| log = open(os.path.join(args.out, "log.jsonl"), "a") |
| t0 = time.time() - prev_elapsed * 60 |
| t_last, tok_acc, gen_t = time.time(), 0, 0.0 |
| ema = None |
| model.train() |
| while step < args.steps: |
| if args.max_minutes and (time.time() - t0) / 60 > args.max_minutes: |
| print("time limit reached", flush=True) |
| break |
| prof = step % args.log_every == 0 |
| if prof and dev.type == "cuda": |
| torch.cuda.synchronize() |
| tg = time.time() |
| b = stream.next_batch() |
| if prof and dev.type == "cuda": |
| torch.cuda.synchronize() |
| if prof: |
| gen_t += (time.time() - tg) * args.log_every |
| ss = step - step0 |
| if args.time_schedule and args.max_minutes: |
| frac = (time.time() - t0) / 60 / args.max_minutes |
| lr = lr_at(ss, max(ss + 1, int(ss / max(frac, 1e-6))) if ss > args.warmup else args.steps, args.warmup, args.lr, args.min_lr_frac) |
| else: |
| lr = lr_at(ss, args.steps, args.warmup, args.lr, args.min_lr_frac) |
| for g in opt.param_groups: |
| g["lr"] = lr |
| try: |
| with amp: |
| loss, st = fwd(b["values"], b["observed"], b["roles"], b["s"]) |
| except Exception as e: |
| if fwd is model.forward_train: |
| raise |
| print("compiled forward failed, falling back to eager:", repr(e)[:500], flush=True) |
| fwd = model.forward_train |
| with amp: |
| loss, st = fwd(b["values"], b["observed"], b["roles"], b["s"]) |
| if not torch.isfinite(loss): |
| print(f"non-finite loss at step {step}; skipping batch", flush=True) |
| opt.zero_grad(set_to_none=True) |
| step += 1 |
| continue |
| loss.backward() |
| gn = torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip) |
| opt.step() |
| opt.zero_grad(set_to_none=True) |
| tok_acc += b["values"].shape[0] * b["values"].shape[1] * (args.ctx // cfg.patch_size) |
| step += 1 |
| lv = loss.item() if step % args.log_every == 0 or step == 1 else None |
| if lv is not None: |
| ema = lv if ema is None else 0.9 * ema + 0.1 * lv |
| if dev.type == "cuda": |
| torch.cuda.synchronize() |
| dt = time.time() - t_last |
| rec = {"step": step, "loss": lv, "loss_ema": ema, "lr": lr, "gnorm": float(gn), |
| "enc_tok_s": tok_acc / dt, "gen_frac": gen_t / dt, "elapsed_min": (time.time() - t0) / 60} |
| if step % args.val_every == 0: |
| rec["val"] = evaluate() |
| if dev.type == "cuda": |
| rec["mem_gb"] = torch.cuda.max_memory_allocated() / 1e9 |
| print(json.dumps(rec), flush=True) |
| log.write(json.dumps(rec) + "\n") |
| log.flush() |
| t_last, tok_acc, gen_t = time.time(), 0, 0.0 |
| if step % args.ckpt_every == 0: |
| torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step, "args": vars(args), "elapsed_min": (time.time() - t0) / 60}, ck) |
| export(model, os.path.join(args.out, "export")) |
| torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": step, "args": vars(args), "elapsed_min": (time.time() - t0) / 60}, ck) |
| export(model, os.path.join(args.out, "export")) |
| fin = {"final_step": step, "val": evaluate(), "params": pc, "minutes": (time.time() - t0) / 60} |
| print("FINAL", json.dumps(fin), flush=True) |
| json.dump(fin, open(os.path.join(args.out, "final.json"), "w")) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
