Buckets:
| """Training-speed benchmark on the real model and real V2 examples: same examples, same tokens, several setups, | |
| each checked against the current one before its speed counts. | |
| A current: every micro-batch checkpointed, 16K-token micro-batches | |
| B A + Liger RMSNorm and SwiGLU kernels | |
| C no checkpointing for micro-batches up to 8K padded tokens (8K micro-batches; a longer example goes alone, | |
| checkpointed) | |
| D B + C | |
| E D with 12K micro-batches, tried only when D passes with room to spare | |
| Equivalence, on fixed micro-batches (readout, chains, one long chain) with a seeded non-zero adapter, against R, the | |
| same computation in fp32 without autocast: a setup passes when its loss and LoRA-gradient errors relative to R are no | |
| more than 1.5x those of A (the current bf16 setup) plus a small floor. bf16 rounding sets A's error; a kernel that | |
| only reorders the rounding lands near it, a wrong kernel far above it. The fastest passing setup under 72 GB peak | |
| memory wins if it beats A by at least 3%; otherwise A stays. Writes results/speed_bench.json with the train.py flags. | |
| python bench_speed.py --model Qwen/Qwen3.5-4B | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import collections | |
| import json | |
| import os | |
| import pickle | |
| import random | |
| import subprocess | |
| import sys | |
| import time | |
| CONFIGS = {"R": {"liger": False, "micro": 16384, "ckpt_above": 0}, | |
| "A": {"liger": False, "micro": 16384, "ckpt_above": 0}, | |
| "B": {"liger": True, "micro": 16384, "ckpt_above": 0}, | |
| "C": {"liger": False, "micro": 8192, "ckpt_above": 8192}, | |
| "D": {"liger": True, "micro": 8192, "ckpt_above": 8192}, | |
| "E": {"liger": True, "micro": 12288, "ckpt_above": 12288}} | |
| STEP_TOKENS = 32768 | |
| CE_CHUNK = 2048 # train.py's default | |
| TIMED_TOKENS, WARM_TOKENS = 260_000, 50_000 | |
| def flags(c: dict) -> str: | |
| return " ".join(f for f in ["--liger" if c["liger"] else "", | |
| f"--micro-tokens {c['micro']}" if c["micro"] != 16384 else "", | |
| f"--ckpt-above {c['ckpt_above']}" if c["ckpt_above"] else ""] if f) | |
| def prepare(args) -> None: | |
| from transformers import AutoTokenizer | |
| from common import read_jsonl | |
| from train import build_examples | |
| tok = AutoTokenizer.from_pretrained(args.model) | |
| pool = {x["id"]: x for x in read_jsonl(args.pool)} | |
| gen = read_jsonl(args.gen) | |
| rng = random.Random(5) | |
| ids = list(pool) | |
| rng.shuffle(ids) | |
| gen_s = rng.sample(gen, min(600, len(gen))) # about the full run's readout : chain-item ratio | |
| counts = collections.Counter() | |
| ex = build_examples(pool, gen_s, tok, "joint", 0.5, random.Random(6), counts, readout_ids=ids[:2400], | |
| min_correct=2, second_chain_min=3, rep_max=20, post_soft=True, eval_order_frac=0.5, | |
| drop_recall=True, strip_cw=True) | |
| random.Random(7).shuffle(ex) | |
| timed, warm, n = [], [], 0 | |
| for e in ex: | |
| L = len(e["input_ids"]) | |
| if L > 16000 * args.scale: | |
| continue | |
| if n < TIMED_TOKENS: | |
| timed.append(e) | |
| elif n < TIMED_TOKENS + WARM_TOKENS: | |
| warm.append(e) | |
| else: | |
| break | |
| n += L | |
| ro = [e for e in ex if e["kind"] == "readout"] | |
| ch = sorted((e for e in ex if e["kind"] == "reasoner"), key=lambda e: len(e["input_ids"])) | |
| ch = [e for e in ch if len(e["input_ids"]) < 3500 * args.scale] or ch[:1] | |
| long_ = [e for e in ex if e["kind"] == "reasoner" and 8192 * args.scale < len(e["input_ids"]) <= 16000 * args.scale] | |
| check = [ro[:6], ch[-2:]] + ([long_[:1]] if long_ else []) | |
| with open(args.examples, "wb") as f: | |
| pickle.dump({"timed": timed, "warm": warm, "check": check}, f) | |
| print(json.dumps({"timed_examples": len(timed), "timed_tokens": sum(len(e["input_ids"]) for e in timed), | |
| "timed_reasoner": sum(e["kind"] == "reasoner" for e in timed), "warm_examples": len(warm), | |
| "check_microbatch_tokens": [max(len(e["input_ids"]) for e in mb) * len(mb) for mb in check]}), flush=True) | |
| def worker(args) -> None: | |
| import torch | |
| from transformers import AutoTokenizer | |
| from train import hidden_forward, load_model, make_steps, micro_loss, set_checkpointing | |
| c = CONFIGS[args.worker] | |
| data = pickle.load(open(args.examples, "rb")) | |
| tok = AutoTokenizer.from_pretrained(args.model) | |
| pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.convert_tokens_to_ids("<|endoftext|>") | |
| device = torch.device("cuda") | |
| torch.manual_seed(17) | |
| ref = args.worker == "R" | |
| model = load_model(args.model, 32, 64, ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], | |
| device=device, liger=c["liger"], dtype=torch.float32 if ref else torch.bfloat16) | |
| forward = hidden_forward(model) | |
| lm_weight = model.get_base_model().lm_head.weight | |
| named = sorted((n, p) for n, p in model.named_parameters() if p.requires_grad) | |
| with torch.no_grad(): # the same non-zero adapter in every setup | |
| for i, (n, p) in enumerate(named): | |
| g = torch.Generator().manual_seed(i) | |
| p.copy_(torch.randn(p.shape, generator=g) * (0.02 if "lora_A" in n else 0.002)) | |
| model.train() | |
| def ckpt(mb): | |
| if c["ckpt_above"]: | |
| set_checkpointing(model, max(len(e["input_ids"]) for e in mb) * len(mb) > c["ckpt_above"]) | |
| losses = [] | |
| for mb in data["check"]: | |
| ckpt(mb) | |
| loss, _ = micro_loss(forward, lm_weight, mb, pad_id, device, chunk=CE_CHUNK, autocast=not ref) | |
| loss.backward() | |
| losses.append(float(loss)) | |
| grad = torch.cat([p.grad.detach().float().flatten().cpu() for _, p in named]) | |
| torch.save(grad, os.path.join(os.path.dirname(args.examples), f"bench_grad_{args.worker}.pt")) | |
| if ref: | |
| print("RESULT " + json.dumps({"config": "R", "check_losses": losses}), flush=True) | |
| return | |
| for _, p in named: | |
| p.grad = None | |
| params = [p for _, p in named] | |
| opt = torch.optim.AdamW(params, lr=1e-5, weight_decay=0.0) | |
| def run(examples, seed): | |
| toks = 0 | |
| for step in make_steps(examples, c["micro"], STEP_TOKENS, random.Random(seed)): | |
| n_ex = sum(len(mb) for mb in step) | |
| for mb in step: | |
| ckpt(mb) | |
| loss, _ = micro_loss(forward, lm_weight, mb, pad_id, device, chunk=CE_CHUNK) | |
| (loss / n_ex).backward() | |
| toks += sum(len(e["input_ids"]) for e in mb) | |
| torch.nn.utils.clip_grad_norm_(params, 1.0) | |
| opt.step() | |
| opt.zero_grad(set_to_none=True) | |
| return toks | |
| torch.cuda.reset_peak_memory_stats() | |
| run(data["warm"], 1) # kernel compilation, not timed | |
| torch.cuda.synchronize() | |
| t0 = time.time() | |
| toks = run(data["timed"], 2) | |
| torch.cuda.synchronize() | |
| secs = time.time() - t0 | |
| print("RESULT " + json.dumps({"config": args.worker, **c, "tok_per_s": toks / secs, "seconds": secs, "tokens": toks, | |
| "peak_gb": torch.cuda.max_memory_allocated() / 1e9, "check_losses": losses}), flush=True) | |
| 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("--examples", default="/root/bench_tmp/examples.pkl", help="scratch: examples and gradients (~1.5 GB), kept out of results/ and the bucket") | |
| ap.add_argument("--out", default="results/speed_bench.json") | |
| ap.add_argument("--worker", default=None) | |
| ap.add_argument("--scale", type=float, default=1.0, help="smoke tests only: shrink micro-batches and token counts") | |
| args = ap.parse_args() | |
| if args.scale != 1.0: | |
| global TIMED_TOKENS, WARM_TOKENS, STEP_TOKENS | |
| global CE_CHUNK | |
| TIMED_TOKENS, WARM_TOKENS, STEP_TOKENS = (int(x * args.scale) for x in (TIMED_TOKENS, WARM_TOKENS, STEP_TOKENS)) | |
| CE_CHUNK = max(128, int(CE_CHUNK * args.scale)) | |
| for c in CONFIGS.values(): | |
| c["micro"], c["ckpt_above"] = int(c["micro"] * args.scale), int(c["ckpt_above"] * args.scale) | |
| if args.worker: | |
| return worker(args) | |
| import torch | |
| os.makedirs(os.path.dirname(args.examples), exist_ok=True) | |
| prepare(args) | |
| res = {} | |
| def launch(name): | |
| p = subprocess.run([sys.executable, "-u", __file__, "--worker", name, "--model", args.model, "--examples", args.examples, | |
| "--scale", str(args.scale)], | |
| capture_output=True, text=True) | |
| line = next((l for l in p.stdout.splitlines() if l.startswith("RESULT ")), None) | |
| if line is None: | |
| res[name] = {**CONFIGS[name], "failed": (p.stderr or p.stdout)[-1500:]} | |
| print(json.dumps({"config": name, "failed": res[name]["failed"][-300:]}), flush=True) | |
| return | |
| r = json.loads(line[7:]) | |
| if name != "R" and "R" in res and "failed" not in res["R"]: | |
| gr = torch.load(os.path.join(os.path.dirname(args.examples), "bench_grad_R.pt")).double() | |
| g = torch.load(os.path.join(os.path.dirname(args.examples), f"bench_grad_{name}.pt")).double() | |
| r["grad_cosine_vs_fp32"] = float(g @ gr / (g.norm() * gr.norm())) | |
| r["grad_err_vs_fp32"] = float((g - gr).norm() / gr.norm()) | |
| r["loss_err_vs_fp32"] = max(abs(a - b) / abs(b) for a, b in zip(r["check_losses"], res["R"]["check_losses"])) | |
| a = res.get("A", r) | |
| r["equivalent"] = ("failed" not in a and r["grad_err_vs_fp32"] <= 1.5 * a["grad_err_vs_fp32"] + 1e-3 | |
| and r["loss_err_vs_fp32"] <= 1.5 * a["loss_err_vs_fp32"] + 1e-4) | |
| else: | |
| r["equivalent"] = name != "R" and "R" in res and "failed" not in res["R"] | |
| res[name] = r | |
| print(json.dumps(r), flush=True) | |
| for name in ["R", "A", "B", "C", "D"]: | |
| launch(name) | |
| d = res.get("D", {}) | |
| if d.get("equivalent") and d.get("peak_gb", 99) < 48: | |
| launch("E") | |
| ok = {n: r for n, r in res.items() if n != "R" and "failed" not in r and r.get("equivalent") and r["peak_gb"] < 72} | |
| base = res.get("A", {}).get("tok_per_s") | |
| best = max(ok, key=lambda n: ok[n]["tok_per_s"]) if ok else "A" | |
| if base is None or best == "A" or ok[best]["tok_per_s"] < 1.03 * base: | |
| best = "A" | |
| out = {"configs": res, "chosen": best, "flags": flags(CONFIGS[best]), | |
| "speedup_vs_A": (ok[best]["tok_per_s"] / base) if base and best in ok else 1.0} | |
| json.dump(out, open(args.out, "w"), indent=1) | |
| print("CHOSEN " + json.dumps({k: out[k] for k in ("chosen", "flags", "speedup_vs_A")}), flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 10.9 kB
- Xet hash:
- 9bcfde36d98cf12d0df4b8db2f2f071f2dbc4eb676a2ec8d9e1f8f50803f8665
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.