| |
| """T1.4 FP8 acceptance check on the exact generation config (DECISIONS.md S-20). |
| |
| Spec §6 criteria: teacher-forced KL(bf16 || FP8) <= 0.02 nats/token, and GSM8K-250 accuracy within |
| 1.5 points of bf16. "FP8" is the whole generation config in configs/gen/teacher_fp8.yaml — FP8 block |
| weights AND FP8 KV cache, router bf16, Triton MoE and attention — run on gen's own kernels. The bf16 |
| reference is configs/gen/teacher_bf16_ref.yaml (same kernels, bf16 weights and KV, part of the |
| weights in pinned CPU memory). |
| |
| KL is exact over the full 151,936-token vocabulary: the capture step records the final hidden state |
| at every position (the input of vLLM's compute_logits), and the kl step rebuilds both models' logits |
| through the shared, unquantised bf16 LM head in strict fp32. Triton attention reads K/V from the |
| (FP8) cache in prefill as in decode, so teacher-forced positions carry the same KV quantisation as |
| generation. |
| |
| Subcommands: prep-gsm8k, prep-math, prep-tf, generate, seqs-from, capture, kl, score-gsm8k, |
| analyze-trunc. Every vLLM step runs as its own process through the gen queue. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import os |
| import random |
| import urllib.request |
| from collections import Counter, defaultdict |
| from pathlib import Path |
|
|
| import numpy as np |
| import yaml |
|
|
| from lowbit.eval.gsm8k import MATH_SUFFIX, extract_answer, same_number |
| from lowbit.gen.loops import loop_score |
|
|
| REPO = Path(__file__).resolve().parent.parent |
| DATA = Path(os.environ.get("MT_DATA", "/path/to/data")) |
| EOS_DOC = 151643 |
| PROD_CFG = REPO / "configs" / "gen" / "teacher_fp8.yaml" |
| PARQUET = "https://huggingface.co/datasets/{ds}/resolve/refs%2Fconvert%2Fparquet/{cfg}/{split}/0000.parquet" |
| MARCO_AYA = {"en": "English", "zh": "Simplified Chinese", "ar": "Standard Arabic", "de": "German", "es": "Spanish", |
| "fr": "French", "ko": "Korean", "ja": "Japanese", "pt": "Portuguese", "tr": "Turkish", |
| "id": "Indonesian", "it": "Italian", "nl": "Dutch", "pl": "Polish", "ru": "Russian", |
| "vi": "Vietnamese", "th": "Thai", "uk": "Ukrainian", "ms": "Standard Malay", "bn": "Bengali", |
| "ur": "Urdu", "el": "Greek", "hu": "Hungarian", "ne": "Nepali"} |
|
|
|
|
| def fetch(url: str, dest: Path) -> str: |
| dest.parent.mkdir(parents=True, exist_ok=True) |
| if not dest.exists(): |
| urllib.request.urlretrieve(url, dest) |
| return hashlib.sha256(dest.read_bytes()).hexdigest() |
|
|
|
|
| def write_jsonl(path: Path, rows) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with path.open("w") as f: |
| for r in rows: |
| f.write(json.dumps(r, ensure_ascii=False) + "\n") |
|
|
|
|
| def read_jsonl(path) -> list[dict]: |
| return [json.loads(l) for l in open(path) if l.strip()] |
|
|
|
|
| def load_tokenizer(model: str): |
| from transformers import AutoTokenizer |
| return AutoTokenizer.from_pretrained(model) |
|
|
|
|
| |
|
|
|
|
| def prep_gsm8k(a) -> None: |
| import pyarrow.parquet as pq |
| dest = DATA / "evalsets/gsm8k/main_test.parquet" |
| sha = fetch(PARQUET.format(ds="openai/gsm8k", cfg="main", split="test"), dest) |
| rows = pq.read_table(dest).to_pylist() |
| ids = sorted(random.Random(a.seed).sample(range(len(rows)), 250)) |
| subset = {"dataset": "openai/gsm8k", "config": "main", "split": "test", "dataset_sha": a.dataset_sha, |
| "parquet_sha256": sha, "seed": a.seed, "n": 250, "ids": ids} |
| out = REPO / "configs/eval/subsets/gsm8k_250.json" |
| out.parent.mkdir(parents=True, exist_ok=True) |
| out.write_text(json.dumps(subset, indent=1) + "\n") |
| write_jsonl(DATA / "fp8check/gsm8k_250.jsonl", |
| [{"id": f"gsm8k-{i}", "text": rows[i]["question"] + MATH_SUFFIX, |
| "gold": rows[i]["answer"].split("####")[-1].strip().replace(",", ""), "kind": "gsm8k"} |
| for i in ids]) |
| print(f"gsm8k: {len(rows)} test items, 250 sampled (seed {a.seed}) -> {out}") |
|
|
|
|
| def prep_math(a) -> None: |
| import pyarrow.parquet as pq |
| subjects = ["algebra", "counting_and_probability", "geometry", "intermediate_algebra", "number_theory", |
| "prealgebra", "precalculus"] |
| pool = [] |
| for s in subjects: |
| dest = DATA / f"evalsets/hendrycks_math/{s}_train.parquet" |
| fetch(PARQUET.format(ds="EleutherAI/hendrycks_math", cfg=s, split="train"), dest) |
| pool += [dict(r, subject=s) for r in pq.read_table(dest).to_pylist() if r["level"] in ("Level 4", "Level 5")] |
| pick = random.Random(a.seed).sample(pool, a.n) |
| write_jsonl(DATA / "fp8check/math_long.jsonl", |
| [{"id": f"math-{i}", "text": r["problem"] + MATH_SUFFIX, "kind": "math", "level": r["level"], |
| "subject": r["subject"]} for i, r in enumerate(pick)]) |
| print(f"math: {len(pool)} level-4/5 train problems, {a.n} sampled") |
|
|
|
|
| def prep_tf(a) -> None: |
| """Aya human-written prompt+completion text in Marco's languages, packed into fixed windows.""" |
| import pyarrow.parquet as pq |
| tok = load_tokenizer(a.model) |
| t = pq.read_table(DATA / "bench/aya_train.parquet", columns=["inputs", "targets", "language"]).to_pylist() |
| rng = random.Random(a.seed) |
| by_lang = defaultdict(list) |
| names = {v: k for k, v in MARCO_AYA.items()} |
| for r in t: |
| if r["language"] in names: |
| by_lang[names[r["language"]]].append(r["inputs"].strip() + "\n\n" + r["targets"].strip()) |
| per_lang = a.tokens // len(by_lang) |
| docs = [] |
| for lang, texts in sorted(by_lang.items()): |
| rng.shuffle(texts) |
| n = 0 |
| for txt in texts: |
| ids = tok(txt, add_special_tokens=False)["input_ids"] |
| docs.append((lang, ids)) |
| n += len(ids) + 1 |
| if n >= per_lang: |
| break |
| rng.shuffle(docs) |
| windows, cur, segs = [], [], [] |
| for lang, ids in docs: |
| ids = ids + [EOS_DOC] |
| while ids: |
| room = a.window - len(cur) |
| segs.append([len(cur), lang]) |
| cur += ids[:room] |
| ids = ids[room:] |
| if len(cur) == a.window: |
| windows.append({"id": f"aya-{len(windows)}", "kind": "aya", "token_ids": cur, "segments": segs}) |
| cur, segs = [], [] |
| write_jsonl(DATA / "fp8check/tf_aya.jsonl", windows) |
| langs = Counter() |
| for w in windows: |
| bounds = [s[0] for s in w["segments"]] + [a.window] |
| for (start, lang), end in zip(w["segments"], bounds[1:]): |
| langs[lang] += end - start |
| print(f"prep-tf: {len(windows)} windows x {a.window} = {len(windows) * a.window:,} tokens; per language: {dict(langs)}") |
|
|
|
|
| |
|
|
|
|
| def build_llm(cfg_path: str, overrides: list[str], in_process: bool): |
| cfg = yaml.safe_load(open(cfg_path)) |
| os.environ.update({k: str(v) for k, v in cfg.get("env", {}).items()}) |
| if in_process: |
| os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" |
| from vllm import LLM |
| eng = dict(cfg["engine"]) |
| for kv in overrides: |
| k, v = kv.split("=", 1) |
| eng[k] = yaml.safe_load(v) |
| return LLM(model=cfg["model"], seed=0, **eng), cfg |
|
|
|
|
| def spec_sampling(path: str | Path = PROD_CFG) -> dict: |
| """Spec §6 sampling, always from the production generation config: an FP8-vs-bf16 comparison |
| must sample both engines identically, and the bf16 reference config carries no sampling block.""" |
| return yaml.safe_load(open(path))["sampling"] |
|
|
|
|
| def generate(a) -> None: |
| llm, _ = build_llm(a.engine, a.override, in_process=False) |
| from vllm import SamplingParams |
| from vllm.inputs import TokensPrompt |
| tok = llm.get_tokenizer() |
| rows = read_jsonl(a.prompts)[: a.limit or None] |
| prompts = [TokensPrompt(prompt_token_ids=tok(tok.apply_chat_template([{"role": "user", "content": r["text"]}], |
| add_generation_prompt=True, tokenize=False), |
| add_special_tokens=False)["input_ids"]) for r in rows] |
| if a.sampling == "greedy": |
| sp = SamplingParams(temperature=0.0, max_tokens=a.max_tokens) |
| else: |
| s = spec_sampling(a.sampling_config) |
| sp = SamplingParams(temperature=s["temperature"], top_p=s["top_p"], max_tokens=a.max_tokens, seed=0) |
| outs = llm.generate(prompts, sp) |
| write_jsonl(Path(a.out), [{"id": r.get("id", f"p{i}"), **{k: v for k, v in r.items() if k != "text"}, |
| "prompt_token_ids": o.prompt_token_ids, |
| "output_token_ids": list(o.outputs[0].token_ids), "text": o.outputs[0].text, |
| "finish_reason": o.outputs[0].finish_reason} for i, (r, o) in enumerate(zip(rows, outs))]) |
| n_tok = sum(len(o.outputs[0].token_ids) for o in outs) |
| print(json.dumps({"n": len(outs), "out_tokens": n_tok, |
| "truncated": sum(o.outputs[0].finish_reason == "length" for o in outs)})) |
|
|
|
|
| def seqs_from(a) -> None: |
| """Teacher-forcing sequences from generate outputs: chat prompt + response, response start marked.""" |
| rows = read_jsonl(a.generated) |
| write_jsonl(Path(a.out), [{"id": r["id"], "kind": r.get("kind", "gen"), |
| "token_ids": r["prompt_token_ids"] + r["output_token_ids"], |
| "response_start": len(r["prompt_token_ids"]), "finish_reason": r["finish_reason"]} |
| for r in rows if r["output_token_ids"]]) |
|
|
|
|
| def install_capture(store: dict): |
| """Wrap the V2 runner's prompt-logprob step to keep each prefill row's final hidden state.""" |
| from vllm.v1.worker.gpu.sample import prompt_logprob as pl |
| orig = pl.PromptLogprobsWorker.compute_prompt_logprobs |
|
|
| def wrapped(self, logits_fn, hidden_states, input_batch, all_token_ids, num_computed_tokens, prompt_lens): |
| plens = prompt_lens[input_batch.idx_mapping_np] |
| computed = input_batch.num_computed_prefill_tokens_np |
| qsl = input_batch.query_start_loc_np |
| for i, rid in enumerate(input_batch.req_ids): |
| if computed[i] >= plens[i]: |
| continue |
| s, e = int(qsl[i]), int(qsl[i + 1]) |
| store[rid].append((int(computed[i]), hidden_states[s:e].detach().to("cpu", copy=True))) |
| return orig(self, logits_fn, hidden_states, input_batch, all_token_ids, num_computed_tokens, prompt_lens) |
|
|
| pl.PromptLogprobsWorker.compute_prompt_logprobs = wrapped |
|
|
|
|
| def capture(a) -> None: |
| store: dict[str, list] = defaultdict(list) |
| llm, _ = build_llm(a.engine, a.override, in_process=True) |
| install_capture(store) |
| import torch |
| from vllm import SamplingParams |
| from vllm.inputs import TokensPrompt |
| seqs = [] |
| for f in a.sequences: |
| seqs += read_jsonl(f) |
| seqs = seqs[: a.limit or None] |
| sp = SamplingParams(temperature=0.0, max_tokens=1, prompt_logprobs=0) |
| outs = llm.generate([TokensPrompt(prompt_token_ids=s["token_ids"]) for s in seqs], sp) |
| out_dir = Path(a.out) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| total = sum(len(s["token_ids"]) for s in seqs) |
| hid = np.lib.format.open_memmap(out_dir / "hidden_bf16.npy", mode="w+", dtype=np.uint16, shape=(total, 1024)) |
| vllm_lp = np.full(total, np.nan, dtype=np.float32) |
| index, off = [], 0 |
| for s, o in zip(seqs, outs): |
| rid = o.request_id |
| chunks = store.get(rid) or next((v for k, v in store.items() if k.startswith(rid + "-") or rid.startswith(k)), None) |
| if chunks is None: |
| raise RuntimeError(f"no captured hidden states for request {rid} ({s['id']}); ids seen: {list(store)[:5]}") |
| chunks.sort(key=lambda c: c[0]) |
| h = torch.cat([c[1] for c in chunks]) |
| n = len(s["token_ids"]) |
| starts = [c[0] for c in chunks] |
| if h.shape[0] != n or starts[0] != 0: |
| raise RuntimeError(f"{s['id']}: captured {h.shape[0]} rows from {starts}, expected {n}") |
| hid[off:off + n] = h.to(torch.bfloat16).view(torch.int16).numpy().view(np.uint16) |
| for p, d in enumerate(o.prompt_logprobs or []): |
| if d is not None: |
| vllm_lp[off + p] = d[s["token_ids"][p]].logprob |
| index.append({k: v for k, v in s.items() if k != "token_ids"} | {"offset": off, "length": n}) |
| off += n |
| hid.flush() |
| np.save(out_dir / "vllm_token_logprob.npy", vllm_lp) |
| write_jsonl(out_dir / "index.jsonl", index) |
| (out_dir / "meta.json").write_text(json.dumps({"engine": a.engine, "override": a.override, "positions": total, |
| "sequences": len(seqs)}, indent=1)) |
| print(json.dumps({"captured_positions": total, "sequences": len(seqs)})) |
|
|
|
|
| def kl(a) -> None: |
| import torch |
| from safetensors import safe_open |
| torch.backends.cuda.matmul.allow_tf32 = False |
| torch.set_float32_matmul_precision("highest") |
| emb = None |
| for f in sorted(Path(a.model).glob("*.safetensors")): |
| with safe_open(f, "pt") as st: |
| if "model.embed_tokens.weight" in st.keys(): |
| emb = st.get_tensor("model.embed_tokens.weight") |
| dev = a.device |
| E = emb.to(dev, torch.float32) |
| ref_dir, test_dir = Path(a.ref), Path(a.test) |
| ref_idx = {r["id"]: r for r in read_jsonl(ref_dir / "index.jsonl")} |
| test_idx = {r["id"]: r for r in read_jsonl(test_dir / "index.jsonl")} |
| seqs = {} |
| for f in a.sequences: |
| seqs.update({s["id"]: s for s in read_jsonl(f)}) |
| H_ref = np.load(ref_dir / "hidden_bf16.npy", mmap_mode="r") |
| H_test = np.load(test_dir / "hidden_bf16.npy", mmap_mode="r") |
| lp_ref_vllm = np.load(ref_dir / "vllm_token_logprob.npy") |
| cols = defaultdict(list) |
| for sid, r in ref_idx.items(): |
| if sid not in test_idx: |
| continue |
| t, s = test_idx[sid], seqs[sid] |
| n = r["length"] |
| assert t["length"] == n and len(s["token_ids"]) == n, sid |
| hb = torch.from_numpy(np.array(H_ref[r["offset"]:r["offset"] + n], copy=True).view(np.int16)).view(torch.bfloat16) |
| hf = torch.from_numpy(np.array(H_test[t["offset"]:t["offset"] + n], copy=True).view(np.int16)).view(torch.bfloat16) |
| lang = np.full(n, "", dtype=object) |
| if "segments" in s: |
| bounds = [x[0] for x in s["segments"]] + [n] |
| for (st_, lg), en in zip(s["segments"], bounds[1:]): |
| lang[st_:en] = lg |
| toks = torch.tensor(s["token_ids"], device=dev) |
| parts = defaultdict(list) |
| for c in range(0, n, a.chunk): |
| e = min(n, c + a.chunk) |
| lb = torch.log_softmax(hb[c:e].to(dev, torch.float32) @ E.T, -1) |
| lf = torch.log_softmax(hf[c:e].to(dev, torch.float32) @ E.T, -1) |
| parts["kl"].append((lb.exp() * (lb - lf)).sum(-1).cpu()) |
| parts["rkl"].append((lf.exp() * (lf - lb)).sum(-1).cpu()) |
| parts["agree"].append((lb.argmax(-1) == lf.argmax(-1)).cpu()) |
| nxt = torch.arange(c + 1, e + 1, device=dev).clamp(max=n - 1) |
| parts["our_lp"].append(lb.gather(1, toks[nxt].unsqueeze(1)).squeeze(1).cpu()) |
| our = torch.cat(parts["our_lp"]).numpy().astype(np.float64) |
| our[-1] = np.nan |
| vl = np.full(n, np.nan) |
| vl[:-1] = lp_ref_vllm[r["offset"] + 1:r["offset"] + n] |
| pos = np.arange(n) |
| cols["kind"].append(np.full(n, s.get("kind", ""), dtype=object)) |
| cols["lang"].append(lang) |
| cols["pos"].append(pos) |
| cols["is_resp"].append(pos >= s.get("response_start", n + 1)) |
| cols["kl"].append(torch.cat(parts["kl"]).numpy().astype(np.float64)) |
| cols["rkl"].append(torch.cat(parts["rkl"]).numpy().astype(np.float64)) |
| cols["agree"].append(torch.cat(parts["agree"]).numpy()) |
| cols["our_lp"].append(our) |
| cols["vl"].append(vl) |
| kinds, langs, pos, is_resp = (np.concatenate(cols[k]) for k in ("kind", "lang", "pos", "is_resp")) |
| klv, rkl, agree, our_lp, vl = (np.concatenate(cols[k]) for k in ("kl", "rkl", "agree", "our_lp", "vl")) |
| out_dir = Path(a.out) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(out_dir / "per_position.npz", kind=kinds.astype(str), lang=langs.astype(str), pos=pos, |
| is_response=is_resp, kl=klv, rkl=rkl, top1_agree=agree) |
|
|
| def stats(mask): |
| m = mask & np.isfinite(klv) |
| if not m.any(): |
| return None |
| return {"positions": int(m.sum()), "kl_mean": float(klv[m].mean()), "kl_p50": float(np.median(klv[m])), |
| "kl_p99": float(np.quantile(klv[m], 0.99)), "kl_max": float(klv[m].max()), |
| "reverse_kl_mean": float(rkl[m].mean()), "top1_agreement": float(agree[m].mean())} |
| every = np.ones_like(klv, dtype=bool) |
| summary = {"ref": str(ref_dir), "test": str(test_dir), "threshold_nats_per_token": 0.02, |
| "all": stats(every), "by_kind": {k: stats(kinds == k) for k in sorted(set(kinds))}, |
| "by_position_bucket": {f"{b * 512}-{b * 512 + 511}": stats((pos // 512) == b) |
| for b in range(int(pos.max()) // 512 + 1)}, |
| "by_language": {l: stats(langs == l) for l in sorted(set(langs)) if l}, |
| "math_response_by_bucket": {f"{b * 512}-{b * 512 + 511}": stats((kinds == "math") & is_resp & ((pos // 512) == b)) |
| for b in range(int(pos.max()) // 512 + 1)}} |
| ok = np.isfinite(our_lp) & np.isfinite(vl) |
| summary["lm_head_reconstruction"] = {"positions": int(ok.sum()), "max_abs_diff_vs_vllm_logprob": float(np.abs(our_lp[ok] - vl[ok]).max()), |
| "mean_abs_diff": float(np.abs(our_lp[ok] - vl[ok]).mean())} |
| (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) |
| print(json.dumps({k: summary[k] for k in ("all", "by_kind", "lm_head_reconstruction")}, indent=1)) |
|
|
|
|
| |
|
|
|
|
| def score_gsm8k(a) -> None: |
| gold = {r["id"]: r["gold"] for r in read_jsonl(DATA / "fp8check/gsm8k_250.jsonl")} |
| runs = {} |
| for spec in a.pred: |
| name, path = spec.split("=", 1) |
| runs[name] = {r["id"]: r for r in read_jsonl(path)} |
| res = {} |
| correct = {} |
| for name, rows in runs.items(): |
| c = {i: same_number(extract_answer(r["text"]), gold[i]) for i, r in rows.items()} |
| correct[name] = c |
| res[name] = {"n": len(c), "accuracy": 100 * sum(c.values()) / len(c), |
| "truncated": sum(r["finish_reason"] == "length" for r in rows.values()), |
| "mean_out_tokens": float(np.mean([len(r["output_token_ids"]) for r in rows.values()]))} |
| if len(runs) == 2: |
| (na, ca), (nb, cb) = correct.items() |
| ids = sorted(set(ca) & set(cb)) |
| res["paired"] = {"a": na, "b": nb, "diff_points": res[na]["accuracy"] - res[nb]["accuracy"], |
| "both_right": sum(ca[i] and cb[i] for i in ids), "only_a": sum(ca[i] and not cb[i] for i in ids), |
| "only_b": sum(cb[i] and not ca[i] for i in ids), "both_wrong": sum(not ca[i] and not cb[i] for i in ids), |
| "identical_text": sum(runs[na][i]["text"] == runs[nb][i]["text"] for i in ids), |
| "threshold_points": 1.5} |
| res["paired"]["pass"] = abs(res["paired"]["diff_points"]) <= 1.5 |
| Path(a.out).write_text(json.dumps(res, indent=1)) |
| print(json.dumps(res, indent=1)) |
|
|
|
|
| def analyze_trunc(a) -> None: |
| rows = read_jsonl(a.generated) |
| trunc = [r for r in rows if r["finish_reason"] == "length"] |
| scores = [loop_score(r["output_token_ids"]) for r in trunc] |
| loops = sum(s >= 0.5 for s in scores) |
| by_lang = Counter(r.get("language", "") for r in trunc) |
| tot_lang = Counter(r.get("language", "") for r in rows) |
| out = {"n": len(rows), "truncated": len(trunc), "truncated_share": len(trunc) / max(1, len(rows)), |
| "truncated_loops(score>=0.5)": loops, "truncated_non_loops": len(trunc) - loops, |
| "loop_score_quartiles": [float(q) for q in np.quantile(scores, [0.25, 0.5, 0.75])] if scores else None, |
| "truncated_by_language": {k: f"{by_lang[k]}/{tot_lang[k]}" for k in sorted(tot_lang)}, |
| "output_tokens_total": sum(len(r["output_token_ids"]) for r in rows), |
| "output_tokens_in_truncated": sum(len(r["output_token_ids"]) for r in trunc)} |
| examples = [{"id": r["id"], "language": r.get("language"), "loop_score": round(s, 3), "tail": r["text"][-300:]} |
| for r, s in sorted(zip(trunc, scores), key=lambda x: x[1])][:: max(1, len(trunc) // 8)] |
| Path(a.out).write_text(json.dumps({"summary": out, "examples": examples}, indent=1, ensure_ascii=False)) |
| print(json.dumps(out, indent=1)) |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser() |
| sub = p.add_subparsers(dest="op", required=True) |
| s = sub.add_parser("prep-gsm8k"); s.add_argument("--seed", type=int, default=0); s.add_argument("--dataset-sha", default="740312add88f781978c0658806c59bc2815b9866") |
| s = sub.add_parser("prep-math"); s.add_argument("--seed", type=int, default=0); s.add_argument("--n", type=int, default=200) |
| s = sub.add_parser("prep-tf"); s.add_argument("--model", required=True); s.add_argument("--tokens", type=int, default=1_000_000) |
| s.add_argument("--window", type=int, default=4000); s.add_argument("--seed", type=int, default=0) |
| for name in ("generate", "capture"): |
| s = sub.add_parser(name) |
| s.add_argument("--engine", required=True) |
| s.add_argument("--override", nargs="*", default=[], help="engine KEY=VAL (capacity knobs; diagnostic arms)") |
| s.add_argument("--limit", type=int, default=0) |
| s.add_argument("--out", required=True) |
| if name == "generate": |
| s.add_argument("--prompts", required=True) |
| s.add_argument("--sampling", choices=["greedy", "spec"], required=True) |
| s.add_argument("--max-tokens", type=int, required=True) |
| s.add_argument("--sampling-config", default=str(PROD_CFG), help="config whose sampling block --sampling spec uses") |
| else: |
| s.add_argument("--sequences", nargs="+", required=True) |
| s = sub.add_parser("seqs-from"); s.add_argument("--generated", required=True); s.add_argument("--out", required=True) |
| s = sub.add_parser("kl"); s.add_argument("--ref", required=True); s.add_argument("--test", required=True) |
| s.add_argument("--sequences", nargs="+", required=True); s.add_argument("--model", required=True) |
| s.add_argument("--chunk", type=int, default=512); s.add_argument("--out", required=True); s.add_argument("--device", default="cuda") |
| s = sub.add_parser("score-gsm8k"); s.add_argument("--pred", nargs="+", required=True, help="name=path"); s.add_argument("--out", required=True) |
| s = sub.add_parser("analyze-trunc"); s.add_argument("--generated", required=True); s.add_argument("--out", required=True) |
| a = p.parse_args() |
| {"prep-gsm8k": prep_gsm8k, "prep-math": prep_math, "prep-tf": prep_tf, "generate": generate, |
| "seqs-from": seqs_from, "capture": capture, "kl": kl, "score-gsm8k": score_gsm8k, |
| "analyze-trunc": analyze_trunc}[a.op](a) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
