evaluate_quality.py

from __future__ import annotations

import argparse
import dataclasses
import os
from pathlib import Path

import torch
from accelerate import Accelerator
from bigcode_eval.arguments import EvalArguments
from bigcode_eval.evaluator import Evaluator
from transformers import AutoModelForCausalLM, AutoTokenizer

from scripts.inference_config import MODEL, QUANTIZE_LM_HEAD
from scripts.quantization import DequantizedLinearINT4, QuantizedLinearINT4, replace_linear_layers
from scripts.results_io import RESULTS_DIR, save_results
from scripts.roofline import QWEN25_CODER_05B, ceiling_ms, model_bytes_int4

def bigcode_args(limit: int | None, max_length: int, generations_path: Path, metrics_path: Path) -> argparse.Namespace:
    args = dataclasses.asdict(EvalArguments()) 
    args.update(
        model=MODEL,
        modeltype="causal",
        peft_model=None,
        revision=None,
        use_auth_token=False,
        trust_remote_code=False,
        tasks="humaneval",
        instruction_tokens=None,
        batch_size=1,
        max_length_generation=max_length,
        precision="bf16",
        load_in_8bit=False,
        load_in_4bit=False,
        left_padding=False,
        limit=limit,
        limit_start=0,
        save_every_k_tasks=-1,
        postprocess=True,
        allow_code_execution=True,
        generation_only=False,
        load_generations_path=None,
        load_data_path=None,
        metric_output_path=str(metrics_path),
        save_generations=True,
        load_generations_intermediate_paths=None,
        save_generations_path=str(generations_path),
        save_references=False,
        save_references_path="references.json",
        prompt="prompt",
        max_memory_per_gpu=None,
        check_references=False,
    )
    args.update(do_sample=False, n_samples=1, temperature=None, top_p=None, top_k=None) 
    return argparse.Namespace(**args)

HUMANEVAL_DATASET = "openai/openai_humaneval"

def patch_humaneval_task(num_workers: int = 1) -> None:
    import bigcode_eval.tasks as tasks

    base = tasks.TASK_REGISTRY["humaneval"]
    if getattr(base, "_codealign_patched", False):
        return

    class PatchedHumanEval(base):
        _codealign_patched = True
        DATASET_PATH = HUMANEVAL_DATASET

        def __init__(self, **kwargs):
            super().__init__(**kwargs)
            if not hasattr(self, "dataset"):  
                raise RuntimeError(f"could not load {HUMANEVAL_DATASET} from the Hugging Face Hub (see the warning above)")
            self.num_workers = num_workers

    tasks.TASK_REGISTRY["humaneval"] = PatchedHumanEval

PRECISIONS = ("bf16", "int4", "int4-dequant")

def result_tag(precision: str, quantize_lm_head: bool, limit: int | None) -> str:
    tag = precision.replace("-", "_")
    if precision != "bf16" and not quantize_lm_head:
        tag += "_bf16head"
    return tag + ("" if limit is None else f"_limit{limit}")

def run_evaluation(precision: str, limit: int | None, max_length: int, quantize_lm_head: bool) -> dict:
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    print(f"Loading {MODEL} ({dtype})...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL, truncation_side="left", padding_side="right")
    tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=dtype, device_map="cuda")
    model.eval()

    replaced = 0
    linear_impl = "torch (bf16)"
    if precision in ("int4", "int4-dequant"):
        skip = () if quantize_lm_head else ("lm_head",)
        if precision == "int4":
            factory, linear_impl = QuantizedLinearINT4.from_linear, "CUDA kernels (INT4 GEMV v2 / Split-K GEMM)"
        else:
            torch.backends.cuda.matmul.allow_tf32 = False   
            factory, linear_impl = DequantizedLinearINT4.from_linear, "torch F.linear on dequantized fp32 weights"
        replaced = replace_linear_layers(model, skip=skip, factory=factory)
        torch.cuda.empty_cache()
        print(f"Quantized {replaced} nn.Linear layers to INT4 [{linear_impl}] "
              f"(lm_head {'INT4' if quantize_lm_head else 'kept in ' + str(dtype)})")

    tag = result_tag(precision, quantize_lm_head, limit)
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    args = bigcode_args(limit, max_length, RESULTS_DIR / f"humaneval_generations_{tag}.json", RESULTS_DIR / f"humaneval_metrics_{tag}.json")

    os.environ["HF_ALLOW_CODE_EVAL"] = "1"
    patch_humaneval_task()
    evaluator = Evaluator(Accelerator(), model, tokenizer, args)
    print("Starting HumanEval evaluation (greedy pass@1)...")
    results = evaluator.evaluate("humaneval")
    if results is None:
        raise RuntimeError("bigcode Evaluator returned no results")

    pass_at_1 = results["pass@1"] if "pass@1" in results else results.get("humaneval", {}).get("pass@1")
    print(f"\n=== HumanEval pass@1 ({tag}): {pass_at_1} ===")
    save_results(f"humaneval_{tag}", {
        "model": MODEL, "precision": precision, "tag": tag, "linear_impl": linear_impl, "dtype": str(dtype),
        "int4_layers": replaced, "quantize_lm_head": quantize_lm_head if precision != "bf16" else None,
        "limit": limit, "max_length_generation": max_length, "decoding": "greedy, n_samples=1",
        "pass@1": pass_at_1, "raw": results,
        "weights_bytes_per_token": model_bytes_int4(QWEN25_CODER_05B, quantize_lm_head) if precision != "bf16" else None,
        "tpot_ceiling_ms": ceiling_ms(model_bytes_int4(QWEN25_CODER_05B, quantize_lm_head)) if precision != "bf16" else None,
    })
    return results

def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--precision", choices=PRECISIONS, default="int4", help="int4: INT4 on the CUDA kernels; int4-dequant: same INT4 weights on torch matmul")
    parser.add_argument("--limit", type=int, default=None, help="evaluate only the first N problems")
    parser.add_argument("--max-length", type=int, default=512, help="max prompt + completion tokens")
    parser.add_argument("--keep-lm-head", action="store_true", help="do not quantize the lm_head (tag gets _bf16head: never overwrites the int4 results)")
    args = parser.parse_args()
    if args.precision == "bf16" and args.keep_lm_head:
        parser.error("--keep-lm-head only applies to int4 / int4-dequant")
    run_evaluation(args.precision, args.limit, args.max_length, QUANTIZE_LM_HEAD and not args.keep_lm_head)

if __name__ == "__main__":
    main()