Buckets:
| """Train one arm (E1 / E2 / E3) from S1. Single GPU, fp32 master weights, bf16 autocast. | |
| Per optimizer step: 32 examples (seeded shuffle), one K drawn for the whole step (own seeded RNG, so arms | |
| share the K sequence), examples length-sorted into micro-batches whose padded tokens x K stay under a | |
| budget; loss = sum of S1's per-decision losses / 32 (CE with soft targets, + 0.25 RPS for score). | |
| """ | |
| import argparse | |
| import json | |
| import math | |
| import os | |
| import random | |
| import sys | |
| import time | |
| import numpy as np | |
| import torch | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from common import DATA, S1_PATH, Renderer, load_canonical, tokenizer_and_mapping # noqa: E402 | |
| from model import ThinkDecide, collate, item_logits, load_backbone # noqa: E402 | |
| from shared_yesno_supervised import supervised_terms # noqa: E402 | |
| def micro_batches(items, K, budget, max_items, waste=0.15, min_tokens=2048): | |
| """Length-sorted greedy split: padded_len * n * K <= budget, and a new micro-batch starts once padding | |
| would exceed `waste` of it (if the current one already holds >= min_tokens). Padded/real ~1.17 at every | |
| K (vs 1.3-3.0 with the budget rule alone); batching does not change the loss (per-item sum / batch).""" | |
| items = sorted(items, key=lambda x: len(x["ids"])) | |
| out, cur, tot = [], [], 0 | |
| for x in items: | |
| L, n = len(x["ids"]), len(cur) + 1 | |
| if cur and (L * n * K > budget or len(cur) >= max_items | |
| or (L * n - (tot + L) > waste * L * n and L * len(cur) >= min_tokens)): | |
| out.append(cur) | |
| cur, tot = [], 0 | |
| cur.append(x) | |
| tot += L | |
| if cur: | |
| out.append(cur) | |
| return out | |
| def decision_losses(model, items, K, device, pad_id, rps_weight, distill=None): | |
| batch = collate(items, device, pad_id) | |
| s = model(batch, K=K)[-1] | |
| terms = [] | |
| for lg, x in zip(item_logits(s, batch), items): | |
| t = torch.tensor([x["target"]], dtype=torch.float32, device=device) | |
| terms.append(supervised_terms(lg, t, x["primitive"], rps_weight)) | |
| return terms | |
| def evaluate_dev(model, items, Ks, device, pad_id, budget=48000): | |
| model.eval() | |
| res = {K: dict(ce=0.0, kl=0.0, correct=0, hard=0, n=0) for K in Ks} | |
| for mb in micro_batches(items, max(Ks), budget * 4, 64): | |
| batch = collate(mb, device, pad_id) | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| outs = model(batch, K=max(Ks), read_at=Ks) | |
| for K, s in zip(sorted(set(Ks)), outs): | |
| for lg, x in zip(item_logits(s, batch), mb): | |
| t = torch.tensor([x["target"]], dtype=torch.float32, device=device) | |
| tm = supervised_terms(lg.float(), t, x["primitive"], 0.0) | |
| r = res[K] | |
| r["ce"] += tm["ce"].item() | |
| r["kl"] += tm["kl"].item() | |
| r["n"] += 1 | |
| if x["kind"] == "hard": | |
| r["hard"] += 1 | |
| r["correct"] += int(lg.argmax().item() == int(np.argmax(x["target"]))) | |
| model.train() | |
| return {K: dict(ce=r["ce"] / r["n"], kl=r["kl"] / r["n"], acc=r["correct"] / max(1, r["hard"]), n=r["n"]) | |
| for K, r in res.items()} | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--n-think", type=int, default=0) | |
| ap.add_argument("--no-head", action="store_true") | |
| ap.add_argument("--k-choices", default="1") | |
| ap.add_argument("--epochs", type=float, default=1.0) | |
| ap.add_argument("--batch", type=int, default=32) | |
| ap.add_argument("--lr-bb", type=float, default=1e-5) | |
| ap.add_argument("--lr-new", type=float, default=1e-4) | |
| ap.add_argument("--min-lr", type=float, default=1e-6) | |
| ap.add_argument("--wd", type=float, default=0.01) | |
| ap.add_argument("--warmup", type=float, default=0.05) | |
| ap.add_argument("--rps-weight", type=float, default=0.25) | |
| ap.add_argument("--max-length", type=int, default=2048) | |
| ap.add_argument("--budget", type=int, default=40000, help="padded tokens x K per micro-batch") | |
| ap.add_argument("--max-items", type=int, default=32) | |
| ap.add_argument("--ckpt", choices=["auto", "on", "off"], default="off") | |
| ap.add_argument("--ckpt-above", type=int, default=1 << 30, help="auto: checkpoint micro-batches above this tokens x K") | |
| ap.add_argument("--dev-every", type=int, default=250) | |
| ap.add_argument("--dev-ks", default="1,4") | |
| ap.add_argument("--seed", type=int, default=42) | |
| ap.add_argument("--train-limit", type=int, default=0) | |
| ap.add_argument("--dev-limit", type=int, default=0) | |
| ap.add_argument("--max-steps", type=int, default=0) | |
| ap.add_argument("--distill", default="", help="E3: teacher cache (.pt)") | |
| ap.add_argument("--distill-l1", type=float, default=1.0) | |
| ap.add_argument("--distill-kl", type=float, default=0.5) | |
| ap.add_argument("--distill-layers", default="14,28") | |
| ap.add_argument("--no-save", action="store_true") | |
| ap.add_argument("--no-compile", action="store_true") | |
| ap.add_argument("--train-file", default="") | |
| ap.add_argument("--dev-file", default="") | |
| ap.add_argument("--fb-mode", choices=["hidden", "soft"], default="hidden") | |
| ap.add_argument("--gate-init", type=float, default=0.0) | |
| ap.add_argument("--gate-fixed", action="store_true") | |
| a = ap.parse_args() | |
| os.makedirs(a.out, exist_ok=True) | |
| torch.manual_seed(a.seed) | |
| random.seed(a.seed) | |
| torch.backends.cuda.matmul.allow_tf32 = True | |
| torch.backends.cudnn.allow_tf32 = True | |
| device = torch.device("cuda") | |
| tok, mapping, mask_id, mask_text = tokenizer_and_mapping() | |
| rend = Renderer(tok, mask_id, mask_text, a.n_think) | |
| train, dropped = load_canonical(a.train_file or DATA + "/train.jsonl", rend, a.max_length, a.train_limit) | |
| dev, _ = load_canonical(a.dev_file or DATA + "/dev.jsonl", rend, 10**9, a.dev_limit) | |
| Ks = [int(k) for k in a.k_choices.split(",")] | |
| dev_ks = sorted({int(k) for k in a.dev_ks.split(",")}) if a.n_think else [1] | |
| teacher = None | |
| if a.distill: | |
| from distill import attach_teacher | |
| teacher = attach_teacher(train, a.distill) | |
| print(json.dumps(dict(event="distill", with_target=sum("teacher" in x for x in train))), flush=True) | |
| base = load_backbone(S1_PATH, torch.float32, device) | |
| model = ThinkDecide(base, mapping["Yes"], mapping["No"], n_think=a.n_think, use_head=not a.no_head, | |
| fb_mode=a.fb_mode, gate_init=a.gate_init, gate_fixed=a.gate_fixed).to(device) | |
| model.train() | |
| if not a.no_compile: | |
| model.compile_layers() | |
| opt = torch.optim.AdamW(model.param_groups(a.lr_bb, a.lr_new, a.wd), betas=(0.9, 0.999), eps=1e-8, fused=True) | |
| steps_per_epoch = len(train) // a.batch | |
| total = int(round(steps_per_epoch * a.epochs)) | |
| if a.max_steps: | |
| total = min(total, a.max_steps) | |
| warm = math.ceil(total * a.warmup) | |
| def lr_factor(step, peak): | |
| if step < warm: | |
| return (step + 1) / warm | |
| prog = min(1.0, (step - warm) / max(1, total - warm)) | |
| floor = a.min_lr / peak | |
| return floor + (1 - floor) * 0.5 * (1 + math.cos(math.pi * prog)) | |
| peaks = [g["lr"] for g in opt.param_groups] | |
| config = dict(vars(a), train_count=len(train), dev_count=len(dev), dropped_overlong=dropped, total_steps=total, | |
| warmup_steps=warm, Ks=Ks, n_params=sum(p.numel() for p in model.parameters()), | |
| n_new=sum(p.numel() for n, p in model.named_parameters() if not n.startswith("base.")), | |
| git=os.popen("git -C /workspace/think-decide rev-parse --short HEAD 2>/dev/null").read().strip()) | |
| json.dump(config, open(os.path.join(a.out, "config.json"), "w"), indent=1) | |
| print(json.dumps(dict(event="start", **{k: config[k] for k in ("train_count", "dev_count", "total_steps", | |
| "n_params", "n_new", "dropped_overlong")})), | |
| flush=True) | |
| hist = open(os.path.join(a.out, "train_log.jsonl"), "a") | |
| def log(rec): | |
| hist.write(json.dumps(rec) + "\n") | |
| hist.flush() | |
| print(json.dumps(rec), flush=True) | |
| def dev_eval(step): | |
| t = time.time() | |
| r = evaluate_dev(model, dev, dev_ks, device, tok.pad_token_id) | |
| log(dict(event="dev", step=step, seconds=round(time.time() - t, 1), | |
| gate=model.gate.item() if a.n_think else None, | |
| **{"K%d" % K: {k: round(v, 5) for k, v in m.items()} for K, m in r.items()})) | |
| dev_eval(0) | |
| order_rng = random.Random(a.seed) | |
| k_rng = random.Random(a.seed + 1) | |
| t0, tokens_done, passes_done = time.time(), 0, 0 | |
| params = [p for p in model.parameters() if p.requires_grad] | |
| n_oom = 0 | |
| def big(mb, K): # snapshot grads only for micro-batches near the memory budget | |
| return max(len(x["ids"]) for x in mb) * len(mb) * K > 0.6 * a.budget | |
| step, epoch, order, pos = 0, 0, [], 0 | |
| torch.cuda.reset_peak_memory_stats() | |
| while step < total: | |
| if pos + a.batch > len(order): | |
| order = list(range(len(train))) | |
| order_rng.shuffle(order) | |
| pos = 0 | |
| epoch += 1 | |
| idx = order[pos:pos + a.batch] | |
| pos += a.batch | |
| K = k_rng.choice(Ks) | |
| for g, peak in zip(opt.param_groups, peaks): | |
| g["lr"] = peak * lr_factor(step, peak) | |
| sums = dict(loss=0.0, ce=0.0, kl=0.0, rps=0.0, l1=0.0, dkl=0.0) | |
| n_ex = len(idx) | |
| for mb in micro_batches([train[i] for i in idx], K, a.budget, a.max_items): | |
| L = max(len(x["ids"]) for x in mb) | |
| ckpt = a.ckpt == "on" or (a.ckpt == "auto" and L * len(mb) * K > a.ckpt_above) | |
| snap = [p.grad.clone() if p.grad is not None else None for p in params] if big(mb, K) else None | |
| for attempt in (0, 1): | |
| model.ckpt = ckpt or attempt == 1 | |
| try: | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| if teacher is not None: | |
| from distill import distill_losses | |
| terms, extra = distill_losses(model, mb, K, device, tok.pad_token_id, a) | |
| else: | |
| terms, extra = decision_losses(model, mb, K, device, tok.pad_token_id, a.rps_weight), None | |
| loss = sum(t["loss"] for t in terms) / n_ex | |
| if extra is not None: | |
| loss = loss + extra["loss"] / n_ex | |
| loss.backward() | |
| break | |
| except torch.OutOfMemoryError: | |
| if attempt or snap is None: | |
| raise | |
| terms = extra = loss = None | |
| for p, g in zip(params, snap): # undo any partial backward, redo with checkpointing | |
| p.grad = g | |
| torch.cuda.empty_cache() | |
| n_oom += 1 | |
| log(dict(event="oom_retry", step=step, K=K, tokens=L * len(mb))) | |
| for t in terms: | |
| for k in ("loss", "ce", "kl", "rps"): | |
| sums[k] += t[k].item() | |
| if extra is not None: | |
| sums["l1"] += extra["l1"] | |
| sums["dkl"] += extra["dkl"] | |
| tokens_done += sum(len(x["ids"]) for x in mb) | |
| passes_done += sum(len(x["ids"]) for x in mb) * K | |
| gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) | |
| if not torch.isfinite(gn): | |
| raise FloatingPointError("non-finite grad norm at step %d" % step) | |
| opt.step() | |
| opt.zero_grad(set_to_none=True) | |
| step += 1 | |
| if step % 10 == 0 or step == total: | |
| el = time.time() - t0 | |
| rec = dict(step=step, total=total, K=K, lr_bb=opt.param_groups[0]["lr"], grad_norm=round(gn.item(), 4), | |
| **{k: round(v / n_ex, 5) for k, v in sums.items()}, elapsed=round(el, 1), | |
| tok_per_s=round(tokens_done / el), pass_tok_per_s=round(passes_done / el), | |
| peak_gib=round(torch.cuda.max_memory_allocated() / 2**30, 2), epoch=epoch) | |
| if a.n_think: | |
| rec["gate"] = round(model.gate.item(), 6) | |
| rec["W_dev"] = round((model.fb_proj.weight - torch.eye(model.fb_proj.weight.shape[0], | |
| device=device)).norm().item(), 4) | |
| if hasattr(model, "head"): | |
| rec["head_o"] = round(model.head.o.weight.norm().item(), 4) | |
| log(rec) | |
| if step % a.dev_every == 0 and step != total: | |
| dev_eval(step) | |
| dev_eval(step) | |
| log(dict(event="finished", step=step, elapsed=round(time.time() - t0, 1), oom_retries=n_oom, | |
| peak_gib=round(torch.cuda.max_memory_allocated() / 2**30, 2))) | |
| if not a.no_save: | |
| from safetensors.torch import save_file | |
| sd = {k: v.detach().to(torch.bfloat16).contiguous() for k, v in model.state_dict().items()} | |
| # tied embeddings: drop the duplicate lm_head entry (restored on load) | |
| sd.pop("base.lm_head.weight", None) | |
| sd = {k: v for k, v in sd.items() if not k.startswith("bb.")} # bb is base.model under a second name | |
| save_file(sd, os.path.join(a.out, "model.safetensors")) | |
| print(json.dumps(dict(event="saved", path=os.path.join(a.out, "model.safetensors"))), flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 13.4 kB
- Xet hash:
- 268d2f34e3c32105b4adc08b48b9e3d594edce50d7db61afb2d6256934f790f9
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.