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 | |
| def build_examples(pool: dict, gen: list[dict], tok, mode: str, alpha: float, rng: random.Random, | |
| counts: collections.Counter) -> list[dict]: | |
| im_end = tok.convert_tokens_to_ids("<|im_end|>") | |
| lids = letter_ids(tok) | |
| out = [] | |
| for rec in gen: | |
| item = pool.get(rec["id"]) | |
| if item is None: | |
| continue | |
| k = len(item["options"]) | |
| gold = np.asarray(item["target"], dtype=np.float64) | |
| # readout example | |
| votes, n = vote_share(item, rec) | |
| 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_fallback"] += 1 | |
| 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": | |
| continue | |
| # reasoner example: shortest correct chain that closed on its own | |
| best = None | |
| for s in rec["samples"]: | |
| if s["closed"] and s["answer"] is not None and rec["order"][s["answer"]] == item["gold"]: | |
| if best is None or s["tokens"] < best["tokens"]: | |
| best = s | |
| if best is None: | |
| counts["reasoner_no_correct_chain"] += 1 | |
| continue | |
| prompt = render(tok, item, think=True, order=rec["order"]) | |
| p_ids = tok(prompt, add_special_tokens=False)["input_ids"] | |
| c_ids = tok(best["text"], add_special_tokens=False)["input_ids"] + [im_end] | |
| out.append({"kind": "reasoner", "id": item["id"], "input_ids": p_ids + c_ids, "n_prompt": len(p_ids)}) | |
| counts["reasoner_examples"] += 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) -> 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("cuda", dtype=torch.bfloat16): | |
| 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] | |
| 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 / hs.shape[0] | |
| stats["chain_loss"] += float(loss) | |
| stats["chain_n"] += 1 | |
| stats["chain_tokens"] += int(hs.shape[0]) | |
| 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 load_model(name: str, rank: int, lora_alpha: int, targets: list[str]): | |
| from peft import LoraConfig, get_peft_model | |
| model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.bfloat16, attn_implementation="sdpa").cuda() | |
| 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)) | |
| 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() | |
| examples = build_examples(pool, gen, tok, args.mode, args.alpha, rng, counts) | |
| n_tokens = sum(len(e["input_ids"]) for e in examples) | |
| print(json.dumps({"examples": len(examples), "tokens": n_tokens, **counts}), flush=True) | |
| model = load_model(args.model, args.rank, args.lora_alpha, args.targets.split(",")) | |
| model.print_trainable_parameters() | |
| forward = hidden_forward(model) | |
| lm_weight = model.get_base_model().lm_head.weight # frozen; only letter rows / chunks are used | |
| device = torch.device("cuda") | |
| 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() | |
| 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: | |
| loss, st = micro_loss(forward, lm_weight, mb, pad_id, device) | |
| (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, | |
| "mem_gb": torch.cuda.max_memory_allocated() / 1e9} | |
| 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] | |
| # 10-step timing probe on a throwaway copy of the schedule, then reset the adapter | |
| steps = make_steps(examples, args.micro_tokens, args.step_tokens, random.Random(args.seed)) | |
| probe_n = min(10, len(steps)) | |
| if args.max_steps is None and probe_n >= 3: | |
| 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) | |
| _, ptoks, psecs, _ = run(steps, opt, sched, None, None, max_steps=probe_n) | |
| 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) | |
| meta = {"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:
- 15.6 kB
- Xet hash:
- 1f4c876070bbe303415ca6a27d80b474ec31d09d92e2a7c7244c9b55493fa20e
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.