1#!/usr/bin/env python2"""Stage A: backbone pretraining (plain causal decoder; workspace disabled) on a uint16 token stream.34Usage: python -m training.pretrain --config configs/experiments/pretrain_103m.yaml [--dry-run] [--device mps]5The experiment YAML has keys: model (path to model yaml or inline dict), data (train_bin, val_bin), train6(batch_size, seq_len, steps, lr, warmup, weight_decay, grad_clip, eval_every, eval_batches, ckpt_every), seed.7"""8from __future__ import annotations910import argparse, json, math, sys, time11from pathlib import Path1213import torch1415sys.path.insert(0, str(Path(__file__).resolve().parents[1]))16from prisme.config import PrismeConfig, load_config # noqa: E40217from prisme.model import PrismeModel # noqa: E40218from prisme.losses import lm_loss # noqa: E40219from training.common import (Run, TokenBin, cosine_lr, load_checkpoint, load_yaml, make_adamw, peak_memory_bytes, # noqa: E40220 pick_device, save_checkpoint, seed_all, sha256_file, sync)212223def build_model_cfg(exp: dict) -> PrismeConfig:24 m = exp["model"]25 cfg = load_config(m) if isinstance(m, str) else PrismeConfig.from_dict(m)26 for k, v in exp.get("model_overrides", {}).items():27 setattr(cfg, k, v)28 cfg.validate()29 return cfg303132@torch.no_grad()33def evaluate(model, bin_: TokenBin, batch_size, device, max_batches):34 model.eval()35 tot, n = 0.0, 036 for x, y in bin_.sequential_batches(batch_size, device, max_batches):37 logits = model(x, None)["logits"]38 tot += lm_loss(logits, y).item() * y.numel(); n += y.numel()39 model.train()40 return tot / max(n, 1)414243def main(argv=None) -> int:44 ap = argparse.ArgumentParser(description=__doc__)45 ap.add_argument("--config", required=True)46 ap.add_argument("--dry-run", action="store_true", help="build model/data, run 2 steps on CPU-sized batch, exit")47 ap.add_argument("--device", default="auto")48 ap.add_argument("--seed", type=int, default=None)49 ap.add_argument("--resume", default=None, help="checkpoint dir to resume weights from")50 ap.add_argument("--run-root", default="experiments/runs")51 args = ap.parse_args(argv)52 exp = load_yaml(args.config)53 seed = args.seed if args.seed is not None else exp.get("seed", 1)54 seed_all(seed)55 device = pick_device(args.device)56 cfg = build_model_cfg(exp)57 cfg.use_workspace = False # Stage A: plain backbone58 tr = exp["train"]59 if args.dry_run:60 tr = {**tr, "steps": 2, "batch_size": 2, "seq_len": min(64, tr["seq_len"]), "eval_batches": 1, "eval_every": 1, "ckpt_every": 10**9}61 device = torch.device("cpu")62 model = PrismeModel(cfg).to(device)63 n_params = sum(p.numel() for p in model.parameters())64 train_bin = TokenBin(exp["data"]["train_bin"], tr["seq_len"])65 val_bin = TokenBin(exp["data"]["val_bin"], tr["seq_len"])66 run = Run(exp.get("name", "pretrain"), {"experiment": exp, "model": cfg.to_dict(), "resolved_train": tr}, root=args.run_root, seed=seed)67 run.write_manifest(device=str(device), backend=device.type, precision="fp32", n_params=n_params,68 dataset_manifest_sha256=sha256_file(Path(exp["data"]["train_bin"]).parent / "manifest.json") if (Path(exp["data"]["train_bin"]).parent / "manifest.json").exists() else None,69 checkpoint_parent=args.resume)70 if args.resume:71 meta = load_checkpoint(args.resume, model, strict=False)72 run.log(event="resumed", missing=meta["_missing"], unexpected=meta["_unexpected"])73 opt = make_adamw(model.parameters(), tr["lr"], tr.get("weight_decay", 0.1))74 steps = tr["steps"]75 tokens_per_step = tr["batch_size"] * tr["seq_len"]76 batches = train_bin.batches(tr["batch_size"], seed, device)77 model.train()78 t0 = time.time(); seen = 0; best_val = math.inf79 for step in range(steps):80 lr = cosine_lr(step, tr.get("warmup", 200), steps, tr["lr"])81 for g in opt.param_groups:82 g["lr"] = lr83 x, y = next(batches)84 loss = lm_loss(model(x, None)["logits"], y)85 opt.zero_grad(set_to_none=True)86 loss.backward()87 gn = torch.nn.utils.clip_grad_norm_(model.parameters(), tr.get("grad_clip", 1.0))88 opt.step()89 seen += tokens_per_step90 if step % tr.get("log_every", 10) == 0 or step == steps - 1:91 sync(device)92 el = time.time() - t093 run.log(step=step, loss=loss.item(), lr=lr, grad_norm=float(gn), tokens=seen, tok_per_s=seen / max(el, 1e-6))94 print(f"step {step} loss {loss.item():.4f} lr {lr:.2e} gn {float(gn):.2f} tok/s {seen/max(el,1e-6):.0f}", flush=True)95 if (step + 1) % tr["eval_every"] == 0 or step == steps - 1:96 vl = evaluate(model, val_bin, tr["batch_size"], device, tr.get("eval_batches", 20))97 run.log(step=step, val_loss=vl, val_ppl=math.exp(vl), tokens=seen)98 print(f" val loss {vl:.4f} ppl {math.exp(vl):.2f}", flush=True)99 if vl < best_val:100 best_val = vl101 if (step + 1) % tr["ckpt_every"] == 0 or step == steps - 1:102 sha = save_checkpoint(run.dir / "checkpoint", model, cfg.to_dict(), {"step": step + 1, "tokens": seen, "run_id": run.run_id})103 run.log(step=step, checkpoint_sha256=sha)104 sync(device)105 run.finish(train_tokens=seen, estimated_train_flops=6 * n_params * seen, peak_memory_bytes=peak_memory_bytes(device),106 final_val_loss=best_val, metrics_file=str(run.dir / "metrics.jsonl"))107 print("run", run.run_id, "done; val", best_val)108 return 0109110111if __name__ == "__main__":112 raise SystemExit(main())113