| """Context-length SWEEP of HumanEval pass@1: BF16 vs INT4-KIVI on Laguna-XS.2. |
| |
| Runs the long-prefix HumanEval protocol at increasing prefix lengths in ONE |
| process (model loaded once) and prints a context-vs-accuracy table. Reuses the |
| validated prefix builder, prompt construction, extraction and execution from |
| ``scripts.longctx_code_bench`` / ``scripts.humaneval_bench`` -- this file only |
| orchestrates the sweep. |
| |
| For each prefix size we report: BF16 pass@1, INT4-KIVI pass@1, per-problem |
| agreement, actual mean context length, INT4 cache compression ratio, mean gen |
| time, and (at the largest cell that ran) the INT4 vs BF16 KV memory. |
| |
| Usage |
| ----- |
| .venv/bin/python scripts/longctx_sweep.py --n 8 --max-new 128 \ |
| --prefixes 8000,16000,32000,64000,128000 |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| import time |
|
|
| import torch |
| from datasets import load_dataset |
| from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache |
|
|
| ROOT = "/home/alex/poolside-hackathon-kv-quant" |
| sys.path.insert(0, ROOT) |
|
|
| from scripts.humaneval_bench import extract_code, run_tests |
| from scripts.longctx_code_bench import build_prefix_text, build_input_ids |
| from int4_kivi.hf_cache import Int4KiviCache |
|
|
| MODEL = "poolside/Laguna-XS.2" |
|
|
|
|
| @torch.no_grad() |
| def generate(model, input_ids, max_new, mode, config): |
| cache = Int4KiviCache(config=config) if mode == "int4" else DynamicCache() |
| out = model.generate( |
| input_ids, max_new_tokens=max_new, past_key_values=cache, |
| use_cache=True, do_sample=False, num_beams=1, |
| ) |
| return out[0, input_ids.shape[1]:].tolist(), cache |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--n", type=int, default=8) |
| ap.add_argument("--max-new", type=int, default=128) |
| ap.add_argument("--prefixes", type=str, default="8000,16000,32000,64000,128000") |
| ap.add_argument("--n-large", type=int, default=0, |
| help="if >0, override --n for prefixes >= --large-thresh") |
| ap.add_argument("--large-thresh", type=int, default=100000) |
| args = ap.parse_args() |
|
|
| prefixes = [int(x) for x in args.prefixes.split(",") if x] |
|
|
| print(f"[load] {MODEL} ...", flush=True) |
| tok = AutoTokenizer.from_pretrained(MODEL) |
| model = AutoModelForCausalLM.from_pretrained( |
| MODEL, dtype=torch.bfloat16, device_map="auto" |
| ).eval() |
| device = next(model.parameters()).device |
|
|
| full_ds = load_dataset("openai/openai_humaneval", split="test") |
|
|
| table = [] |
| |
|
|
| for ptarget in prefixes: |
| n = args.n |
| if args.n_large and ptarget >= args.large_thresh: |
| n = args.n_large |
| ds = full_ds.select(range(n)) |
|
|
| print(f"\n{'':=<92}") |
| print(f" PREFIX TARGET {ptarget} tokens (n={n}, max_new={args.max_new})") |
| print(f"{'':=<92}", flush=True) |
| prefix_text, prefix_tok = build_prefix_text(tok, ptarget) |
| print(f" [prefix] actual prefix length: {prefix_tok} tokens", flush=True) |
| print(f" {'#':>3} {'task_id':<16} {'BF16':>5} {'INT4':>5} {'ctx':>8} " |
| f"{'bf16_s':>7} {'int4_s':>7}", flush=True) |
|
|
| res = {"bf16": [], "int4": []} |
| ctxs, ratios, times = [], [], [] |
| last_int4_bytes = last_bf16_bytes = 0 |
| for idx, prob in enumerate(ds): |
| prompt, tests, entry = prob["prompt"], prob["test"], prob["entry_point"] |
| input_ids = build_input_ids(tok, prompt, prefix_text, device) |
| ctx = int(input_ids.shape[1]); ctxs.append(ctx) |
| row, tnote = {}, {} |
| for mode in ("bf16", "int4"): |
| t0 = time.time() |
| ids, cache = generate(model, input_ids, args.max_new, mode, |
| model.config) |
| dt = time.time() - t0; tnote[mode] = dt |
| resp = tok.decode(ids, skip_special_tokens=True) |
| code = extract_code(resp, prompt) |
| passed, _ = run_tests(code, tests, entry) |
| res[mode].append(passed); row[mode] = passed |
| if mode == "int4": |
| try: |
| ratios.append(cache.compression_ratio_vs_bf16()) |
| last_int4_bytes = cache.nbytes() |
| last_bf16_bytes = cache.bf16_nbytes() |
| except Exception: |
| pass |
| times.append(dt) |
| print(f" {idx+1:>3} {prob['task_id']:<16} " |
| f"{'PASS' if row['bf16'] else 'fail':>5} " |
| f"{'PASS' if row['int4'] else 'fail':>5} {ctx:>8} " |
| f"{tnote['bf16']:>6.1f}s {tnote['int4']:>6.1f}s", flush=True) |
|
|
| bf, iq = sum(res["bf16"]), sum(res["int4"]) |
| agree = sum(a == b for a, b in zip(res["bf16"], res["int4"])) |
| ctx_m = sum(ctxs) // len(ctxs) |
| ratio_m = (sum(ratios) / len(ratios)) if ratios else 0.0 |
| t_m = (sum(times) / len(times)) if times else 0.0 |
| print(f" -> ctx~{ctx_m} BF16={bf}/{n} INT4={iq}/{n} " |
| f"agree={agree}/{n} ratio={ratio_m:.2f}x " |
| f"int4_KV={last_int4_bytes/2**20:.0f}MiB " |
| f"bf16_KV={last_bf16_bytes/2**20:.0f}MiB", flush=True) |
| table.append((ptarget, ctx_m, n, bf, iq, agree, ratio_m, t_m, |
| last_int4_bytes, last_bf16_bytes)) |
|
|
| |
| print(f"\n{'':#<100}") |
| print(" CONTEXT-LENGTH SWEEP -- HumanEval pass@1, BF16 vs INT4-KIVI " |
| "(executed, greedy, batch-1)") |
| print(f"{'':#<100}") |
| print(f" {'ctx_tok':>8} {'n':>3} {'BF16':>10} {'INT4-KIVI':>12} " |
| f"{'agree':>8} {'ratio':>7} {'int4_KV':>9} {'bf16_KV':>9} {'t/gen':>7}") |
| print(f"{'':-<100}") |
| for (pt, ctx, n, bf, iq, ag, ratio, tm, ib, bb) in table: |
| print(f" {ctx:>8} {n:>3} {bf:>3}/{n:<3}({100*bf/n:>3.0f}%) " |
| f"{iq:>3}/{n:<3}({100*iq/n:>3.0f}%) {ag:>3}/{n:<3} " |
| f"{ratio:>6.2f}x {ib/2**20:>7.0f}M {bb/2**20:>7.0f}M {tm:>6.1f}s") |
| print(f"{'':-<100}") |
| print("\n[done]", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
