Add lm-eval correctness test

Author: dbarbuzzi

tests/accuracy/lm-eval-tasks.yaml

@@ -0,0 +1,73 @@
+# Llama 2 7B: FP16, FP16 sparse, marlin
+- model_name: "NousResearch/Llama-2-7b-chat-hf"
+  tasks:
+  - name: "gsm8k"
+    metrics:
+    - name: "exact_match,strict-match"
+      value: 0.2266868840030326
+    - name: "exact_match,flexible-extract"
+      value: 0.22820318423047764
+- model_name: "neuralmagic/Llama-2-7b-pruned50-retrained-ultrachat"
+  tasks:
+  - name: "gsm8k"
+    metrics:
+    - name: "exact_match,strict-match"
+      value: 0.09855951478392722
+    - name: "exact_match,flexible-extract"
+      value: 0.10083396512509477
+  extra_args:
+    --sparsity: "sparse_w16a16"
+- model_name: "neuralmagic/llama-2-7b-chat-marlin"
+  tasks:
+  - name: "gsm8k"
+    metrics:
+    - name: "exact_match,strict-match"
+      value: 0.14101592115238817
+    - name: "exact_match,flexible-extract"
+      value: 0.1652767247915087
+# Mistral 7B: FP16, FP16 sparse, marlin
+- model_name: "teknium/OpenHermes-2.5-Mistral-7B"
+  tasks:
+  - name: "gsm8k"
+    metrics:
+    - name: "exact_match,strict-match"
+      value: 0.6004548900682335
+    - name: "exact_match,flexible-extract"
+      value: 0.6482183472327521
+- model_name: "neuralmagic/OpenHermes-2.5-Mistral-7B-pruned50"
+  tasks:
+  - name: "gsm8k"
+    metrics:
+    - name: "exact_match,strict-match"
+      value: 0.4935557240333586
+    - name: "exact_match,flexible-extract"
+      value: 0.5269143290371494
+  extra_args:
+    --sparsity: "sparse_w16a16"
+- model_name: "neuralmagic/OpenHermes-2.5-Mistral-7B-marlin"
+  tasks:
+  - name: "gsm8k"
+    metrics:
+    - name: "exact_match,strict-match"
+      value: 0.4935557240333586
+    - name: "exact_match,flexible-extract"
+      value: 0.5868081880212282
+# Phi 2: marlin
+- model_name: "neuralmagic/phi-2-super-marlin"
+  tasks:
+  - name: "gsm8k"
+    metrics:
+    - name: "exact_match,strict-match"
+      value: 0.49962092494313876
+    - name: "exact_match,flexible-extract"
+      value: 0.5041698256254739
+# Mixtral: FP16
+- model_name: "mistralai/Mixtral-8x7B-Instruct-v0.1"
+  tasks:
+  - name: "gsm8k"
+    metrics:
+    - name: "exact_match,strict-match"
+      value: 0.6550416982562547
+    - name: "exact_match,flexible-extract"
+      value: 0.6603487490523123
+  enable_tensor_parallel: true

tests/accuracy/test_lm_eval_correctness.py

@@ -0,0 +1,102 @@
+import logging
+from pathlib import Path
+from typing import Any, Dict, List, TypedDict
+
+import lm_eval
+import numpy
+import pytest
+import torch
+import yaml
+
+from tests.utils.server import ServerContext
+
+
+class Metric(TypedDict):
+    name: str
+    value: float
+
+
+class Task(TypedDict):
+    name: str
+    metrics: List[Metric]
+
+
+# to support python3.8 typing prior to adding `Required`/`NotRequired`, this class
+# stores the optional keys and the `EvalDefinition` subclass inherits those alongside
+# the required keys it defines.
+class EvalTaskDefinitionOpts(TypedDict, total=False):
+    enable_tensor_parallel: bool
+    extra_args: Dict[str, Any]
+
+
+class EvalTaskDefinition(EvalTaskDefinitionOpts):
+    model_name: str
+    tasks: List[Task]
+
+
+TEST_DATA_FILE = Path(__file__).parent / "lm-eval-tasks.yaml"
+TEST_DATA = yaml.safe_load(TEST_DATA_FILE.read_text(encoding="utf-8"))
+TEST_DATA: List[EvalTaskDefinition] = [
+    pytest.param(eval_def, id=eval_def["model_name"]) for eval_def in TEST_DATA
+]
+
+
+@pytest.mark.parametrize("eval_data", TEST_DATA)
+def test_lm_eval_correctness(
+    eval_data: EvalTaskDefinition,
+    logger: logging.Logger,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    monkeypatch.setenv("TOKENIZERS_PARALLELISM", "false")
+    monkeypatch.setenv("OPENAI_API_KEY", "dummy")
+
+    model_name = eval_data["model_name"]
+    logger.info("building server startup args")
+    vllm_args = {"--model": model_name, "--disable-log-requests": None}
+
+    if eval_data.get("enable_tensor_parallel") is True:
+        tp = torch.cuda.device_count()
+        logger.info("Enabling tensor parallelism with %d devices", tp)
+        vllm_args["--tensor-parallel-size"] = tp
+
+    if extra_args := eval_data.get("extra_args"):
+        vllm_args.update(extra_args)
+
+    openai_args = ",".join(
+        [
+            f"model={model_name}",
+            "tokenizer_backend=huggingface",
+            "base_url=http://localhost:8000/v1",
+        ]
+    )
+
+    logger.info("launching server")
+    with ServerContext(vllm_args, logger=logger) as _:
+        task_names = [t["name"] for t in eval_data["tasks"]]
+        logger.info("getting results for task_names=%s", task_names)
+        results = lm_eval.simple_evaluate(
+            model="local-completions",
+            model_args=openai_args,
+            tasks=task_names,
+            batch_size=64,
+        )
+
+    logger.info("clearing torch cache")
+    lm_eval.models.utils.clear_torch_cache()
+
+    for task in eval_data["tasks"]:
+        logger.info("checking metrics for task=%s", task["name"])
+        for metric in task["metrics"]:
+            ground_truth = metric["value"]
+            measured_value = results["results"][task["name"]][metric["name"]]
+            logger.info(
+                "%s %s:\nground_truth=%s measured_value=%s",
+                task["name"],
+                metric["name"],
+                ground_truth,
+                measured_value,
+            )
+
+            # Metrics must be within 1% of the larger of the two values. This
+            # corresponds to a 99% accuracy threshold.
+            assert numpy.isclose(ground_truth, measured_value, rtol=0.01)