Buckets:
| """Joint and readout-only LoRA training. | |
| Reasoner examples: the shortest correct, self-closed chain of an item, trained whole (cross-entropy on the | |
| chain, the answer letter and the end-of-turn token). Readout examples: the no-think prompt with shuffled | |
| options; cross-entropy of the K-way letter softmax at the last position against | |
| q = alpha * gold + (1 - alpha) * votes. | |
| Loss per optimizer step = mean over its examples of each example's own loss (lambda = 1). | |
| Nothing is truncated: a long example gets a micro-batch to itself. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import collections | |
| import json | |
| import math | |
| import os | |
| import random | |
| import time | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from torch.utils.checkpoint import checkpoint | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, get_cosine_schedule_with_warmup | |
| from common import LETTERS, letter_ids, read_jsonl, render, shuffled_order, split_chain | |
| LORA_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] | |
| LINEAR_ATTN_TARGETS = ["in_proj_qkv", "in_proj_z", "out_proj"] # Qwen3.5 gated-delta layers | |
| # ------------------------------------------------------------------ examples | |
| def vote_share(item: dict, rec: dict) -> tuple[np.ndarray | None, int]: | |
| """Vote distribution over original option indices, and the number of parsed answers.""" | |
| k = len(item["options"]) | |
| votes = np.zeros(k) | |
| n = 0 | |
| for s in rec["samples"]: | |
| if s["answer"] is not None and s["closed"]: | |
| votes[rec["order"][s["answer"]]] += 1 | |
| n += 1 | |
| return (votes / n if n else None), n | |
| import re | |
| # Qwen3.5 sometimes pretends to look an item up ("Wait, I found similar examples online"): invented knowledge, not reasoning | |
| RECALL_RE = re.compile(r"found (?:a|an|the) (?:similar|same|exact) (?:example|item|question)|(?:examples?|answers?|this) (?:found )?online" | |
| r"|\bon the internet\b|\bI recall (?:that )?(?:this|the) (?:dataset|example|item)|known (?:example|item) from" | |
| r"|\bfound (?:a|an) (?:reference|source) (?:to|for)\b" | |
| r"|\bI (?:recall|remember) (?:seeing |a |this |similar |the |that |specific )*(?:dataset|example|benchmark|item|question|sentence)s?\b" | |
| r"|\bsearch(?:ed)? for (?:the|this) exact\b", re.I) | |
| CW_RE = re.compile(r"\.cw(\s*</think>)") # a stray '.cw' the model writes right before </think> | |
| def repeated_line(text: str, min_len: int = 25) -> int: | |
| """How often the most repeated substantial line (>= min_len characters) occurs in a chain.""" | |
| c = collections.Counter(l.strip() for l in text.split("\n") if len(l.strip()) >= min_len) | |
| return max(c.values()) if c else 0 | |
| def build_examples(pool: dict, gen: list[dict], tok, mode: str, alpha: float, rng: random.Random, | |
| counts: collections.Counter, readout_ids: list[str] | None = None, min_correct: int = 1, | |
| second_chain_min: int = 0, rep_max: int = 0, post_soft: bool = False, | |
| eval_order_frac: float = 0.0, drop_recall: bool = False, strip_cw: bool = False) -> list[dict]: | |
| """Readout examples for the generated items (v1) or for every id in readout_ids (V2: gold where an item has no | |
| chains, alpha * gold + (1 - alpha) * votes where it has at least 2 parsed answers); reasoner examples in joint mode. | |
| V2 options (defaults reproduce v1): min_correct (an item needs this many correct chains), second_chain_min (also | |
| train the second-shortest correct chain when at least this many are correct), rep_max (drop chains whose most | |
| repeated long line occurs this often), post_soft (the letter after a chain gets the soft target | |
| 0.5 * own answer + 0.5 * votes instead of being one more hard token), eval_order_frac (share of readout examples | |
| shown in dataset option order instead of shuffled).""" | |
| im_end = tok.convert_tokens_to_ids("<|im_end|>") | |
| think_end = tok.convert_tokens_to_ids("</think>") | |
| lids = letter_ids(tok) | |
| by_gen = {r["id"]: r for r in gen} | |
| out = [] | |
| for iid in ([r["id"] for r in gen] if readout_ids is None else readout_ids): | |
| item = pool.get(iid) | |
| if item is None: | |
| counts["readout_item_missing"] += 1 | |
| continue | |
| k = len(item["options"]) | |
| gold = np.asarray(item["target"], dtype=np.float64) | |
| rec = by_gen.get(iid) | |
| votes, n = vote_share(item, rec) if rec else (None, 0) | |
| if votes is not None and n >= 2: | |
| q = alpha * gold + (1 - alpha) * votes | |
| counts["readout_from_votes"] += 1 | |
| else: | |
| q = gold | |
| counts["readout_gold_only" if rec is None else "readout_gold_fallback"] += 1 | |
| if eval_order_frac and rng.random() < eval_order_frac: | |
| order = list(range(k)) | |
| counts["readout_eval_order"] += 1 | |
| else: | |
| order = shuffled_order(item, rng) | |
| ids = tok(render(tok, item, think=False, order=order), add_special_tokens=False)["input_ids"] | |
| out.append({"kind": "readout", "id": item["id"], "input_ids": ids, | |
| "letters": lids[:k], "q": [float(q[j]) for j in order]}) | |
| if mode != "joint": | |
| return out | |
| for rec in gen: | |
| item = pool.get(rec["id"]) | |
| if item is None: | |
| continue | |
| k = len(item["options"]) | |
| correct = sorted((s for s in rec["samples"] if s["closed"] and s["answer"] is not None | |
| and rec["order"][s["answer"]] == item["gold"]), key=lambda s: s["tokens"]) | |
| if not correct: | |
| counts["reasoner_no_correct_chain"] += 1 | |
| continue | |
| if len(correct) < min_correct: | |
| counts["reasoner_too_few_correct"] += 1 | |
| continue | |
| if drop_recall: | |
| keep = [s for s in correct if not RECALL_RE.search(s["text"])] | |
| counts["reasoner_chains_dropped_invented_recall"] += len(correct) - len(keep) | |
| correct = keep | |
| if not correct: | |
| continue | |
| if rep_max: | |
| keep = [s for s in correct if repeated_line(s["text"]) < rep_max] | |
| counts["reasoner_chains_dropped_repetition"] += len(correct) - len(keep) | |
| correct = keep | |
| if not correct: | |
| continue | |
| chosen = correct[:2] if (second_chain_min and len(correct) >= max(2, second_chain_min)) else correct[:1] | |
| votes, n = vote_share(item, rec) | |
| p_ids = tok(render(tok, item, think=True, order=rec["order"]), add_special_tokens=False)["input_ids"] | |
| for s in chosen: | |
| text = CW_RE.sub(r".\1", s["text"]) if strip_cw else s["text"] | |
| counts["reasoner_cw_stripped"] += text != s["text"] | |
| c_ids = tok(text, add_special_tokens=False)["input_ids"] + [im_end] | |
| ex = {"kind": "reasoner", "id": item["id"], "input_ids": p_ids + c_ids, "n_prompt": len(p_ids)} | |
| if post_soft: | |
| pos = None | |
| if think_end in c_ids: | |
| t = len(c_ids) - 1 - c_ids[::-1].index(think_end) | |
| pos = next((j for j in range(t + 1, len(c_ids)) if c_ids[j] in lids[:k]), None) | |
| if pos is not None and lids.index(c_ids[pos]) == s["answer"]: | |
| own = np.zeros(k) | |
| own[s["answer"]] = 1.0 | |
| v = np.array([votes[rec["order"][i]] for i in range(k)]) if (votes is not None and n >= 2) else own | |
| ex.update(letter_pos=len(p_ids) + pos, letters=lids[:k], q_letter=[float(x) for x in 0.5 * own + 0.5 * v]) | |
| counts["reasoner_soft_letter"] += 1 | |
| else: | |
| counts["reasoner_letter_not_found"] += 1 | |
| out.append(ex) | |
| counts["reasoner_examples"] += 1 | |
| counts["reasoner_items"] += 1 | |
| return out | |
| def make_steps(examples: list[dict], micro_tokens: int, step_tokens: int, rng: random.Random) -> list[list[list[dict]]]: | |
| """Length-grouped micro-batches (padded tokens <= micro_tokens, or one long example alone), shuffled, then | |
| grouped into optimizer steps of about step_tokens real tokens.""" | |
| exs = sorted(examples, key=lambda e: len(e["input_ids"])) | |
| micro, cur, cur_max = [], [], 0 | |
| for e in exs: | |
| L = len(e["input_ids"]) | |
| if cur and max(cur_max, L) * (len(cur) + 1) > micro_tokens: | |
| micro.append(cur) | |
| cur, cur_max = [], 0 | |
| cur.append(e) | |
| cur_max = max(cur_max, L) | |
| if cur: | |
| micro.append(cur) | |
| rng.shuffle(micro) | |
| steps, cur, tok_count = [], [], 0 | |
| for mb in micro: | |
| cur.append(mb) | |
| tok_count += sum(len(e["input_ids"]) for e in mb) | |
| if tok_count >= step_tokens: | |
| steps.append(cur) | |
| cur, tok_count = [], 0 | |
| if cur: | |
| steps.append(cur) | |
| return steps | |
| # ------------------------------------------------------------------ loss | |
| def _chunk_ce(h: torch.Tensor, weight: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: | |
| logits = (h @ weight.t()).float() | |
| return F.cross_entropy(logits, labels, reduction="sum") | |
| def micro_loss(forward, lm_weight, mb: list[dict], pad_id: int, device, chunk: int = 2048, | |
| autocast: bool = True) -> tuple[torch.Tensor, dict]: | |
| """Sum over the micro-batch of each example's own loss (chain: mean token CE; readout: soft CE). | |
| forward(ids, attention_mask) returns the final normalised hidden states.""" | |
| L = max(len(e["input_ids"]) for e in mb) | |
| ids = torch.full((len(mb), L), pad_id, dtype=torch.long) | |
| att = torch.zeros((len(mb), L), dtype=torch.long) | |
| for i, e in enumerate(mb): # right padding | |
| n = len(e["input_ids"]) | |
| ids[i, :n] = torch.tensor(e["input_ids"]) | |
| att[i, :n] = 1 | |
| ids, att = ids.to(device), att.to(device) | |
| with torch.autocast(device.type, dtype=torch.bfloat16, enabled=autocast): | |
| hidden = forward(ids, att) | |
| total = hidden.new_zeros((), dtype=torch.float32) | |
| stats = collections.Counter() | |
| for i, e in enumerate(mb): | |
| n = len(e["input_ids"]) | |
| if e["kind"] == "readout": | |
| h = hidden[i, n - 1] | |
| z = (lm_weight[e["letters"]] @ h).float() | |
| q = torch.tensor(e["q"], device=device, dtype=torch.float32) | |
| logp = F.log_softmax(z, dim=-1) | |
| loss = -(q * logp).sum() | |
| stats["readout_loss"] += float(loss) | |
| stats["readout_n"] += 1 | |
| stats["readout_argmax_hits"] += int(int(z.argmax()) == int(q.argmax())) | |
| stats["readout_top_prob"] += float(logp.max().exp()) | |
| stats["readout_target_entropy"] += float(-(q * torch.log(q.clamp_min(1e-12))).sum()) | |
| else: | |
| s = e["n_prompt"] # predict tokens s..n-1 from positions s-1..n-2 | |
| hs = hidden[i, s - 1:n - 1] | |
| labels = ids[i, s:n].clone() | |
| if "letter_pos" in e: # the letter is scored by its own soft target below | |
| labels[e["letter_pos"] - s] = -100 | |
| n_tok = int((labels != -100).sum()) | |
| loss = hidden.new_zeros((), dtype=torch.float32) | |
| for a in range(0, hs.shape[0], chunk): | |
| loss = loss + checkpoint(_chunk_ce, hs[a:a + chunk], lm_weight, labels[a:a + chunk], use_reentrant=False) | |
| loss = loss / max(n_tok, 1) | |
| stats["chain_loss"] += float(loss) | |
| stats["chain_n"] += 1 | |
| stats["chain_tokens"] += int(hs.shape[0]) | |
| if "letter_pos" in e: | |
| z = (lm_weight[e["letters"]] @ hidden[i, e["letter_pos"] - 1]).float() | |
| ql = torch.tensor(e["q_letter"], device=device, dtype=torch.float32) | |
| lsoft = -(ql * F.log_softmax(z, dim=-1)).sum() | |
| loss = loss + lsoft | |
| stats["letter_loss"] += float(lsoft) | |
| stats["letter_n"] += 1 | |
| stats["letter_top_prob"] += float(F.softmax(z, dim=-1).max()) | |
| total = total + loss | |
| return total, stats | |
| # ------------------------------------------------------------------ model | |
| def hidden_forward(peft_model): | |
| """Final normalised hidden states of the LoRA model; the head is applied only where a loss needs it.""" | |
| decoder = peft_model.get_base_model().model | |
| def forward(ids, att): | |
| return decoder(input_ids=ids, attention_mask=att, use_cache=False, return_dict=True).last_hidden_state | |
| return forward | |
| def dev_readout(forward, lm_weight, tok, items, pad_id, device, micro_tokens) -> dict: | |
| """Readout accuracy, NLL and ECE (raw probabilities, dataset option order) on held-out training-source questions.""" | |
| from common import ece | |
| lids = letter_ids(tok) | |
| enc = sorted(((it, tok(render(tok, it, think=False), add_special_tokens=False)["input_ids"]) for it in items), key=lambda x: len(x[1])) | |
| ok, nl, conf = [], [], [] | |
| i = 0 | |
| while i < len(enc): | |
| j, L = i, 0 | |
| while j < len(enc) and max(L, len(enc[j][1])) * (j - i + 1) <= micro_tokens: | |
| L = max(L, len(enc[j][1])) | |
| j += 1 | |
| j = max(j, i + 1) | |
| batch, i = enc[i:j], j | |
| L = max(len(x) for _, x in batch) | |
| ids = torch.full((len(batch), L), pad_id, dtype=torch.long) | |
| att = torch.zeros((len(batch), L), dtype=torch.long) | |
| for b, (_, x) in enumerate(batch): | |
| ids[b, :len(x)] = torch.tensor(x) | |
| att[b, :len(x)] = 1 | |
| with torch.autocast(device.type, dtype=torch.bfloat16): | |
| hidden = forward(ids.to(device), att.to(device)) | |
| for b, (it, x) in enumerate(batch): | |
| lp = F.log_softmax((lm_weight[lids[:len(it["options"])]] @ hidden[b, len(x) - 1]).float(), dim=-1) | |
| ok.append(int(lp.argmax()) == it["gold"]) | |
| nl.append(-float(lp[it["gold"]])) | |
| conf.append(float(lp.max().exp())) | |
| return {"dev_n": len(ok), "dev_acc": float(np.mean(ok)), "dev_nll": float(np.mean(nl)), "dev_ece15": ece(np.array(conf), np.array(ok))} | |
| def set_checkpointing(model, on: bool) -> None: | |
| """Switch activation checkpointing of every decoder layer on or off (between micro-batches).""" | |
| for m in model.modules(): | |
| if hasattr(m, "gradient_checkpointing") and hasattr(m, "_gradient_checkpointing_func"): | |
| m.gradient_checkpointing = on | |
| def load_model(name: str, rank: int, lora_alpha: int, targets: list[str], device=None, liger: bool = False, | |
| dtype=torch.bfloat16): | |
| from peft import LoraConfig, get_peft_model | |
| model = AutoModelForCausalLM.from_pretrained(name, dtype=dtype, attn_implementation="sdpa").to(device or "cuda") | |
| if liger: # fused RMSNorm and SwiGLU kernels; the loss stays ours (only letter rows / chunks) | |
| from liger_kernel.transformers import apply_liger_kernel_to_qwen3_5 | |
| apply_liger_kernel_to_qwen3_5(rope=False, cross_entropy=False, fused_linear_cross_entropy=False, | |
| rms_norm=True, swiglu=True, model=model) | |
| present = {n.split(".")[-1] for n, m in model.named_modules() if isinstance(m, torch.nn.Linear)} | |
| targets = [t for t in targets if t in present] + [t for t in LINEAR_ATTN_TARGETS if t in present and t not in targets] | |
| model.config.use_cache = False | |
| model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) | |
| model.enable_input_require_grads() | |
| cfg = LoraConfig(r=rank, lora_alpha=lora_alpha, lora_dropout=0.0, target_modules=targets, bias="none", | |
| task_type="CAUSAL_LM") | |
| model = get_peft_model(model, cfg) | |
| return model | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--model", default="Qwen/Qwen3.5-4B") | |
| ap.add_argument("--pool", default="data/train_pool.jsonl") | |
| ap.add_argument("--gen", default="data/gen.jsonl") | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--mode", choices=["joint", "readout_only"], default="joint") | |
| ap.add_argument("--alpha", type=float, default=0.5) | |
| ap.add_argument("--lr", type=float, default=1e-4) | |
| ap.add_argument("--rank", type=int, default=32) | |
| ap.add_argument("--lora-alpha", type=int, default=64) | |
| ap.add_argument("--micro-tokens", type=int, default=16384) | |
| ap.add_argument("--step-tokens", type=int, default=32768) | |
| ap.add_argument("--budget-min", type=float, default=75.0, help="time box for the whole training stage") | |
| ap.add_argument("--deadline", type=float, default=None, help="unix time; stop taking steps after it") | |
| ap.add_argument("--seed", type=int, default=17) | |
| ap.add_argument("--max-steps", type=int, default=None, help="preflight only") | |
| ap.add_argument("--targets", default=",".join(LORA_TARGETS)) | |
| # V2 (defaults reproduce v1) | |
| ap.add_argument("--readout-pool", choices=["gen", "all"], default="gen", help="readout examples for generated items only, or every pool item") | |
| ap.add_argument("--subset", default=None, help="file of item ids (one per line) the readout examples are restricted to") | |
| ap.add_argument("--min-correct", type=int, default=1) | |
| ap.add_argument("--second-chain-min", type=int, default=0) | |
| ap.add_argument("--rep-max", type=int, default=0) | |
| ap.add_argument("--post-soft", action="store_true") | |
| ap.add_argument("--eval-order-frac", type=float, default=0.0) | |
| ap.add_argument("--drop-recall", action="store_true", help="drop chains that pretend to look the item up") | |
| ap.add_argument("--strip-cw", action="store_true", help="remove the stray '.cw' before </think>") | |
| ap.add_argument("--dev", default=None, help="held-out questions scored by the trained readout after training") | |
| ap.add_argument("--device", default="cuda") | |
| ap.add_argument("--ce-chunk", type=int, default=2048, help="rows of full-vocabulary logits per loss chunk") | |
| ap.add_argument("--liger", action="store_true", help="Liger RMSNorm + SwiGLU kernels") | |
| ap.add_argument("--ckpt-above", type=int, default=0, | |
| help="checkpoint activations only for micro-batches with more padded tokens than this (0: always)") | |
| ap.add_argument("--no-probe", action="store_true", help="skip the timing probe (short runs that surely fit)") | |
| args = ap.parse_args() | |
| t_start = time.time() | |
| os.makedirs(args.out, exist_ok=True) | |
| rng = random.Random(args.seed) | |
| torch.manual_seed(args.seed) | |
| tok = AutoTokenizer.from_pretrained(args.model) | |
| pool = {x["id"]: x for x in read_jsonl(args.pool)} | |
| gen = read_jsonl(args.gen) | |
| counts = collections.Counter() | |
| readout_ids = None | |
| if args.readout_pool == "all" or args.subset: | |
| readout_ids = list(pool) | |
| if args.subset: | |
| keep = {l.strip() for l in open(args.subset) if l.strip()} | |
| readout_ids = [i for i in readout_ids if i in keep] | |
| examples = build_examples(pool, gen, tok, args.mode, args.alpha, rng, counts, readout_ids=readout_ids, | |
| min_correct=args.min_correct, second_chain_min=args.second_chain_min, rep_max=args.rep_max, | |
| post_soft=args.post_soft, eval_order_frac=args.eval_order_frac, | |
| drop_recall=args.drop_recall, strip_cw=args.strip_cw) | |
| n_tokens = sum(len(e["input_ids"]) for e in examples) | |
| print(json.dumps({"examples": len(examples), "tokens": n_tokens, **counts}), flush=True) | |
| device = torch.device(args.device) | |
| model = load_model(args.model, args.rank, args.lora_alpha, args.targets.split(","), device=device, liger=args.liger) | |
| model.print_trainable_parameters() | |
| forward = hidden_forward(model) | |
| lm_weight = model.get_base_model().lm_head.weight # frozen; only letter rows / chunks are used | |
| pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.convert_tokens_to_ids("<|endoftext|>") | |
| def run(steps, opt, sched, deadline, log, max_steps=None): | |
| model.train() | |
| trainable = [p for p in model.parameters() if p.requires_grad] | |
| done_tokens, t0 = 0, time.time() | |
| for si, step in enumerate(steps): | |
| if (deadline and time.time() >= deadline) or (max_steps and si >= max_steps): | |
| return si, done_tokens, time.time() - t0, True | |
| n_ex = sum(len(mb) for mb in step) | |
| agg = collections.Counter() | |
| for mb in step: | |
| saved = None | |
| if args.ckpt_above: | |
| on = max(len(e["input_ids"]) for e in mb) * len(mb) > args.ckpt_above | |
| set_checkpointing(model, on) | |
| if not on: # kept so an out-of-memory micro-batch can be redone | |
| saved = [None if p.grad is None else p.grad.clone() for p in trainable] | |
| try: | |
| loss, st = micro_loss(forward, lm_weight, mb, pad_id, device, chunk=args.ce_chunk) | |
| (loss / n_ex).backward() | |
| except torch.cuda.OutOfMemoryError: | |
| if saved is None: | |
| raise | |
| loss = st = None | |
| if st is None: | |
| # out of memory without checkpointing: drop this micro-batch's partial gradient, redo it checkpointed | |
| for p, g in zip(trainable, saved): | |
| p.grad = g | |
| torch.cuda.empty_cache() | |
| counts["oom_retries"] += 1 | |
| print(json.dumps({"oom_retry": counts["oom_retries"], "step": si + 1, | |
| "microbatch_tokens": max(len(e["input_ids"]) for e in mb) * len(mb)}), flush=True) | |
| set_checkpointing(model, True) | |
| loss, st = micro_loss(forward, lm_weight, mb, pad_id, device, chunk=args.ce_chunk) | |
| (loss / n_ex).backward() | |
| agg.update(st) | |
| done_tokens += sum(len(e["input_ids"]) for e in mb) | |
| gnorm = torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0) | |
| opt.step() | |
| sched.step() | |
| opt.zero_grad(set_to_none=True) | |
| if log is not None: | |
| el = time.time() - t0 | |
| rn, cn = agg["readout_n"], agg["chain_n"] | |
| rec = {"step": si + 1, "of": len(steps), "lr": sched.get_last_lr()[0], "grad_norm": float(gnorm), | |
| "examples": n_ex, "readout_examples": rn, "reasoner_examples": cn, | |
| "step_tokens": sum(len(e["input_ids"]) for mb in step for e in mb), | |
| "chain_tokens": agg["chain_tokens"], "micro_batches": len(step), | |
| "tok_per_s": done_tokens / el, "elapsed_min": el / 60, | |
| "eta_min": el / (si + 1) * (len(steps) - si - 1) / 60, | |
| "chain_loss": agg["chain_loss"] / cn if cn else None, | |
| "readout_loss": agg["readout_loss"] / rn if rn else None, | |
| "readout_argmax_matches_target": agg["readout_argmax_hits"] / rn if rn else None, | |
| "readout_mean_top_prob": agg["readout_top_prob"] / rn if rn else None, | |
| "readout_target_entropy": agg["readout_target_entropy"] / rn if rn else None, | |
| "letter_loss": agg["letter_loss"] / agg["letter_n"] if agg["letter_n"] else None, | |
| "letter_top_prob": agg["letter_top_prob"] / agg["letter_n"] if agg["letter_n"] else None, | |
| "mem_gb": torch.cuda.max_memory_allocated() / 1e9 if device.type == "cuda" else None} | |
| log.write(json.dumps(rec) + "\n") | |
| log.flush() | |
| if (si + 1) % 10 == 0 or si < 3: | |
| print(json.dumps(rec), flush=True) | |
| return len(steps), done_tokens, time.time() - t0, False | |
| params = [p for p in model.parameters() if p.requires_grad] | |
| # timing probe on a throwaway copy of the schedule: 3 warm-up steps (kernel compilation, not timed), then 10 | |
| # timed steps; then reset the adapter | |
| steps = make_steps(examples, args.micro_tokens, args.step_tokens, random.Random(args.seed)) | |
| probe_n = min(10, len(steps) - 3) | |
| if args.max_steps is None and probe_n >= 3 and not args.no_probe: | |
| init = {n: p.detach().clone() for n, p in model.named_parameters() if p.requires_grad} | |
| opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=0.0) | |
| sched = get_cosine_schedule_with_warmup(opt, 1, 100) | |
| run(steps[:3], opt, sched, None, None) | |
| _, ptoks, psecs, _ = run(steps[3:3 + probe_n], opt, sched, None, None) | |
| speed = ptoks / psecs | |
| with torch.no_grad(): | |
| for n, p in model.named_parameters(): | |
| if p.requires_grad: | |
| p.copy_(init[n]) | |
| del opt, sched, init | |
| remaining = args.budget_min * 60 - (time.time() - t_start) - 120 # 2 min for saving | |
| fit_tokens = speed * remaining | |
| print(json.dumps({"probe_tok_per_s": speed, "fit_tokens": fit_tokens, "have_tokens": n_tokens}), flush=True) | |
| if fit_tokens < n_tokens: # drop whole examples at random, never shorten | |
| keep = fit_tokens / n_tokens | |
| before = len(examples) | |
| examples = [e for e in examples if rng.random() < keep] | |
| counts["dropped_for_time"] = before - len(examples) | |
| print(json.dumps({"kept_fraction": keep, "examples_after_drop": len(examples)}), flush=True) | |
| steps = make_steps(examples, args.micro_tokens, args.step_tokens, random.Random(args.seed + 1)) | |
| opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=0.0) | |
| total = len(steps) if args.max_steps is None else min(len(steps), args.max_steps) | |
| sched = get_cosine_schedule_with_warmup(opt, max(1, int(0.03 * total)), total) | |
| deadline = args.deadline or (t_start + args.budget_min * 60 - 120) | |
| with open(os.path.join(args.out, "train_log.jsonl"), "w") as log: | |
| n_done, toks, secs, stopped = run(steps, opt, sched, deadline, log, max_steps=args.max_steps) | |
| model.save_pretrained(args.out) | |
| tok.save_pretrained(args.out) | |
| dev = {} | |
| if args.dev: | |
| model.eval() | |
| dev = dev_readout(forward, lm_weight, tok, read_jsonl(args.dev), pad_id, device, args.micro_tokens) | |
| print(json.dumps(dev), flush=True) | |
| meta = {**dev, "v2_options": {k: getattr(args, k) for k in ("readout_pool", "subset", "min_correct", "second_chain_min", "rep_max", "post_soft", "eval_order_frac", "drop_recall", "strip_cw")}, | |
| "speed_options": {"liger": args.liger, "ckpt_above": args.ckpt_above, "micro_tokens": args.micro_tokens, | |
| "no_probe": args.no_probe}, | |
| "mode": args.mode, "alpha": args.alpha, "lr": args.lr, "rank": args.rank, "lora_alpha": args.lora_alpha, | |
| "seed": args.seed, "examples": len(examples), "steps_done": n_done, "steps_planned": len(steps), | |
| "stopped_at_deadline": stopped, "train_tokens": toks, "train_seconds": secs, | |
| "tok_per_s": toks / max(secs, 1e-9), "counts": dict(counts), | |
| "reasoner_examples": sum(e["kind"] == "reasoner" for e in examples), | |
| "readout_examples": sum(e["kind"] == "readout" for e in examples)} | |
| with open(os.path.join(args.out, "train_meta.json"), "w") as f: | |
| json.dump(meta, f, indent=1) | |
| print(json.dumps(meta), flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 27.5 kB
- Xet hash:
- 72eaa5b8d210535173b84ba20c0fc75c96a89e740fe7802da7344e370f1423df
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.