kv-landlords / scripts /humaneval_bench.py

Upload kv-quant (INT4/NVFP4 KIVI) work + vLLM fork source

6e668dc verified 5 months ago

9.78 kB

"""HumanEval subset: pass@1 for BF16 vs INT4-simulated KV cache on Laguna-XS.2.
Loads the first --n HumanEval problems, generates completions under two cache
regimes, executes each solution against the reference test suite, and reports
pass@1 for both modes.
INT4 simulation: after every decode step, each layer's K and V tensors are
quantized to INT4 (MSE-optimal blockwise scale) and immediately dequantized back
— the worst-case accuracy test; errors accumulate across the full generation.
Usage:
    python -m scripts.humaneval_bench [--n 20] [--max-new 512]
"""
from __future__ import annotations

import argparse
import contextlib
import io
import re
import signal
import subprocess
import sys
import tempfile
import textwrap
import time

import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache

sys.path.insert(0, "/home/alex/poolside-hackathon-kv-quant")
from kv_quant import BLOCK, _mse_optimal_scale, quantize_block, dequantize_block

MODEL = "poolside/Laguna-XS.2"

# ---------------------------------------------------------------------------
# INT4 simulation helpers
# ---------------------------------------------------------------------------

def _int4_round_trip(x: torch.Tensor) -> torch.Tensor:
    B, H, S, D = x.shape
    xf = x.float().reshape(B, H, S, D // BLOCK, BLOCK)
    s = _mse_optimal_scale(xf)
    return dequantize_block(quantize_block(xf, s), s).reshape(B, H, S, D).to(x.dtype)


def _quantize_cache(cache: DynamicCache) -> None:
    for layer in cache.layers:
        layer.keys   = _int4_round_trip(layer.keys)
        layer.values = _int4_round_trip(layer.values)


# ---------------------------------------------------------------------------
# Generation
# ---------------------------------------------------------------------------

def _generate(model, input_ids: torch.Tensor, max_new: int, int4: bool) -> str:
    tok_obj = getattr(model, "_tokenizer_ref", None)
    device = input_ids.device
    cache = DynamicCache()
    L = input_ids.shape[1]

    with torch.no_grad():
        if int4:
            # Manual decode loop so we can quantize the cache each step.
            cp = torch.arange(L, device=device)
            out = model(input_ids=input_ids, past_key_values=cache, use_cache=True,
                        cache_position=cp, position_ids=cp.unsqueeze(0))
            _quantize_cache(cache)
            tokens = [out.logits[0, -1].argmax().item()]
            abs_pos = L
            eos = getattr(model.config, "eos_token_id", None)
            eos_set = set(eos) if isinstance(eos, (list, tuple)) else ({eos} if eos else set())
            for _ in range(max_new - 1):
                cp2 = torch.tensor([abs_pos], device=device)
                out = model(input_ids=torch.tensor([[tokens[-1]]], device=device),
                            past_key_values=cache, use_cache=True,
                            cache_position=cp2, position_ids=cp2.unsqueeze(0))
                _quantize_cache(cache)
                t = out.logits[0, -1].argmax().item()
                tokens.append(t)
                abs_pos += 1
                if t in eos_set:
                    break
            return tokens
        else:
            out = model.generate(input_ids, max_new_tokens=max_new,
                                 past_key_values=cache, use_cache=True, do_sample=False)
            return out[0, L:].tolist()


# ---------------------------------------------------------------------------
# Code extraction
# ---------------------------------------------------------------------------

def extract_code(response: str, prompt: str) -> str:
    """Return the best completion string to append to the HumanEval prompt."""
    # Prefer fenced code blocks
    m = re.search(r"```(?:python)?\n(.*?)```", response, re.DOTALL)
    if m:
        block = m.group(1)
        # If the block re-states the function signature, use it wholesale.
        if prompt.split("def ", 1)[-1].split("(")[0].strip() in block:
            return block
        return prompt + block

    # Otherwise, strip everything before the first indented line (function body).
    lines = response.splitlines()
    body = []
    in_body = False
    for line in lines:
        if not in_body:
            if line.startswith("    ") or line.startswith("\t"):
                in_body = True
        if in_body:
            # Stop at a new top-level def (next function)
            if line.startswith("def ") and body:
                break
            body.append(line)
    if body:
        return prompt + "\n".join(body) + "\n"

    # Fallback: append raw response
    return prompt + response


# ---------------------------------------------------------------------------
# Test execution
# ---------------------------------------------------------------------------

_TIMEOUT = 10  # seconds per test

def run_tests(solution_code: str, test_code: str, entry_point: str) -> tuple[bool, str]:
    """Execute solution + tests in a subprocess. Returns (passed, error_msg)."""
    full = solution_code + "\n\n" + test_code + f"\ncheck({entry_point})\n"
    with tempfile.NamedTemporaryFile(suffix=".py", mode="w", delete=False) as f:
        f.write(full)
        fname = f.name
    try:
        r = subprocess.run(
            [sys.executable, fname],
            capture_output=True, text=True, timeout=_TIMEOUT,
        )
        if r.returncode == 0:
            return True, ""
        return False, (r.stderr or r.stdout).strip()[-300:]
    except subprocess.TimeoutExpired:
        return False, "timeout"
    except Exception as e:
        return False, str(e)
    finally:
        import os; os.unlink(fname)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n",        type=int, default=20,  help="number of HumanEval problems")
    ap.add_argument("--max-new",  type=int, default=512, help="max new tokens per completion")
    ap.add_argument("--bf16-only", action="store_true",  help="skip INT4 simulation")
    args = ap.parse_args()

    print(f"[load] {MODEL} ...", flush=True)
    tok = AutoTokenizer.from_pretrained(MODEL)
    model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16, device_map="auto")
    model.eval()
    device = next(model.parameters()).device

    print(f"[data] loading HumanEval ({args.n} problems) ...", flush=True)
    ds = load_dataset("openai/openai_humaneval", split="test").select(range(args.n))

    modes = ["bf16"] + ([] if args.bf16_only else ["int4"])
    results: dict[str, list[bool]] = {m: [] for m in modes}
    errors:  dict[str, list[str]]  = {m: [] for m in modes}

    print(f"\n{'':─<72}")
    print(f"  {'#':>3}  {'task_id':<30}  {'BF16':>6}  {'INT4':>6}  {'note'}")
    print(f"{'':─<72}")

    for idx, prob in enumerate(ds):
        task   = prob["task_id"]
        prompt = prob["prompt"]          # function sig + docstring
        tests  = prob["test"]            # check() function body
        entry  = prob["entry_point"]

        sys_msg = (
            "You are a Python coding assistant. Complete the function below. "
            "Return a fenced ```python``` code block containing the complete "
            "function (including signature and docstring)."
        )
        user_msg = f"Complete this Python function:\n\n```python\n{prompt}```"
        msgs = [{"role": "system", "content": sys_msg},
                {"role": "user",   "content": user_msg}]
        input_ids = tok.apply_chat_template(
            msgs, add_generation_prompt=True, return_tensors="pt", return_dict=False,
        ).to(device)

        row_pass, row_note = {}, {}
        for mode in modes:
            t0 = time.time()
            token_ids = _generate(model, input_ids, args.max_new, int4=(mode == "int4"))
            elapsed = time.time() - t0
            response = tok.decode(token_ids, skip_special_tokens=True)
            code = extract_code(response, prompt)
            passed, err = run_tests(code, tests, entry)
            results[mode].append(passed)
            errors[mode].append("" if passed else err[:80])
            row_pass[mode] = passed
            row_note[mode] = f"{elapsed:.0f}s"

        bf16_sym = "✓" if row_pass.get("bf16") else "✗"
        int4_sym  = ("✓" if row_pass.get("int4") else "✗") if "int4" in modes else "─"
        print(f"  {idx+1:>3}  {task:<30}  {bf16_sym:>6}  {int4_sym:>6}  {row_note.get('bf16','')}", flush=True)

    print(f"{'':─<72}")

    # Summary
    print(f"\n{'':═<72}")
    print("  RESULTS")
    print(f"{'':═<72}")
    for mode in modes:
        passed = sum(results[mode])
        n = len(results[mode])
        fails = [ds[i]["task_id"] for i, p in enumerate(results[mode]) if not p]
        print(f"  {mode.upper():<8}  pass@1 = {passed}/{n}  ({100*passed/n:.0f}%)")
        if fails:
            print(f"            failed: {', '.join(fails)}")
    print()

    if "bf16" in modes and "int4" in modes:
        agree = sum(a == b for a, b in zip(results["bf16"], results["int4"]))
        print(f"  BF16 vs INT4 agreement: {agree}/{args.n} problems ({100*agree/args.n:.0f}%)")
        # Problems where INT4 passed but BF16 failed (or vice versa)
        diff = [(ds[i]["task_id"], results["bf16"][i], results["int4"][i])
                for i in range(args.n) if results["bf16"][i] != results["int4"][i]]
        if diff:
            print("  Differences:")
            for tid, b, q in diff:
                print(f"    {tid}  bf16={'✓' if b else '✗'}  int4={'✓' if q else '✗'}")
    print()


if __name__ == "__main__":
    main()

	"""HumanEval subset: pass@1 for BF16 vs INT4-simulated KV cache on Laguna-XS.2.

	Loads the first --n HumanEval problems, generates completions under two cache
	regimes, executes each solution against the reference test suite, and reports
	pass@1 for both modes.

	INT4 simulation: after every decode step, each layer's K and V tensors are
	quantized to INT4 (MSE-optimal blockwise scale) and immediately dequantized back
	— the worst-case accuracy test; errors accumulate across the full generation.

	Usage:
	python -m scripts.humaneval_bench [--n 20] [--max-new 512]
	"""
	from __future__ import annotations

	import argparse
	import contextlib
	import io
	import re
	import signal
	import subprocess
	import sys
	import tempfile
	import textwrap
	import time

	import torch
	from datasets import load_dataset
	from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache

	sys.path.insert(0, "/home/alex/poolside-hackathon-kv-quant")
	from kv_quant import BLOCK, _mse_optimal_scale, quantize_block, dequantize_block

	MODEL = "poolside/Laguna-XS.2"

	# ---------------------------------------------------------------------------
	# INT4 simulation helpers
	# ---------------------------------------------------------------------------

	def _int4_round_trip(x: torch.Tensor) -> torch.Tensor:
	B, H, S, D = x.shape
	xf = x.float().reshape(B, H, S, D // BLOCK, BLOCK)
	s = _mse_optimal_scale(xf)
	return dequantize_block(quantize_block(xf, s), s).reshape(B, H, S, D).to(x.dtype)


	def _quantize_cache(cache: DynamicCache) -> None:
	for layer in cache.layers:
	layer.keys = _int4_round_trip(layer.keys)
	layer.values = _int4_round_trip(layer.values)


	# ---------------------------------------------------------------------------
	# Generation
	# ---------------------------------------------------------------------------

	def _generate(model, input_ids: torch.Tensor, max_new: int, int4: bool) -> str:
	tok_obj = getattr(model, "_tokenizer_ref", None)
	device = input_ids.device
	cache = DynamicCache()
	L = input_ids.shape[1]

	with torch.no_grad():
	if int4:
	# Manual decode loop so we can quantize the cache each step.
	cp = torch.arange(L, device=device)
	out = model(input_ids=input_ids, past_key_values=cache, use_cache=True,
	cache_position=cp, position_ids=cp.unsqueeze(0))
	_quantize_cache(cache)
	tokens = [out.logits[0, -1].argmax().item()]
	abs_pos = L
	eos = getattr(model.config, "eos_token_id", None)
	eos_set = set(eos) if isinstance(eos, (list, tuple)) else ({eos} if eos else set())
	for _ in range(max_new - 1):
	cp2 = torch.tensor([abs_pos], device=device)
	out = model(input_ids=torch.tensor([[tokens[-1]]], device=device),
	past_key_values=cache, use_cache=True,
	cache_position=cp2, position_ids=cp2.unsqueeze(0))
	_quantize_cache(cache)
	t = out.logits[0, -1].argmax().item()
	tokens.append(t)
	abs_pos += 1
	if t in eos_set:
	break
	return tokens
	else:
	out = model.generate(input_ids, max_new_tokens=max_new,
	past_key_values=cache, use_cache=True, do_sample=False)
	return out[0, L:].tolist()


	# ---------------------------------------------------------------------------
	# Code extraction
	# ---------------------------------------------------------------------------

	def extract_code(response: str, prompt: str) -> str:
	"""Return the best completion string to append to the HumanEval prompt."""
	# Prefer fenced code blocks
	m = re.search(r"```(?:python)?\n(.*?)```", response, re.DOTALL)
	if m:
	block = m.group(1)
	# If the block re-states the function signature, use it wholesale.
	if prompt.split("def ", 1)[-1].split("(")[0].strip() in block:
	return block
	return prompt + block

	# Otherwise, strip everything before the first indented line (function body).
	lines = response.splitlines()
	body = []
	in_body = False
	for line in lines:
	if not in_body:
	if line.startswith(" ") or line.startswith("\t"):
	in_body = True
	if in_body:
	# Stop at a new top-level def (next function)
	if line.startswith("def ") and body:
	break
	body.append(line)
	if body:
	return prompt + "\n".join(body) + "\n"

	# Fallback: append raw response
	return prompt + response


	# ---------------------------------------------------------------------------
	# Test execution
	# ---------------------------------------------------------------------------

	_TIMEOUT = 10 # seconds per test

	def run_tests(solution_code: str, test_code: str, entry_point: str) -> tuple[bool, str]:
	"""Execute solution + tests in a subprocess. Returns (passed, error_msg)."""
	full = solution_code + "\n\n" + test_code + f"\ncheck({entry_point})\n"
	with tempfile.NamedTemporaryFile(suffix=".py", mode="w", delete=False) as f:
	f.write(full)
	fname = f.name
	try:
	r = subprocess.run(
	[sys.executable, fname],
	capture_output=True, text=True, timeout=_TIMEOUT,
	)
	if r.returncode == 0:
	return True, ""
	return False, (r.stderr or r.stdout).strip()[-300:]
	except subprocess.TimeoutExpired:
	return False, "timeout"
	except Exception as e:
	return False, str(e)
	finally:
	import os; os.unlink(fname)


	# ---------------------------------------------------------------------------
	# Main
	# ---------------------------------------------------------------------------

	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--n", type=int, default=20, help="number of HumanEval problems")
	ap.add_argument("--max-new", type=int, default=512, help="max new tokens per completion")
	ap.add_argument("--bf16-only", action="store_true", help="skip INT4 simulation")
	args = ap.parse_args()

	print(f"[load] {MODEL} ...", flush=True)
	tok = AutoTokenizer.from_pretrained(MODEL)
	model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16, device_map="auto")
	model.eval()
	device = next(model.parameters()).device

	print(f"[data] loading HumanEval ({args.n} problems) ...", flush=True)
	ds = load_dataset("openai/openai_humaneval", split="test").select(range(args.n))

	modes = ["bf16"] + ([] if args.bf16_only else ["int4"])
	results: dict[str, list[bool]] = {m: [] for m in modes}
	errors: dict[str, list[str]] = {m: [] for m in modes}

	print(f"\n{'':─<72}")
	print(f" {'#':>3} {'task_id':<30} {'BF16':>6} {'INT4':>6} {'note'}")
	print(f"{'':─<72}")

	for idx, prob in enumerate(ds):
	task = prob["task_id"]
	prompt = prob["prompt"] # function sig + docstring
	tests = prob["test"] # check() function body
	entry = prob["entry_point"]

	sys_msg = (
	"You are a Python coding assistant. Complete the function below. "
	"Return a fenced ```python``` code block containing the complete "
	"function (including signature and docstring)."
	)
	user_msg = f"Complete this Python function:\n\n```python\n{prompt}```"
	msgs = [{"role": "system", "content": sys_msg},
	{"role": "user", "content": user_msg}]
	input_ids = tok.apply_chat_template(
	msgs, add_generation_prompt=True, return_tensors="pt", return_dict=False,
	).to(device)

	row_pass, row_note = {}, {}
	for mode in modes:
	t0 = time.time()
	token_ids = _generate(model, input_ids, args.max_new, int4=(mode == "int4"))
	elapsed = time.time() - t0
	response = tok.decode(token_ids, skip_special_tokens=True)
	code = extract_code(response, prompt)
	passed, err = run_tests(code, tests, entry)
	results[mode].append(passed)
	errors[mode].append("" if passed else err[:80])
	row_pass[mode] = passed
	row_note[mode] = f"{elapsed:.0f}s"

	bf16_sym = "✓" if row_pass.get("bf16") else "✗"
	int4_sym = ("✓" if row_pass.get("int4") else "✗") if "int4" in modes else "─"
	print(f" {idx+1:>3} {task:<30} {bf16_sym:>6} {int4_sym:>6} {row_note.get('bf16','')}", flush=True)

	print(f"{'':─<72}")

	# Summary
	print(f"\n{'':═<72}")
	print(" RESULTS")
	print(f"{'':═<72}")
	for mode in modes:
	passed = sum(results[mode])
	n = len(results[mode])
	fails = [ds[i]["task_id"] for i, p in enumerate(results[mode]) if not p]
	print(f" {mode.upper():<8} pass@1 = {passed}/{n} ({100*passed/n:.0f}%)")
	if fails:
	print(f" failed: {', '.join(fails)}")
	print()

	if "bf16" in modes and "int4" in modes:
	agree = sum(a == b for a, b in zip(results["bf16"], results["int4"]))
	print(f" BF16 vs INT4 agreement: {agree}/{args.n} problems ({100*agree/args.n:.0f}%)")
	# Problems where INT4 passed but BF16 failed (or vice versa)
	diff = [(ds[i]["task_id"], results["bf16"][i], results["int4"][i])
	for i in range(args.n) if results["bf16"][i] != results["int4"][i]]
	if diff:
	print(" Differences:")
	for tid, b, q in diff:
	print(f" {tid} bf16={'✓' if b else '✗'} int4={'✓' if q else '✗'}")
	print()


	if __name__ == "__main__":
	main()