diff --git a/tests/e2e_matrix.py b/tests/e2e_matrix.py new file mode 100644 index 000000000..85f53e3c9 --- /dev/null +++ b/tests/e2e_matrix.py @@ -0,0 +1,479 @@ +"""Configurable E2E matrix harness (plan §8). + +One matrix-driven entry point replacing the hardcoded shell sweeps. Each +axis is a comma-separated multi-value flag; the harness takes the Cartesian +product, runs every cell as subprocesses (reusing the existing runners and +comparators -- no inference logic is reimplemented), and writes one JSON +record per cell. + + python -m tests.e2e_matrix \ + --backend hf,vllm --model gpt2,qwen3 \ + --mode eager,cuda_graph --standard transport_bitwise \ + --hooks vllm-full --tp 1 --out results/e2e.jsonl + +Cell dispatch +------------- +- ``vllm`` + ``bitwise`` / ``transport_bitwise`` + vllm_ref_runner (RefDiskWorker, D2D->disk) + vllm_monitored_runner + (ring->ClickHouse) -> vllm_identical_comparator. +- ``vllm`` + ``row_count`` + vllm_monitored_runner -> vllm_rowcnt_comparator (row-count checks only; + value comparison is skipped because no reference tensors are captured). +- ``hf`` + ``allclose`` + hf_reference_runner + hf_monitored_runner -> hf_comparator + (token-id equality + max-abs-diff <= E2E_TOLERANCE on hidden states and + logits). ``bitwise`` and ``row_count`` are not implemented for HF. + +Unsupported combinations (``vllm`` + ``allclose``, ``hf`` + ``bitwise`` / +``row_count``) raise ``ValueError`` at plan time so the matrix never reports +a misleading cell. + +The public ``E2E_HOOK_SELECTION`` input is translated to the internal +``DMX_HOOK_SELECTION`` runtime contract in each subprocess env (plan §2). + +``--dry-run`` prints the planned cells + dispatch commands without touching +CUDA / ClickHouse, so the expansion and env translation are unit-testable on +a CPU-only box. +""" +from __future__ import annotations + +import argparse +import itertools +import json +import os +import shutil +import subprocess +import sys +import tempfile +from dataclasses import dataclass, field +from typing import List, Optional + +from tests.lib.report import CellResult, checks_from_legacy_result, write_jsonl, human_table + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + +# model key -> vLLM ref model source filename (under model_executor/models). +_VLLM_REF_FILES = { + "gpt2": "gpt2_ref.py", + "qwen2_moe": "qwen2_moe_ref.py", + "qwen3": "qwen3_ref.py", + "llama": "llama_ref.py", +} + +# Standards that compare reference D2D buffers against ring/ClickHouse output. +_VLLM_IDENTICAL_STANDARDS = {"bitwise", "transport_bitwise"} +# Standards that only verify row counts (no reference tensors captured). +_VLLM_ROWCOUNT_STANDARDS = {"row_count"} +# The only standard whose comparator is actually implemented for HF. +_HF_SUPPORTED_STANDARDS = {"allclose"} + + +@dataclass(frozen=True) +class Cell: + """One point in the matrix.""" + + backend: str + model: str + mode: str + standard: str + hooks: str + tp: int = 1 + ring_mb: int = 4096 + dtype: str = "bfloat16" + prompt_set: str = "smoke" + + @property + def enforce_eager(self) -> str: + return "1" if self.mode == "eager" else "0" + + +# --------------------------------------------------------------------------- +# Axis expansion +# --------------------------------------------------------------------------- + + +def _split(s: str) -> List[str]: + return [x.strip() for x in s.split(",") if x.strip()] + + +def build_cells(args: argparse.Namespace) -> List[Cell]: + """Cartesian product over every multi-value axis.""" + cells: List[Cell] = [] + for backend, model, mode, standard, hooks, tp, ring_mb, dtype, pset in itertools.product( + _split(args.backend), _split(args.model), _split(args.mode), + _split(args.standard), _split(args.hooks), _split(str(args.tp)), + _split(str(args.ring_mb)), _split(args.dtype), _split(args.prompt_set), + ): + cells.append(Cell( + backend=backend, model=model, mode=mode, standard=standard, + hooks=hooks, tp=int(tp), ring_mb=int(ring_mb), dtype=dtype, + prompt_set=pset, + )) + return cells + + +def cell_env(cell: Cell, args: argparse.Namespace, base: Optional[dict] = None) -> dict: + """Build the subprocess env for a cell. + + Sets the public ``E2E_*`` knobs *and* the translated internal + ``DMX_HOOK_SELECTION`` (plan §2) so the runners see one consistent + configuration. + """ + env = dict(base if base is not None else os.environ) + env["E2E_MODEL"] = cell.model + env["E2E_ENFORCE_EAGER"] = cell.enforce_eager + env["E2E_CUDA_GRAPHS"] = "0" if cell.mode == "eager" else "1" + env["E2E_DTYPE"] = cell.dtype + env["E2E_TP_SIZE"] = str(cell.tp) + env["E2E_RING_PAYLOAD_MB"] = str(cell.ring_mb) + env["E2E_RING_PINNED_MB"] = str(cell.ring_mb) + env["E2E_PROMPT_SET"] = cell.prompt_set + # Public hook-selection input + internal runtime contract translation. + env["E2E_HOOK_SELECTION"] = cell.hooks + env["DMX_HOOK_SELECTION"] = cell.hooks + env["E2E_NUM_PROMPTS"] = str(args.num_prompts) + env["E2E_MAX_NEW_TOKENS"] = str(args.max_new_tokens) + env["E2E_MAX_MODEL_LEN"] = str(args.max_model_len) + env["E2E_MAX_NUM_BATCHED_TOKENS"] = str(args.max_batched_tokens) + env["E2E_GPU_MEM_UTIL"] = str(args.gpu_mem_util) + env["E2E_TOLERANCE"] = str(args.tolerance) + env["DMX_DB_HOST"] = args.db_host + env["DMX_DB_PORT"] = str(args.db_port) + env["VLLM_DISABLE_COMPILE_CACHE"] = "1" + return env + + +# --------------------------------------------------------------------------- +# Dispatch planning (no side effects -- the dry-run surface) +# --------------------------------------------------------------------------- + + +@dataclass +class Step: + """One planned subprocess: a label + argv (env applied at run time).""" + + label: str + argv: List[str] + + +def _runner(mod: str, *flags: str) -> List[str]: + return [sys.executable, "-m", mod, *flags] + + +def plan_cell(cell: Cell, run_dir: str) -> tuple[List[Step], str, str]: + """Return (steps, comparator_module, result_file) for a cell. + + Pure planning: builds the subprocess argv list without executing, so the + same code path feeds both ``--dry-run`` and the real runner. + """ + ref_dir = os.path.join(run_dir, "ref") + mon_dir = os.path.join(run_dir, "mon") + result_file = os.path.join(run_dir, "result.json") + steps: List[Step] = [] + + if cell.backend == "vllm": + if cell.standard in _VLLM_IDENTICAL_STANDARDS: + config_file = os.path.join(ref_dir, "ref_config.json") + steps.append(Step("enable_ref_hooks", ["", "enable_ref_hooks"])) + steps.append(Step("vllm_ref", _runner("tests.vllm_ref_runner", "--output-dir", ref_dir))) + steps.append(Step("vllm_monitored", _runner("tests.vllm_monitored_runner", "--output-dir", mon_dir))) + steps.append(Step("compare", _runner( + "tests.vllm_identical_comparator", + "--ref-config", config_file, "--mon-dir", mon_dir, + "--result-file", result_file))) + return steps, "tests.vllm_identical_comparator", result_file + # row_count -> monitored-only + rowcnt comparator (value comparison skipped) + if cell.standard not in _VLLM_ROWCOUNT_STANDARDS: + raise ValueError( + f"vllm backend does not support standard={cell.standard!r}. " + f"Supported: {sorted(_VLLM_IDENTICAL_STANDARDS | _VLLM_ROWCOUNT_STANDARDS)}" + ) + steps.append(Step("vllm_monitored", _runner("tests.vllm_monitored_runner", "--output-dir", mon_dir))) + steps.append(Step("compare", _runner( + "tests.vllm_rowcnt_comparator", + "--ref-dir", ref_dir, "--mon-dir", mon_dir, + "--result-file", result_file))) + return steps, "tests.vllm_rowcnt_comparator", result_file + + if cell.backend == "hf": + if cell.standard not in _HF_SUPPORTED_STANDARDS: + raise ValueError( + f"hf backend does not support standard={cell.standard!r}. " + f"Supported: {sorted(_HF_SUPPORTED_STANDARDS)}" + ) + steps.append(Step("hf_ref", _runner("tests.hf_reference_runner", "--output-dir", ref_dir))) + steps.append(Step("hf_monitored", _runner("tests.hf_monitored_runner", "--output-dir", mon_dir))) + steps.append(Step("compare", _runner( + "tests.hf_comparator", + "--ref-dir", ref_dir, "--mon-dir", mon_dir, + "--standard", cell.standard, + "--result-file", result_file))) + return steps, "tests.hf_comparator", result_file + + raise ValueError(f"unknown backend {cell.backend!r}") + + +# --------------------------------------------------------------------------- +# Execution +# --------------------------------------------------------------------------- + + +def _enable_vllm_ref_hooks(cell: Cell, run_dir: str, env: dict) -> tuple[str, str, str]: + """Run the in-process enable_ref_hooks preprocessor for a vLLM cell. + + Returns (model_file, backup_file, config_file). Mirrors the flow in + test_vllm_identical: back up the ref model source, generate the hooked + ref + config, leaving restore to the caller's finally. + """ + ref_dir = os.path.join(run_dir, "ref") + os.makedirs(ref_dir, exist_ok=True) + models_dir = os.path.join( + PROJECT_ROOT, "integration", "vllm", "vllm", + "model_executor", "models") + ref_filename = _VLLM_REF_FILES.get(cell.model) + if ref_filename is None: + raise ValueError(f"no vLLM ref model registered for {cell.model!r}") + model_file = os.path.join(models_dir, ref_filename) + backup_file = os.path.join(run_dir, f"{ref_filename}.bak") + config_file = os.path.join(ref_dir, "ref_config.json") + max_len = int(os.environ.get("E2E_REF_MAX_LEN", "8192")) + + shutil.copy2(model_file, backup_file) + sys.path.insert(0, models_dir) + from enable_ref_hooks import enable_ref_hooks # type: ignore + enable_ref_hooks( + model_file=model_file, hooks=cell.hooks, max_len=max_len, + output_dir=ref_dir, config_out=config_file, + ) + return model_file, backup_file, config_file + + +def _run_steps(steps: List[Step], env: dict, run_dir: str, + *, restore=None, timeout: Optional[float] = None) -> Optional[str]: + """Execute the runner/comparator steps in order. + + ``restore`` is a zero-arg callback run after the reference step (for the + vLLM identical flow, which restores the ref model source before the + monitored run). ``timeout`` bounds each subprocess so a hung runner + fails the cell instead of hanging the whole matrix. Returns an error + string on first failure, else None. + """ + for step in steps: + if step.argv and step.argv[0] == "": + continue # enable_ref_hooks handled by the caller + # vLLM identical: restore the patched ref source after the ref run, + # before the monitored run starts. + if restore is not None and step.label in ("vllm_monitored", "hf_monitored"): + restore() + restore = None + env_step = dict(env) + if step.label in ("vllm_ref",): + env_step["REF_CONFIG"] = os.path.join(run_dir, "ref", "ref_config.json") + try: + proc = subprocess.run( + step.argv, env=env_step, cwd=PROJECT_ROOT, + capture_output=True, text=True, timeout=timeout) + except subprocess.TimeoutExpired: + return f"step {step.label} timed out after {timeout}s" + if proc.returncode != 0: + tail = (proc.stderr or "")[-2000:] + return f"step {step.label} failed (rc={proc.returncode}): {tail}" + return None + + +def run_cell(cell: Cell, args: argparse.Namespace) -> CellResult: + """Run one cell end-to-end and return its :class:`CellResult`.""" + cr = CellResult( + backend=cell.backend, model=cell.model, mode=cell.mode, + standard=cell.standard, hook_selection=cell.hooks, tp=cell.tp, + extra={"ring_mb": cell.ring_mb, "dtype": cell.dtype, + "prompt_set": cell.prompt_set}, + ) + run_dir = tempfile.mkdtemp(prefix="e2e_matrix_") + os.makedirs(os.path.join(run_dir, "ref"), exist_ok=True) + os.makedirs(os.path.join(run_dir, "mon"), exist_ok=True) + env = cell_env(cell, args) + backup_file = model_file = None + try: + steps, _comparator, result_file = plan_cell(cell, run_dir) + restore = None + + if cell.backend == "vllm" and cell.standard in _VLLM_IDENTICAL_STANDARDS: + model_file, backup_file, _config = _enable_vllm_ref_hooks(cell, run_dir, env) + + def restore(): # noqa: E306 -- restore ref source pre-monitored run + shutil.copy2(backup_file, model_file) + + elif cell.backend == "vllm": + # rowcnt comparator still expects a ref meta.json (skipped marker). + with open(os.path.join(run_dir, "ref", "meta.json"), "w") as f: + json.dump({"skipped": True}, f) + + err = _run_steps(steps, env, run_dir, restore=restore, + timeout=args.cell_timeout) + if err is not None: + cr.error = err + return cr.finalize() + + with open(result_file) as f: + legacy = json.load(f) + cr.checks = checks_from_legacy_result(legacy) + return cr.finalize() + + except Exception as exc: # noqa: BLE001 -- one bad cell must not abort the matrix + cr.error = f"{type(exc).__name__}: {exc}" + return cr.finalize() + finally: + if backup_file and model_file and os.path.exists(backup_file): + shutil.copy2(backup_file, model_file) + if not args.keep_artifacts: + shutil.rmtree(run_dir, ignore_errors=True) + + +# --------------------------------------------------------------------------- +# Thin-wrapper support (plan §5) +# --------------------------------------------------------------------------- + + +def matrix_argv_from_env(backend: str, standard: str, *, + mode: Optional[str] = None, + default_tolerance: str = "0.01", + env: Optional[dict] = None) -> List[str]: + """Build a single-cell matrix argv from the legacy ``E2E_*`` env knobs. + + The pytest wrappers preserve the old test names + shell entry points + (verify_hf.sh / verify_vllm.sh) and drive the matrix with the same + configuration the legacy tests honored. ``mode`` defaults to + eager/cuda_graph from ``E2E_ENFORCE_EAGER``; the HF cuda-graph wrapper + passes it explicitly. + """ + e = os.environ if env is None else env + if mode is None: + mode = "eager" if e.get("E2E_ENFORCE_EAGER", "1") == "1" else "cuda_graph" + hooks = e.get("E2E_HOOK_SELECTION", e.get("DMX_HOOK_SELECTION", "vllm-full")) + return [ + "--backend", backend, + "--model", e.get("E2E_MODEL", "gpt2"), + "--mode", mode, + "--standard", standard, + "--hooks", hooks, + "--tp", e.get("E2E_TP_SIZE", "1"), + "--ring-mb", e.get("E2E_RING_PAYLOAD_MB", "4096"), + "--dtype", e.get("E2E_DTYPE", "bfloat16"), + "--num-prompts", e.get("E2E_NUM_PROMPTS", "8"), + "--max-new-tokens", e.get("E2E_MAX_NEW_TOKENS", "20"), + "--max-model-len", e.get("E2E_MAX_MODEL_LEN", "512"), + "--max-batched-tokens", e.get("E2E_MAX_NUM_BATCHED_TOKENS", "512"), + "--gpu-mem-util", e.get("E2E_GPU_MEM_UTIL", "0.5"), + "--tolerance", e.get("E2E_TOLERANCE", default_tolerance), + "--db-host", e.get("DMX_DB_HOST", "localhost"), + "--db-port", e.get("DMX_DB_PORT", "9000"), + ] + + +def run_single(argv: List[str]) -> CellResult: + """Parse ``argv`` into exactly one cell, run it, return its CellResult. + + Raises ``ValueError`` if the axes expand to other than one cell -- the + wrappers must drive a single concrete cell. + """ + args = build_parser().parse_args(argv) + cells = build_cells(args) + if len(cells) != 1: + raise ValueError(f"run_single expected exactly 1 cell, got {len(cells)}") + return run_cell(cells[0], args) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + prog="tests.e2e_matrix", + description="Configurable E2E matrix over backend/model/mode/standard/hooks/tp.") + p.add_argument("--backend", default="vllm", help="comma list: hf,vllm") + p.add_argument("--model", default="gpt2", help="comma list: gpt2,qwen3,llama,qwen2_moe,") + p.add_argument("--mode", default="eager", help="comma list: eager,cuda_graph") + p.add_argument("--standard", default="transport_bitwise", + help=( + "comma list of comparison standards. " + "vllm supports: bitwise, transport_bitwise, row_count. " + "hf supports: allclose. " + "Unsupported combinations raise an error at plan time." + )) + p.add_argument("--hooks", default="vllm-full", + help="comma list: preset (vllm-full,hidden-states) or single hook") + p.add_argument("--tp", default="1", help="comma list of tensor-parallel sizes") + p.add_argument("--ring-mb", dest="ring_mb", default="4096", + help="comma list of ring payload/pinned sizes in MB (default 4096)") + p.add_argument("--dtype", default="bfloat16", help="comma list: bfloat16,float16,float32") + p.add_argument("--prompt-set", dest="prompt_set", default="smoke", + help="comma list: smoke,math,chat,random") + p.add_argument("--num-prompts", type=int, default=8) + p.add_argument("--max-new-tokens", type=int, default=20) + p.add_argument("--max-model-len", type=int, default=512) + p.add_argument("--max-batched-tokens", type=int, default=512) + p.add_argument("--gpu-mem-util", type=float, default=0.5) + p.add_argument("--tolerance", type=float, default=0.01, + help="abs tolerance forwarded to comparators (E2E_TOLERANCE)") + p.add_argument("--db-host", default="localhost") + p.add_argument("--db-port", type=int, default=9000) + p.add_argument("--out", default=None, help="JSONL output path (one record per cell)") + p.add_argument("--cell-timeout", type=float, default=1800.0, + help="per-subprocess timeout in seconds (hung runner fails the cell)") + p.add_argument("--keep-artifacts", action="store_true", + help="keep per-cell temp run dirs") + p.add_argument("--dry-run", action="store_true", + help="print planned cells + dispatch commands; no CUDA/ClickHouse") + return p + + +def _dry_run(cells: List[Cell], args: argparse.Namespace) -> int: + print(f"# {len(cells)} cell(s) planned\n") + for i, cell in enumerate(cells): + steps, comparator, _result = plan_cell(cell, run_dir="") + env = cell_env(cell, args, base={}) + print(f"[{i}] backend={cell.backend} model={cell.model} mode={cell.mode} " + f"standard={cell.standard} hooks={cell.hooks} tp={cell.tp} " + f"ring_mb={cell.ring_mb} dtype={cell.dtype} prompt_set={cell.prompt_set}") + print(f" env: E2E_HOOK_SELECTION={env['E2E_HOOK_SELECTION']} -> " + f"DMX_HOOK_SELECTION={env['DMX_HOOK_SELECTION']} " + f"E2E_ENFORCE_EAGER={env['E2E_ENFORCE_EAGER']}") + for step in steps: + print(f" - {step.label}: {' '.join(step.argv)}") + print(f" comparator: {comparator}\n") + return 0 + + +def main(argv: Optional[List[str]] = None) -> int: + args = build_parser().parse_args(argv) + cells = build_cells(args) + if not cells: + print("no cells to run (check axis values)", file=sys.stderr) + return 2 + + if args.dry_run: + return _dry_run(cells, args) + + results: List[CellResult] = [] + for i, cell in enumerate(cells): + print(f"\n=== cell {i + 1}/{len(cells)}: {cell.backend}/{cell.model}/{cell.mode}/" + f"{cell.standard}/{cell.hooks}/tp{cell.tp} ===", flush=True) + cr = run_cell(cell, args) + verdict = "ERROR" if cr.error else ("PASS" if cr.passed else "FAIL") + print(f" -> {verdict}" + (f": {cr.error}" if cr.error else ""), flush=True) + results.append(cr) + + print("\n" + human_table(results)) + if args.out: + write_jsonl(results, args.out) + print(f"\nwrote {len(results)} record(s) to {args.out}") + + # Exit non-zero if any cell failed or errored, so CI can gate on it. + return 0 if all(r.passed for r in results) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/hf_comparator.py b/tests/hf_comparator.py index f93886e5b..edb93cf43 100644 --- a/tests/hf_comparator.py +++ b/tests/hf_comparator.py @@ -23,8 +23,16 @@ def main(): p.add_argument("--ref-dir", required=True) p.add_argument("--mon-dir", required=True) p.add_argument("--result-file", required=True) + p.add_argument("--standard", default="allclose", + help="comparison standard (only 'allclose' is implemented)") args = p.parse_args() + if args.standard != "allclose": + raise SystemExit( + f"hf_comparator: standard={args.standard!r} is not implemented; " + "only 'allclose' is supported" + ) + from tests.hf_reference import ( _HFRef, _load_hf_refs_from_disk, diff --git a/tests/hf_monitored_runner.py b/tests/hf_monitored_runner.py index fcc4195d1..a9d7f4719 100644 --- a/tests/hf_monitored_runner.py +++ b/tests/hf_monitored_runner.py @@ -45,7 +45,7 @@ def main(): model_cls = HookedGPT2LMHeadModel model = model_cls.from_pretrained( - hf_model_id, attn_implementation="eager", torch_dtype=torch.float16 + hf_model_id, attn_implementation="eager", dtype=torch.float16 ).to(device).eval() tokenizer = AutoTokenizer.from_pretrained(hf_model_id) @@ -99,8 +99,15 @@ def main(): from monitoring._native_engine import RingConfig ring_cfg = RingConfig() ring_cfg.task_ring_entries = int(os.environ.get("E2E_RING_TASK_ENTRIES", "16384")) - ring_cfg.payload_ring_bytes = int(os.environ.get("E2E_RING_PAYLOAD_BYTES", str(4 * 1024**3))) - ring_cfg.pinned_staging_bytes = int(os.environ.get("E2E_RING_PINNED_BYTES", str(4 * 1024**3))) + # E2E_RING_PAYLOAD_MB is the matrix-level knob (set by cell_env via --ring-mb). + # E2E_RING_PAYLOAD_BYTES is a fine-grained override; if neither is set, default + # to 256 MB so the runner does not OOM on GPUs smaller than the ring buffer. + _payload_mb = int(os.environ.get("E2E_RING_PAYLOAD_MB", "256")) + ring_cfg.payload_ring_bytes = int(os.environ.get("E2E_RING_PAYLOAD_BYTES", + str(_payload_mb * 1024 * 1024))) + _pinned_mb = int(os.environ.get("E2E_RING_PINNED_MB", "256")) + ring_cfg.pinned_staging_bytes = int(os.environ.get("E2E_RING_PINNED_BYTES", + str(_pinned_mb * 1024 * 1024))) ring_cfg.drain_poll_timeout_us = int(os.environ.get("E2E_DRAIN_POLL_TIMEOUT_US", "100")) ring_cfg.clone_slices = int(os.environ.get("E2E_CLONE_SLICES", "0")) != 0 ring_cfg.insert_queue_max_bytes = int(os.environ.get("E2E_INSERT_QUEUE_MAX_BYTES", str(512 * 1024**2))) diff --git a/tests/hf_reference_runner.py b/tests/hf_reference_runner.py index 5ad2e1165..41e0dfa07 100644 --- a/tests/hf_reference_runner.py +++ b/tests/hf_reference_runner.py @@ -33,7 +33,7 @@ def main(): # Load ORIGINAL model (NOT hooked) model = AutoModelForCausalLM.from_pretrained( - hf_model_id, attn_implementation="eager", torch_dtype=torch.float16 + hf_model_id, attn_implementation="eager", dtype=torch.float16 ).to(device).eval() tokenizer = AutoTokenizer.from_pretrained(hf_model_id) diff --git a/tests/lib/__init__.py b/tests/lib/__init__.py new file mode 100644 index 000000000..bda07b33a --- /dev/null +++ b/tests/lib/__init__.py @@ -0,0 +1,25 @@ +"""Shared E2E test library (plan §7). + +Consolidates the read / merge / align / compare / report logic that the +configurable matrix (:mod:`tests.e2e_matrix`), the pytest wrappers, and the +numeric-difference study all build on, so each rule lives in exactly one +place. + +Submodules: + align -- left-pad strip, EOS trim, request_id ":" parsing + compare -- Check + bitwise / allclose / row_count / transport_bitwise + report -- Check / CellResult dataclasses -> JSON(L) + human table + clickhouse_io -- read offload rows by request_id, dtype/hook maps, row counts + segments -- merge chunked segments -> dense tensors (segment_merger) + disk_ref -- load .pt / structured reference tensors written by ref workers + hf_reference -- ROL + GEN HF rollouts (re-export of tests.hf_reference) + +Only ``align``, ``compare``, and ``report`` are imported eagerly here; they +are pure-CPU (torch-only). The IO-heavy submodules are imported on demand +to keep ``import tests.lib`` cheap and offline-friendly. +""" +from __future__ import annotations + +from tests.lib import align, compare, report # noqa: F401 + +__all__ = ["align", "compare", "report"] diff --git a/tests/lib/align.py b/tests/lib/align.py new file mode 100644 index 000000000..ca31eca17 --- /dev/null +++ b/tests/lib/align.py @@ -0,0 +1,88 @@ +"""Alignment helpers shared by the comparators / matrix (plan §7). + +Consolidates the left-pad strip, EOS trim, and ``request_id`` parsing that +several comparators reimplement. Pure-CPU, ``torch``-only, unit-tested +without CUDA. +""" +from __future__ import annotations + +import re +from typing import Optional, Tuple + +import torch + +# request_id canonical form is ":" -- the group is the +# batched generate() call, the row the position within that batch. +_REQUEST_ID_RE = re.compile(r"^(\d+):(\d+)$") + +# vLLM appends a "-<8 hex>" UUID suffix to request ids; the ref/disk workers +# strip it so monitored and reference rows key the same. +_VLLM_SUFFIX_RE = re.compile(r"-[0-9a-f]{8}$") + + +def parse_request_id(req_id: str) -> Tuple[int, int]: + """Parse ``":"`` into ``(group_id, row_index)``. + + Raises ``ValueError`` on an unexpected format so a malformed id surfaces + loudly rather than silently sorting wrong. + """ + m = _REQUEST_ID_RE.match(req_id) + if not m: + raise ValueError(f"unexpected request_id format: {req_id!r}") + return int(m.group(1)), int(m.group(2)) + + +def normalize_request_id(req_id: str) -> str: + """Strip a trailing vLLM ``-<8hex>`` UUID suffix, if present.""" + return _VLLM_SUFFIX_RE.sub("", req_id) + + +def strip_left_pad(ids_row: torch.Tensor, attn_row: torch.Tensor) -> torch.Tensor: + """Drop left-padding from a single sequence using its attention mask. + + Returns the last ``attn_row.sum()`` ids (HF left-pads, so the real + tokens are the trailing run). Empty (all-pad) rows return an empty + slice. + """ + true_len = int(attn_row.sum().item()) + if true_len <= 0: + return ids_row[:0] + return ids_row[-true_len:] + + +def trim_eos(ids: torch.Tensor, eos_id: int, + *, keep_eos: bool = False) -> torch.Tensor: + """Trim a 1-D id sequence at the first EOS token. + + With ``keep_eos=False`` (default) the EOS itself is dropped; with + ``keep_eos=True`` it is retained. If no EOS is present the sequence is + returned unchanged. + """ + flat = ids.reshape(-1) + hits = torch.nonzero(flat == eos_id, as_tuple=False) + if hits.numel() == 0: + return flat + first = int(hits[0].item()) + return flat[: first + 1] if keep_eos else flat[:first] + + +def align_to_min_len(a: torch.Tensor, b: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Trim two tensors to a common length along dim 0. + + Used before a value comparison when reference and monitored captures + cover slightly different token spans (e.g. generate() drops the final + never-forwarded token). + """ + n = min(a.shape[0], b.shape[0]) + return a[:n], b[:n] + + +def logits_align_skip(db_len: int, ref_len: int) -> int: + """Rows to skip at the head of a DB ``final_logits`` block to align to ref. + + The DB (``logits_to_keep=0``) stores every position + ``[prompt_0..prompt_{N-1}, decode_0..decode_{G-1}]``; ``generate()``'s + ``output_logits`` yields ``[prefill_last, decode_0..decode_{G-2}]``. + The prefill-last DB row sits at ``prompt_len - 1 = db_len - ref_len - 1``. + """ + return max(0, db_len - ref_len - 1) diff --git a/tests/lib/clickhouse_io.py b/tests/lib/clickhouse_io.py new file mode 100644 index 000000000..a3151409e --- /dev/null +++ b/tests/lib/clickhouse_io.py @@ -0,0 +1,124 @@ +"""ClickHouse read helpers for the comparators / matrix (plan §7). + +Consolidates the row-decode logic, dtype table, and short-hook -> CH +``act_name`` map that ``vllm_identical_comparator``, ``compare_disk_vs_ch``, +and ``vllm_rowcnt_comparator`` each reimplement. + +``clickhouse_driver`` is imported lazily inside the functions so importing +this module stays CPU/offline-friendly (the unit suite never reaches a DB). +""" +from __future__ import annotations + +from typing import Dict, List, Tuple + +import torch + +# CH stores the torch dtype as its ``str(dtype)``; map back on read. +DTYPE_MAP: Dict[str, torch.dtype] = { + "torch.bfloat16": torch.bfloat16, "torch.float": torch.float32, + "torch.float32": torch.float32, "torch.half": torch.float16, + "torch.float16": torch.float16, "torch.int": torch.int32, + "torch.int32": torch.int32, "torch.long": torch.int64, + "torch.int64": torch.int64, "torch.uint8": torch.uint8, + "torch.int8": torch.int8, "torch.short": torch.int16, + "torch.double": torch.float64, "torch.bool": torch.bool, +} + +# Short hook name (the ``_buf_`` suffix / disk filename stem) -> the +# ClickHouse ``act_name``. Must match tensor_meta.h hook_type_name() and the +# p2p make_act_name() convention. +HOOK_TO_CH_ACT: Dict[str, str] = { + "resid_pre": "blocks.hook_resid_pre", + "ln1": "blocks.hook_ln1", + "q": "blocks.attn.hook_q", + "k": "blocks.attn.hook_k", + "v": "blocks.attn.hook_v", + "z": "blocks.attn.hook_z", + "attn_scores": "blocks.attn.hook_attn_scores", + "pattern": "blocks.attn.hook_pattern", + "attn_out": "blocks.hook_attn_out", + "resid_mid": "blocks.hook_resid_mid", + "ln2": "blocks.hook_ln2", + "mlp_in": "blocks.hook_mlp_in", + "mlp_out": "blocks.hook_mlp_out", + "mlp_post": "blocks.hook_mlp_post", + "embed": "hook_embed", + "pos_embed": "hook_pos_embed", + "resid_final": "hook_resid_final", + "final_ln": "hook_final_ln", + "final_logits": "final_logits", + "token_ids": "token_ids", + "router_logits": "blocks.mlp.hook_router_logits", + "topk_ids": "blocks.mlp.hook_topk_ids", + "topk_weights": "blocks.mlp.hook_topk_weights", +} + +# A CH row key: (req_id, act_name, layer_no, shard_rank, start_token, end_token). +RowKey = Tuple[str, str, int, int, int, int] + + +def _decode(v) -> str: + return v.decode() if isinstance(v, bytes) else v + + +def read_offload_rows( + db_host: str, db_port: int, *, + database: str = "default", table: str = "offload", +) -> Tuple[Dict[RowKey, torch.Tensor], int]: + """Read every row from ``.`` into a keyed dict. + + Returns ``(rows_by_key, num_rows)`` where the key is :data:`RowKey` and + the value the decoded CPU tensor. Raises whatever ``clickhouse_driver`` + raises on a connection / query error -- callers decide whether that is a + soft "db unreachable" skip or a hard failure. + """ + import clickhouse_driver + + client = clickhouse_driver.Client(db_host, port=db_port) + raw_rows = client.execute( + "SELECT model_id, request_id, act_name, layer_no, shard_rank, " + "start_token_idx, end_token_idx, dtype, shape, bytes " + f"FROM {database}.{table}", + settings={"strings_as_bytes": True}, + ) + + out: Dict[RowKey, torch.Tensor] = {} + for row in raw_rows: + _, req_id, act_name, layer_no, shard_rank, s, e, dtype_str, shape, payload = row + dt = DTYPE_MAP.get(_decode(dtype_str), torch.float32) + t = torch.frombuffer(bytearray(payload), dtype=dt).reshape(list(shape)) + out[(_decode(req_id), _decode(act_name), int(layer_no), + int(shard_rank), int(s), int(e))] = t + return out, len(raw_rows) + + +def per_hook_counts(rows_by_key: Dict[RowKey, torch.Tensor]) -> Dict[str, int]: + """Count rows per ``act_name`` (input to the ``row_count`` standard).""" + counts: Dict[str, int] = {} + for key in rows_by_key: + act = key[1] + counts[act] = counts.get(act, 0) + 1 + return counts + + +def group_by_request( + rows_by_key: Dict[RowKey, torch.Tensor], +) -> Dict[str, Dict[Tuple[int, str], List[Tuple[int, int, torch.Tensor]]]]: + """Regroup CH rows as ``req_id -> (layer_no, act_name) -> [(s, e, t)]``. + + ``act_name`` is canonicalised by stripping a leading ``"blocks."`` so a + per-layer hook keys as ``(layer_no, "hook_resid_pre")`` and a global hook + as ``(-1, "final_logits")``. Segments are left unsorted; merge callers + sort by start token. + """ + grouped: Dict[str, Dict[Tuple[int, str], List[Tuple[int, int, torch.Tensor]]]] = {} + for (req_id, act_name, layer_no, _shard, s, e), t in rows_by_key.items(): + if act_name.startswith("blocks."): + canon = act_name[len("blocks."):] + lno = layer_no + else: + canon = act_name + lno = -1 + grouped.setdefault(req_id, {}).setdefault((lno, canon), []).append( + (s, e, t.detach().cpu())) + return grouped diff --git a/tests/lib/compare.py b/tests/lib/compare.py new file mode 100644 index 000000000..8ee02a58c --- /dev/null +++ b/tests/lib/compare.py @@ -0,0 +1,211 @@ +"""Comparison standards for the E2E matrix and pytest wrappers (plan §7). + +Four standards, one common interface. Each returns a :class:`Check` +recording pass/fail plus the numeric drift (``max_abs`` / ``mean_abs`` / +``first_diff_pos``) **even on a pass**, so a green cell still surfaces a +"barely passing" trend: + +- ``bitwise`` -- exact equality (raw bytes / ``torch.equal``). +- ``allclose`` -- ``torch.allclose(atol, rtol)`` with a named, + reported threshold. +- ``row_count`` -- schema + segment-count validation only. +- ``transport_bitwise`` -- ``.copy_()`` reference buffers vs ClickHouse ring + output; exact, same engine as ``bitwise`` but a + distinct name so the gating policy (§8) can treat + transport separately from model-output transparency. + +This module is pure-CPU and only depends on ``torch``; it is unit-tested +without CUDA / ClickHouse / vLLM. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Callable, Optional + +import torch + + +@dataclass +class Check: + """One comparison result. + + ``max_abs`` / ``mean_abs`` / ``first_diff_pos`` are recorded whenever + they can be computed (even on pass) so trends stay visible. ``detail`` + is a short human string for tables / assertion messages. + """ + + name: str + passed: bool + max_abs: Optional[float] = None + mean_abs: Optional[float] = None + first_diff_pos: Optional[int] = None + detail: str = "" + + def to_dict(self) -> dict: + d: dict = {"name": self.name, "passed": self.passed} + if self.max_abs is not None: + d["max_abs"] = self.max_abs + if self.mean_abs is not None: + d["mean_abs"] = self.mean_abs + if self.first_diff_pos is not None: + d["first_diff_pos"] = self.first_diff_pos + if self.detail: + d["detail"] = self.detail + return d + + +# --------------------------------------------------------------------------- +# Low-level helpers +# --------------------------------------------------------------------------- + + +_INT_VIEW = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64} + + +def bytes_identical(a: torch.Tensor, b: torch.Tensor) -> bool: + """True iff ``a`` and ``b`` have identical raw bytes. + + Reinterprets as a same-width integer dtype so the compare is a true + bitwise check (no NaN / signed-zero surprises). Falls back to a + storage-bytes compare for exotic element sizes. + """ + a_c = a.contiguous() + b_c = b.contiguous() + if a_c.shape != b_c.shape or a_c.element_size() != b_c.element_size(): + return False + dt = _INT_VIEW.get(a_c.element_size()) + if dt is None: + return bytes(a_c.untyped_storage()) == bytes(b_c.untyped_storage()) + return torch.equal(a_c.view(dt), b_c.view(dt)) + + +def _abs_diff_stats(a: torch.Tensor, b: torch.Tensor) -> tuple[float, float, Optional[int]]: + """Return (max_abs, mean_abs, first_diff_pos) over the flattened tensors. + + ``first_diff_pos`` is the index of the first element that differs in the + flattened view, or ``None`` if the tensors are elementwise equal. + """ + af = a.detach().float().reshape(-1) + bf = b.detach().float().reshape(-1) + n = min(af.numel(), bf.numel()) + af = af[:n] + bf = bf[:n] + diff = (af - bf).abs() + max_abs = float(diff.max().item()) if n else 0.0 + mean_abs = float(diff.mean().item()) if n else 0.0 + ne = torch.nonzero(af != bf, as_tuple=False) + first = int(ne[0].item()) if ne.numel() else None + return max_abs, mean_abs, first + + +def _precheck(a: torch.Tensor, b: torch.Tensor, name: str) -> Optional[Check]: + """Shape/dtype gate shared by the tensor standards. + + Returns a failing :class:`Check` on mismatch, else ``None``. + """ + if a.shape != b.shape: + return Check(name, False, detail=f"shape mismatch: {list(a.shape)} vs {list(b.shape)}") + if a.dtype != b.dtype: + return Check(name, False, detail=f"dtype mismatch: {a.dtype} vs {b.dtype}") + return None + + +# --------------------------------------------------------------------------- +# The four standards +# --------------------------------------------------------------------------- + + +def bitwise(a: torch.Tensor, b: torch.Tensor, name: str = "bitwise") -> Check: + """Exact equality. Records drift stats when it fails.""" + pre = _precheck(a, b, name) + if pre is not None: + return pre + if bytes_identical(a, b): + return Check(name, True, max_abs=0.0, mean_abs=0.0, detail="bitwise equal") + max_abs, mean_abs, first = _abs_diff_stats(a, b) + return Check(name, False, max_abs=max_abs, mean_abs=mean_abs, + first_diff_pos=first, detail=f"max_abs={max_abs:.6e}") + + +def transport_bitwise(a: torch.Tensor, b: torch.Tensor, + name: str = "transport_bitwise") -> Check: + """``.copy_()`` reference buffer vs ring/ClickHouse output -- exact. + + Identical engine to :func:`bitwise`; a distinct standard name so the + §8 gating policy can keep transport bitwise even when model-output + transparency is allowed to use ``allclose`` under CUDA graphs. + """ + return bitwise(a, b, name) + + +def allclose(a: torch.Tensor, b: torch.Tensor, name: str = "allclose", + *, atol: float = 1e-3, rtol: float = 0.0) -> Check: + """``torch.allclose`` with a named, reported threshold. + + Always records max/mean abs diff -- even on a pass -- so the threshold + headroom stays visible. + """ + pre = _precheck(a, b, name) + if pre is not None: + return pre + max_abs, mean_abs, first = _abs_diff_stats(a, b) + passed = bool(torch.allclose(a.float(), b.float(), atol=atol, rtol=rtol)) + return Check( + name, passed, max_abs=max_abs, mean_abs=mean_abs, + first_diff_pos=None if passed else first, + detail=f"max_abs={max_abs:.6e} (atol={atol:g}, rtol={rtol:g})", + ) + + +def row_count(per_hook_counts: dict, name: str = "row_count", *, + min_per_layer_types: int = 10, + require_final_logits: bool = True) -> Check: + """Schema + segment-count validation only (no value comparison). + + ``per_hook_counts`` maps a ClickHouse ``act_name`` to the number of rows + captured for it. Validates that: + + - there are enough per-layer hook types (``blocks.*``), + - every per-layer hook captured the same number of rows, and + - the global ``final_logits`` hook is present (when required). + """ + if not per_hook_counts: + return Check(name, False, detail="no rows") + per_layer = {k: v for k, v in per_hook_counts.items() if k.startswith("blocks.")} + problems: list[str] = [] + if len(per_layer) < min_per_layer_types: + problems.append(f"only {len(per_layer)} per-layer types (<{min_per_layer_types})") + counts = set(per_layer.values()) + if len(counts) > 1: + problems.append(f"uneven per-layer counts: {sorted(counts)}") + if require_final_logits and "final_logits" not in per_hook_counts: + problems.append("final_logits missing") + passed = not problems + detail = "ok" if passed else "; ".join(problems) + return Check(name, passed, detail=detail) + + +# Registry of the tensor-pair standards for matrix dispatch by name. +TENSOR_STANDARDS: dict[str, Callable[..., Check]] = { + "bitwise": bitwise, + "transport_bitwise": transport_bitwise, + "allclose": allclose, +} + +ALL_STANDARDS = tuple(TENSOR_STANDARDS) + ("row_count",) + + +def compare_tensors(a: torch.Tensor, b: torch.Tensor, standard: str, + name: Optional[str] = None, **kwargs) -> Check: + """Dispatch a tensor-pair comparison by standard name. + + ``standard`` must be one of :data:`TENSOR_STANDARDS` (``row_count`` is + not a tensor-pair standard -- call :func:`row_count` directly). + """ + fn = TENSOR_STANDARDS.get(standard) + if fn is None: + raise ValueError( + f"unknown tensor standard {standard!r}; " + f"expected one of {sorted(TENSOR_STANDARDS)}" + ) + return fn(a, b, name or standard, **kwargs) diff --git a/tests/lib/disk_ref.py b/tests/lib/disk_ref.py new file mode 100644 index 000000000..0e988a2b4 --- /dev/null +++ b/tests/lib/disk_ref.py @@ -0,0 +1,84 @@ +"""Load reference tensors written to disk by the ref workers (plan §7). + +Two on-disk reference shapes exist: + +- The vLLM ``RefDiskWorker`` writes per-request ``.pt`` files named + ``{hook}_L{layer}_T{start}_{end}[_SR{rank}].pt`` under + ``//``. :func:`scan_pt_ref_files` parses those. +- The HF reference runner writes a structured dump consumed via + ``tests.hf_reference._load_hf_refs_from_disk``; :func:`load_hf_refs` + re-exports it (lazily, to avoid importing the heavy HF module on CPU). +""" +from __future__ import annotations + +import re +from dataclasses import dataclass +from pathlib import Path +from typing import List, Optional + +# {hook}[_L{layer}]_T{start}_{end}[_SR{rank}].pt +PT_RE = re.compile( + r"^(?P\w+?)(?:_L(?P\d+))?_T(?P\d+)_(?P\d+)" + r"(?:_SR(?P\d+))?\.pt$" +) + + +@dataclass(frozen=True) +class RefFile: + """One parsed reference ``.pt`` file.""" + + req_id: str + hook: str + layer: int # -1 for global hooks + shard: int # 0 when TP == 1 + start: int + end: int + path: str + + @property + def label(self) -> str: + lbl = f"{self.req_id}/{self.hook}" + if self.layer >= 0: + lbl += f"_L{self.layer}" + return lbl + f"_T{self.start}_{self.end}" + + +def parse_pt_name(name: str) -> Optional[dict]: + """Parse a ref ``.pt`` filename into its fields, or ``None`` if unmatched.""" + m = PT_RE.match(name) + if not m: + return None + return { + "hook": m.group("hook"), + "layer": int(m.group("layer")) if m.group("layer") is not None else -1, + "shard": int(m.group("shard")) if m.group("shard") is not None else 0, + "start": int(m.group("start")), + "end": int(m.group("end")), + } + + +def scan_pt_ref_files(ref_dir: str) -> List[RefFile]: + """Walk ``//*.pt`` and return parsed :class:`RefFile`s.""" + root = Path(ref_dir) + out: List[RefFile] = [] + for req_dir in sorted(root.iterdir()): + if not req_dir.is_dir(): + continue + for pt_file in sorted(req_dir.iterdir()): + parsed = parse_pt_name(pt_file.name) + if parsed is None: + continue + out.append(RefFile(req_id=req_dir.name, path=str(pt_file), **parsed)) + return out + + +def load_pt(path: str): + """Load a single reference tensor (CPU, weights-only).""" + import torch + return torch.load(path, weights_only=True, map_location="cpu") + + +def load_hf_refs(ref_dir: str): + """Load the HF structured reference dump (delegates to tests.hf_reference).""" + from tests.hf_reference import _load_hf_refs_from_disk + return _load_hf_refs_from_disk(ref_dir) diff --git a/tests/lib/hf_reference.py b/tests/lib/hf_reference.py new file mode 100644 index 000000000..4bdaf7eb0 --- /dev/null +++ b/tests/lib/hf_reference.py @@ -0,0 +1,28 @@ +"""HF reference rollouts for the matrix / wrappers (plan §7). + +Re-exports the ROL (manual KV-cache rollout: full logits + hidden states + +attn patterns) and GEN (``generate()`` token_ids + decode scores) reference +helpers under stable public names. The canonical implementation still +lives in :mod:`tests.hf_reference`; this shim gives the shared ``tests.lib`` +namespace a single import point without moving the 700-line module (that +relocation is deferred to the legacy-removal PR so this one stays additive). +""" +from __future__ import annotations + +from tests.hf_reference import ( # noqa: F401 (re-exported) + _HFRef as HFRef, + _HFGenRef as HFGenRef, + _hf_greedy_rollout_collect_all_batched as rollout_collect_all, + _hf_generate_collect_scores_batched as generate_collect_scores, + _hf_generate_collect_hidden_states_batched as generate_collect_hidden_states, + _load_hf_refs_from_disk as load_refs_from_disk, +) + +__all__ = [ + "HFRef", + "HFGenRef", + "rollout_collect_all", + "generate_collect_scores", + "generate_collect_hidden_states", + "load_refs_from_disk", +] diff --git a/tests/lib/report.py b/tests/lib/report.py new file mode 100644 index 000000000..7923d9e10 --- /dev/null +++ b/tests/lib/report.py @@ -0,0 +1,126 @@ +"""Machine- and human-readable matrix output (plan §7, §8). + +A :class:`CellResult` is one cell of the E2E matrix -- one +``(backend, model, mode, standard, hook_selection, tp, ...)`` point. It +serialises to a single JSON record (one per line in the JSONL artifact, §8) +and renders into a compact human table. + +Pure-CPU; depends only on the stdlib and :mod:`tests.lib.compare`. +""" +from __future__ import annotations + +import json +from dataclasses import dataclass, field, asdict +from typing import Any, List, Optional + +from tests.lib.compare import Check + + +@dataclass +class CellResult: + """Result of one matrix cell. + + ``checks`` holds the per-tensor / per-hook :class:`Check` objects. + ``passed`` is the cell verdict (all checks passed and no error). + ``error`` carries a setup/dispatch failure message when the cell could + not run at all (distinct from a cell that ran and failed a check). + ``extra`` stashes axis values that don't have a first-class field + (ring sizes, dtype, prompt-set, ...). + """ + + backend: str + model: str + mode: str + standard: str + hook_selection: str + tp: int = 1 + passed: bool = False + checks: List[Check] = field(default_factory=list) + error: Optional[str] = None + extra: dict = field(default_factory=dict) + + def finalize(self) -> "CellResult": + """Set ``passed`` from the checks (no error + every check passed).""" + self.passed = self.error is None and bool(self.checks) and all( + c.passed for c in self.checks) + return self + + def to_record(self) -> dict: + rec: dict[str, Any] = { + "backend": self.backend, + "model": self.model, + "mode": self.mode, + "standard": self.standard, + "hook_selection": self.hook_selection, + "tp": self.tp, + "passed": self.passed, + "checks": [c.to_dict() for c in self.checks], + } + if self.error is not None: + rec["error"] = self.error + if self.extra: + rec["extra"] = self.extra + return rec + + +def checks_from_legacy_result(result: dict) -> List[Check]: + """Adapt a legacy comparator ``result.json`` into :class:`Check` objects. + + The existing comparators emit ``{"tests": [{"name", "passed", "detail"}]}``; + this lets the matrix dispatch to them unchanged and still produce the + new record shape. + """ + out: List[Check] = [] + for t in result.get("tests", []): + out.append(Check( + name=t.get("name", "?"), + passed=bool(t.get("passed", False)), + detail=t.get("detail", ""), + )) + return out + + +def write_jsonl(results: List[CellResult], path: str) -> None: + """Write one JSON record per line (creates parent dirs).""" + import os + os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) + with open(path, "w") as f: + for r in results: + f.write(json.dumps(r.to_record()) + "\n") + + +def read_jsonl(path: str) -> List[dict]: + """Read a JSONL artifact back into a list of records.""" + out: List[dict] = [] + with open(path) as f: + for line in f: + line = line.strip() + if line: + out.append(json.loads(line)) + return out + + +def human_table(results: List[CellResult]) -> str: + """Render the matrix as a compact fixed-width table.""" + header = ("backend", "model", "mode", "standard", "hooks", "tp", "result", "n_fail") + rows: List[tuple] = [] + for r in results: + n_fail = sum(1 for c in r.checks if not c.passed) + verdict = "ERROR" if r.error is not None else ("PASS" if r.passed else "FAIL") + rows.append(( + r.backend, r.model, r.mode, r.standard, r.hook_selection, + str(r.tp), verdict, str(n_fail), + )) + widths = [len(h) for h in header] + for row in rows: + for i, cell in enumerate(row): + widths[i] = max(widths[i], len(cell)) + fmt = " ".join(f"{{:<{w}}}" for w in widths) + lines = [fmt.format(*header), fmt.format(*("-" * w for w in widths))] + lines += [fmt.format(*row) for row in rows] + n_pass = sum(1 for r in results if r.passed) + n_err = sum(1 for r in results if r.error is not None) + lines.append("") + lines.append(f"{n_pass}/{len(results)} cells passed" + + (f", {n_err} errored" if n_err else "")) + return "\n".join(lines) diff --git a/tests/lib/segments.py b/tests/lib/segments.py new file mode 100644 index 000000000..1fc9e066b --- /dev/null +++ b/tests/lib/segments.py @@ -0,0 +1,33 @@ +"""Chunked-segment merge helpers (plan §7). + +Thin wrapper over :mod:`monitoring.segment_merger` so the matrix, the +pytest wrappers, and the numeric study all merge chunked ring/CH segments +through one entry point. Re-exports the canonical implementation rather +than reimplementing the per-act-name merge rules. +""" +from __future__ import annotations + +from typing import List, Optional, Tuple + +import torch + +from monitoring.segment_merger import ( # noqa: F401 (re-exported) + merge_segments, + segment_manager, + parse_internal_id, + get_delta_token_len, +) + + +def merge_request_chunks( + chunks: List[Tuple[int, int, torch.Tensor]], act_name: str, + *, drop_token_cnt_to: Optional[int] = None, +) -> torch.Tensor: + """Sort ``(start, end, tensor)`` chunks by start token and merge them. + + Convenience over :func:`merge_segments` for the + ``group_by_request`` output shape (lists of ``(s, e, t)`` triples). + """ + ordered = sorted(chunks, key=lambda c: c[0]) + return merge_segments( + [t for _, _, t in ordered], act_name, drop_token_cnt_to=drop_token_cnt_to) diff --git a/tests/numeric_study.py b/tests/numeric_study.py new file mode 100644 index 000000000..899ccbfdb --- /dev/null +++ b/tests/numeric_study.py @@ -0,0 +1,794 @@ +"""Per-hook numeric-difference study (plan §9 / Phase 5). + +Enables **one hook at a time** and reports the drift its monitoring path +introduces versus the **unhooked** baseline model. The non-goal carried from +the issue holds: this does not *fix* numeric drift, it makes it *visible, +categorized, and reproducible*. + +Algorithm (plan §9): + + 1. Run the baseline **unhooked** model once; capture token ids + full logits. + 2. For each hook ``H`` in the selection, enable **only** ``H`` and run the + monitored model: + - ``--variant p`` (default): the production ``_p`` Hooked variant driven + with ``hook_selection=H`` -- ``hook_selection`` already isolates a + single hook, so no source patching is needed. + - ``--variant compare``: the ``_compare`` variant under the hardened + :func:`tests.isolate_hook.isolated_hook` context manager (plan §6), + which patches the vendored source so only ``H``'s ``.copy_()`` line + fires and asserts byte-identical restoration on exit. + 3. Compare against the baseline and record, per hook: + - bitwise pass/fail (eager) or allclose-within-threshold (cuda graph), + via the shared standards in :mod:`tests.lib.compare`, + - max abs diff, mean abs diff, first differing token position, + - top-k vocab diffs at the first differing position, + - whether greedy (argmax) token ids diverged. + 4. Emit a machine-readable JSON artifact **and** a human-readable table. + +Alert policy (plan §9), expressed through the §8 standards-by-mode choice: + + - **Eager** -> ``bitwise`` standard: *any* non-bitwise logits diff alerts. + - **CUDA graph** -> ``allclose`` standard with a per-model/per-dtype + threshold: a max abs diff over the threshold alerts. + - **Always** alert on greedy token-id divergence (a hook flipped the + argmax), on a runner error, on a shape mismatch (a possible hook identity + swap), or on an empty capture. + +The capture side reuses the proven subprocess-rollout pattern from +``tests/test_per_hook_isolation.py`` (a clean CUDA context / ring-transport +instance per cell). HF saves full ``[N, vocab]`` logits. vLLM saves a +same-shaped tensor whose non-top-k entries are filled to ``-1e30`` (top-k +logprob capture only); drift metrics are restricted to reported entries. + +The comparison / report / alert core (``compute_drift``, ``format_table``, the +``*_to_dict`` helpers) is pure and CPU-testable; ``torch`` is imported lazily +inside the helpers that need it, so this module loads on a torch-less box. + +CLI:: + + python -m tests.numeric_study \\ + --framework hf --model qwen3 --mode eager \\ + --hooks q,k,resid_pre,final_logits \\ + --out results/numeric_qwen3_eager.json +""" +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +from dataclasses import dataclass, field +from pathlib import Path +from textwrap import dedent +from typing import Any, Dict, List, Optional + +from tests.lib.compare import Check, allclose, bitwise + +REPO_ROOT = Path(__file__).resolve().parent.parent + +# Per-(model, dtype) max-abs-diff thresholds for the CUDA-graph allclose gate. +# Eager ignores the threshold (it gates strictly on bitwise equality); the +# CUDA-graph path tolerates inductor's per-class fusion noise up to this bound. +_DEFAULT_THRESHOLD = 0.15 +_CUDA_GRAPH_THRESHOLDS: Dict[tuple, float] = { + ("gpt2", "float16"): 0.15, + ("qwen3", "float16"): 0.15, + ("qwen2_moe", "float16"): 0.25, +} + +# Default hooks when ``--hooks`` is not given: an attention projection, a +# residual-stream read, and the final-logits read. +_DEFAULT_HOOKS = ["q", "k", "resid_pre", "final_logits"] + + +def cuda_graph_threshold(model_key: str, dtype: str) -> float: + """Resolve the CUDA-graph max-abs-diff alert threshold for a cell.""" + return _CUDA_GRAPH_THRESHOLDS.get((model_key, dtype), _DEFAULT_THRESHOLD) + + +def standard_for_mode(mode: str) -> str: + """The §8 comparison standard for a mode: bitwise (eager) / allclose (cg).""" + return "bitwise" if mode == "eager" else "allclose" + + +# --------------------------------------------------------------------------- +# Result dataclasses (JSON-serializable; no torch types stored) +# --------------------------------------------------------------------------- + + +@dataclass +class VocabDiff: + """One vocab entry's logit drift at the first differing token position.""" + + token_id: int + baseline_logit: float + monitored_logit: float + abs_diff: float + + def to_dict(self) -> Dict[str, Any]: + return { + "token_id": self.token_id, + "baseline_logit": self.baseline_logit, + "monitored_logit": self.monitored_logit, + "abs_diff": self.abs_diff, + } + + +@dataclass +class HookDrift: + """Per-hook drift record vs the unhooked baseline. + + ``check`` is the shared-lib :class:`~tests.lib.compare.Check` for the + logits comparison (carries passed / max_abs / mean_abs / first_diff_pos / + detail). The extra fields capture what the study adds on top. + """ + + hook: str + check: Optional[Check] = None + n_positions: int = 0 + vocab_size: int = 0 + first_diff_pos: int = -1 # token position; -1 == none + token_ids_diverged: bool = False + n_token_diff: int = 0 + topk_vocab_diffs: List[VocabDiff] = field(default_factory=list) + shape_mismatch: bool = False + error: Optional[str] = None + alert: bool = False + alert_reasons: List[str] = field(default_factory=list) + + def to_dict(self) -> Dict[str, Any]: + return { + "hook": self.hook, + "check": self.check.to_dict() if self.check is not None else None, + "n_positions": self.n_positions, + "vocab_size": self.vocab_size, + "first_diff_pos": self.first_diff_pos, + "token_ids_diverged": self.token_ids_diverged, + "n_token_diff": self.n_token_diff, + "topk_vocab_diffs": [v.to_dict() for v in self.topk_vocab_diffs], + "shape_mismatch": self.shape_mismatch, + "error": self.error, + "alert": self.alert, + "alert_reasons": self.alert_reasons, + } + + +@dataclass +class StudyResult: + framework: str + model: str + mode: str + variant: str + dtype: str + standard: str + threshold: float + topk: int + hooks: List[HookDrift] = field(default_factory=list) + + @property + def any_alert(self) -> bool: + return any(h.alert for h in self.hooks) + + def to_dict(self) -> Dict[str, Any]: + return { + "framework": self.framework, + "model": self.model, + "mode": self.mode, + "variant": self.variant, + "dtype": self.dtype, + "standard": self.standard, + "threshold": self.threshold, + "topk": self.topk, + "any_alert": self.any_alert, + "hooks": [h.to_dict() for h in self.hooks], + } + + +# --------------------------------------------------------------------------- +# Pure comparison core (CPU-testable; torch imported lazily) +# --------------------------------------------------------------------------- + + +def _first_diff_row(abs_diff: "Any") -> int: + """Index of the first token position holding any nonzero diff, or -1. + + ``abs_diff`` is a ``[N, vocab]`` tensor of absolute differences. + """ + if abs_diff.numel() == 0: + return -1 + row_has_diff = (abs_diff > 0).any(dim=-1) + nz = row_has_diff.nonzero(as_tuple=False) + return int(nz[0].item()) if nz.numel() > 0 else -1 + + +def compute_drift( + hook: str, + baseline: Dict[str, "Any"], + monitored: Dict[str, "Any"], + *, + mode: str, + threshold: float, + topk: int = 5, + sparse_floor: Optional[float] = None, +) -> HookDrift: + """Compare one monitored capture against the unhooked baseline. + + ``baseline`` / ``monitored`` are dicts with ``token_ids`` (int64 ``[N]``) + and ``logits`` (float ``[N, vocab]``). Uses the shared-lib standard for + ``mode`` (``bitwise`` eager / ``allclose`` cuda graph) for the verdict and + augments it with token-divergence + top-k vocab drift. The §9 alert + policy is applied here. Never raises on bad input: records ``error`` / + ``shape_mismatch`` so the caller can alert. + + ``sparse_floor``: when set (e.g. ``-1e30`` for vLLM top-k logprob captures), + drift metrics (max_abs, mean_abs, first_diff_pos, top-k vocab) are computed + only over entries where at least one tensor has a value above the floor, + preventing the floor fill from dominating the statistics. Greedy-token + argmax is always computed from the raw logits (the floor is too low to win). + """ + import torch + + drift = HookDrift(hook=hook) + + b_logits = baseline.get("logits") + m_logits = monitored.get("logits") + if b_logits is None or m_logits is None: + drift.error = "missing logits in capture" + return _finalize_alert(drift, mode=mode) + + b_logits = b_logits.float() + m_logits = m_logits.float() + drift.n_positions = int(b_logits.shape[0]) if b_logits.ndim >= 1 else 0 + drift.vocab_size = int(b_logits.shape[-1]) if b_logits.ndim >= 1 else 0 + + if drift.n_positions == 0 or m_logits.shape[0] == 0: + drift.error = "empty capture (no positions)" + return _finalize_alert(drift, mode=mode) + + if tuple(b_logits.shape) != tuple(m_logits.shape): + drift.shape_mismatch = True + drift.error = ( + f"logits shape mismatch: baseline {tuple(b_logits.shape)} " + f"vs monitored {tuple(m_logits.shape)}" + ) + # Still run the shared standard so the Check records the mismatch. + drift.check = _run_standard(b_logits, m_logits, mode=mode, threshold=threshold) + return _finalize_alert(drift, mode=mode) + + abs_diff = (b_logits - m_logits).abs() + if sparse_floor is not None: + # Restrict to entries where at least one runner reported a real logit; + # floor-vs-floor pairs contribute zero so they don't inflate max/mean. + real_mask = (b_logits > sparse_floor * 0.5) | (m_logits > sparse_floor * 0.5) + abs_diff = abs_diff.where(real_mask, torch.zeros_like(abs_diff)) + + # Shared-lib verdict (bitwise / allclose) -- carries max/mean/first stats. + drift.check = _run_standard(b_logits, m_logits, mode=mode, threshold=threshold, + masked_abs_diff=abs_diff if sparse_floor is not None else None) + + drift.first_diff_pos = _first_diff_row(abs_diff) + + # Greedy (argmax) token divergence, recomputed from logits and + # cross-checked against the stored ids. + b_arg = b_logits.argmax(dim=-1) + m_arg = m_logits.argmax(dim=-1) + drift.n_token_diff = int((b_arg != m_arg).sum().item()) + drift.token_ids_diverged = drift.n_token_diff > 0 + b_ids = baseline.get("token_ids") + m_ids = monitored.get("token_ids") + if b_ids is not None and m_ids is not None and b_ids.shape == m_ids.shape: + if not torch.equal(b_ids, m_ids): + drift.token_ids_diverged = True + drift.n_token_diff = max(drift.n_token_diff, int((b_ids != m_ids).sum().item())) + + # Top-k vocab diffs at the first differing position. + pos = drift.first_diff_pos + if pos >= 0 and topk > 0: + row = abs_diff[pos] + k = min(topk, int(row.numel())) + top = torch.topk(row, k) + for rank in range(k): + tid = int(top.indices[rank].item()) + drift.topk_vocab_diffs.append( + VocabDiff( + token_id=tid, + baseline_logit=float(b_logits[pos, tid].item()), + monitored_logit=float(m_logits[pos, tid].item()), + abs_diff=float(top.values[rank].item()), + ) + ) + + return _finalize_alert(drift, mode=mode) + + +def _run_standard(b_logits: "Any", m_logits: "Any", *, mode: str, threshold: float, + masked_abs_diff: "Optional[Any]" = None) -> "Check": + """Run the §8 standard for ``mode`` and return its :class:`Check`. + + When ``masked_abs_diff`` is supplied (pre-computed, sparse-masked), the + Check is built directly from it instead of recomputing from the raw tensors + -- ensuring metrics reflect only the meaningful (non-floor) entries. + """ + name = standard_for_mode(mode) + if masked_abs_diff is None: + if name == "bitwise": + return bitwise(b_logits, m_logits, name="logits_bitwise") + return allclose(b_logits, m_logits, name="logits_allclose", atol=threshold, rtol=0.0) + + # Sparse path: derive stats from the already-masked diff tensor. + import torch as _torch + ad = masked_abs_diff.float() + n = ad.numel() + max_v = float(ad.max().item()) if n > 0 else 0.0 + mean_v = float(ad.mean().item()) if n > 0 else 0.0 + ne = (ad > 0).any(dim=-1).nonzero(as_tuple=False) if ad.ndim >= 2 else (ad > 0).nonzero(as_tuple=False) + first: Optional[int] = int(ne[0].item()) if ne.numel() > 0 else None + if name == "bitwise": + passed = max_v == 0.0 + return Check("logits_bitwise", passed, max_abs=max_v, mean_abs=mean_v, + first_diff_pos=None if passed else first, + detail="bitwise equal" if passed else f"max_abs={max_v:.6e}") + passed = max_v <= threshold + return Check("logits_allclose", passed, max_abs=max_v, mean_abs=mean_v, + first_diff_pos=None if passed else first, + detail=f"max_abs={max_v:.6e} (atol={threshold:g}, rtol=0)") + + +def _finalize_alert(drift: HookDrift, *, mode: str) -> HookDrift: + """Apply the §9 alert policy to a populated :class:`HookDrift` in place.""" + reasons: List[str] = [] + if drift.error is not None: + reasons.append(f"capture error: {drift.error}") + if drift.shape_mismatch: + reasons.append("logits shape mismatch (possible hook identity swap)") + if drift.token_ids_diverged: + reasons.append(f"greedy token ids diverged at {drift.n_token_diff} position(s)") + if drift.check is not None and not drift.check.passed and not drift.shape_mismatch: + if mode == "eager": + reasons.append(f"non-bitwise drift in eager mode ({drift.check.detail})") + else: + reasons.append(f"exceeds cuda-graph threshold ({drift.check.detail})") + drift.alert_reasons = reasons + drift.alert = bool(reasons) + return drift + + +# --------------------------------------------------------------------------- +# Human-readable table +# --------------------------------------------------------------------------- + + +def format_table(result: StudyResult) -> str: + """Render a fixed-width human-readable table for the study result.""" + header = ( + f"numeric-difference study: {result.framework}/{result.model} " + f"mode={result.mode} variant={result.variant} dtype={result.dtype} " + f"standard={result.standard}" + + (f" (threshold={result.threshold:g})" if result.standard == "allclose" else "") + ) + cols = ("hook", "verdict", "max_abs", "mean_abs", "first_diff", "tok_div", "alert") + widths = (16, 8, 12, 12, 10, 8, 6) + sep = " " + + def _row(values: tuple) -> str: + return sep.join(str(v).ljust(w) for v, w in zip(values, widths)) + + lines = [header, "-" * len(header), _row(cols)] + for h in result.hooks: + if h.error is not None and h.check is None: + verdict, mx, mn, fd = "ERR", "-", "-", "-" + else: + chk = h.check + verdict = "pass" if (chk and chk.passed) else "fail" + mx = f"{chk.max_abs:.4g}" if chk and chk.max_abs is not None else "-" + mn = f"{chk.mean_abs:.4g}" if chk and chk.mean_abs is not None else "-" + fd = str(h.first_diff_pos) if h.first_diff_pos >= 0 else "-" + lines.append( + _row(( + h.hook, verdict, mx, mn, fd, + str(h.n_token_diff) if h.token_ids_diverged else "0", + "ALERT" if h.alert else "ok", + )) + ) + for vd in h.topk_vocab_diffs: + lines.append( + sep + f" tok {vd.token_id}: base={vd.baseline_logit:.4g} " + f"mon={vd.monitored_logit:.4g} |Δ|={vd.abs_diff:.4g}" + ) + if h.alert: + for reason in h.alert_reasons: + lines.append(sep + f" ! {reason}") + n_alert = sum(1 for h in result.hooks if h.alert) + lines.append("") + lines.append(f"{len(result.hooks) - n_alert}/{len(result.hooks)} hooks clean" + + (f", {n_alert} alerting" if n_alert else "")) + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Subprocess rollout runners +# --------------------------------------------------------------------------- + +# HF rollout saves {token_ids: int64[N], logits: float32[N, vocab]} -- full vocabulary. +# vLLM rollout saves the same shape, but only top-k logprob entries are real; all +# other vocab positions are filled to -1e30. Drift metrics must be restricted to +# entries where at least one runner reported a real logit (see compute_drift sparse_floor). + +_HF_RUNNER = dedent(""" + import argparse, os + import torch + + ap = argparse.ArgumentParser() + ap.add_argument('--model-key', required=True) + ap.add_argument('--hook', required=True) + ap.add_argument('--mode', required=True) + ap.add_argument('--variant', required=True, choices=['p', 'compare']) + ap.add_argument('--rollout', required=True, choices=['baseline', 'hooked']) + ap.add_argument('--max-new-tokens', type=int, default=4) + ap.add_argument('--prompt', default='Hello') + ap.add_argument('--dtype', default='float16') + ap.add_argument('--out', required=True) + args = ap.parse_args() + + MODEL_ALIASES = { + 'gpt2': 'gpt2', + 'qwen3': 'Qwen/Qwen3-0.6B', + 'qwen2_moe': 'Qwen/Qwen1.5-MoE-A2.7B', + 'llama': 'meta-llama/Llama-3.1-8B', + } + hf_id = MODEL_ALIASES[args.model_key] + device = torch.device('cuda') + _DTYPE_MAP = {'float16': torch.float16, 'bfloat16': torch.bfloat16, 'float32': torch.float32} + dtype = _DTYPE_MAP.get(args.dtype, torch.float16) + + from transformers import AutoTokenizer + tok = AutoTokenizer.from_pretrained(hf_id) + if tok.pad_token_id is None: + tok.pad_token_id = tok.eos_token_id + + if args.rollout == 'baseline': + from transformers import AutoModelForCausalLM + model = AutoModelForCausalLM.from_pretrained( + hf_id, torch_dtype=dtype, attn_implementation='eager' + ).to(device).eval() + elif args.variant == 'compare': + # Patched _compare model (driver wraps this subprocess in isolated_hook). + if args.model_key == 'qwen3': + from transformers.models.qwen3_compare.modeling_qwen3 import CompareQwen3ForCausalLM as cls + elif args.model_key == 'gpt2': + from transformers.models.gpt2_compare.modeling_gpt2 import CompareGPT2LMHeadModel as cls + elif args.model_key == 'qwen2_moe': + from transformers.models.qwen2_moe_compare.modeling_qwen2_moe import CompareQwen2MoeForCausalLM as cls + elif args.model_key == 'llama': + from transformers.models.llama_compare.modeling_llama import CompareLlamaForCausalLM as cls + else: + raise ValueError(f'unsupported model_key={args.model_key!r}') + model = cls.from_pretrained( + hf_id, torch_dtype=dtype, attn_implementation='eager' + ).to(device).eval() + model.allocate_compare_buffers(1, 32, dtype=dtype, tp_size=1) + else: # variant p: production _p Hooked variant + hook_selection=H + if args.model_key == 'qwen3': + from transformers.models.qwen3_p.modeling_qwen3 import HookedQwen3ForCausalLM as cls + elif args.model_key == 'gpt2': + from transformers.models.gpt2_p.modeling_gpt2 import HookedGPT2LMHeadModel as cls + elif args.model_key == 'qwen2_moe': + from transformers.models.qwen2_moe_p.modeling_qwen2_moe import HookedQwen2MoeForCausalLM as cls + else: + raise ValueError(f'unsupported model_key={args.model_key!r}') + model = cls.from_pretrained( + hf_id, torch_dtype=dtype, attn_implementation='eager' + ).to(device).eval() + + inputs = tok([args.prompt], return_tensors='pt', padding=True).to(device) + gen_kwargs = dict( + **inputs, + max_new_tokens=args.max_new_tokens, + do_sample=False, + pad_token_id=tok.pad_token_id, + return_dict_in_generate=True, + output_scores=True, + ) + if args.mode == 'cuda_graph': + from transformers import CompileConfig + gen_kwargs['cache_implementation'] = 'static' + gen_kwargs['compile_config'] = CompileConfig(mode='reduce-overhead', fullgraph=False) + + if args.rollout == 'hooked' and args.variant == 'p': + from monitoring import MonitoringEngine, MonitoringConfig + from monitoring.config import CaptureSchedule + from monitoring._native_engine import RingConfig + from integration.hf_adapter import generate_with_monitoring + cfg = MonitoringConfig(schedule=CaptureSchedule(capture_prefill=True, capture_decode=True)) + engine = MonitoringEngine(config=cfg, model_id='numeric_study') + ring_cfg = RingConfig() + ring_cfg.task_ring_entries = 1024 + ring_cfg.payload_ring_bytes = 64 * 1024 * 1024 + ring_cfg.pinned_staging_bytes = 64 * 1024 * 1024 + engine.enable_ring_transport(ring_cfg) + model.monitoring_engine = engine + try: + out = generate_with_monitoring(model, hook_selection=args.hook, **gen_kwargs) + finally: + engine.close() + else: + with torch.no_grad(): + out = model.generate(**gen_kwargs) + + scores = torch.stack(out.scores, dim=0) # [N, 1, vocab] + logits = scores.squeeze(1).float().cpu() # [N, vocab] + token_ids = logits.argmax(dim=-1).to(torch.int64) # [N] + torch.save({'token_ids': token_ids, 'logits': logits}, args.out) + print(f'OK {args.rollout}/{args.variant} N={logits.shape[0]} V={logits.shape[1]} -> {args.out}') +""") + + +_VLLM_RUNNER = dedent(""" + import argparse, os + os.environ.setdefault('VLLM_DISABLE_COMPILE_CACHE', '1') + import torch + from vllm import LLM, SamplingParams + + ap = argparse.ArgumentParser() + ap.add_argument('--model-key', required=True) + ap.add_argument('--hook', required=True) + ap.add_argument('--mode', required=True) + ap.add_argument('--variant', required=True, choices=['p', 'compare']) + ap.add_argument('--rollout', required=True, choices=['baseline', 'hooked']) + ap.add_argument('--max-new-tokens', type=int, default=4) + ap.add_argument('--prompt', default='Hello') + ap.add_argument('--dtype', default='float16') + ap.add_argument('--out', required=True) + args = ap.parse_args() + + MODEL_ALIASES = { + 'gpt2': 'gpt2', + 'qwen3': 'Qwen/Qwen3-0.6B', + 'qwen2_moe': 'Qwen/Qwen1.5-MoE-A2.7B', + } + model_name = MODEL_ALIASES[args.model_key] + + llm_kwargs = dict( + model=model_name, + max_model_len=128, + gpu_memory_utilization=0.5, + enforce_eager=(args.mode == 'eager'), + dtype=args.dtype, + ) + if args.rollout == 'hooked': + additional_config = {'dmx_hook_selection': args.hook, 'dmx_db_host': ''} + if args.variant == 'compare': + llm_kwargs['worker_cls'] = 'tests.compare_worker.CompareWorker' + else: + llm_kwargs['worker_cls'] = 'integration.vllm_adapter.DMXGPUWorker' + llm_kwargs['additional_config'] = additional_config + + llm = LLM(**llm_kwargs) + params = SamplingParams(temperature=0.0, max_tokens=args.max_new_tokens, logprobs=20) + outputs = llm.generate([args.prompt], params) + + completion = outputs[0].outputs[0] + ids = list(completion.token_ids) + step_logprobs = completion.logprobs or [] + # Dense [N, V] from the sparse top-k logprob dicts; unreported entries stay + # at a large negative floor (sufficient for chosen-token + top-k drift). + vocab = int(getattr(llm.llm_engine.model_config.hf_config, 'vocab_size', 0)) or 1 + N = len(ids) + logits = torch.full((N, vocab), -1e30, dtype=torch.float32) + for i in range(N): + if i < len(step_logprobs) and step_logprobs[i] is not None: + for tid, lp in step_logprobs[i].items(): + logits[i, int(tid)] = float(lp.logprob) + token_ids = torch.tensor(ids, dtype=torch.int64) + torch.save({'token_ids': token_ids, 'logits': logits}, args.out) + print(f'OK {args.rollout}/{args.variant} N={N} V={vocab} -> {args.out}') + + try: + llm.collective_rpc('stop_monitoring') + except Exception: + pass +""") + + +def _build_subprocess_env() -> dict: + """Pin CUDA_VISIBLE_DEVICES=0 and put conda lib on LD_LIBRARY_PATH so vLLM + imports resolve (mirrors test_per_hook_isolation / test_no_graph_breaks).""" + env = os.environ.copy() + conda_prefix = env.get("CONDA_PREFIX") + if conda_prefix: + ld = env.get("LD_LIBRARY_PATH", "") + env["LD_LIBRARY_PATH"] = f"{conda_prefix}/lib:{ld}" if ld else f"{conda_prefix}/lib" + env.setdefault("CUDA_VISIBLE_DEVICES", "0") + return env + + +def _run_rollout( + *, + framework: str, + model_key: str, + hook: str, + mode: str, + variant: str, + rollout: str, + out_path: Path, + max_new_tokens: int, + prompt: str, + env: dict, + dtype: str = "float16", + timeout: int = 600, +) -> None: + """Spawn one rollout subprocess; raise with captured output on failure.""" + runner = _HF_RUNNER if framework == "hf" else _VLLM_RUNNER + cmd = [ + sys.executable, "-c", runner, + "--model-key", model_key, + "--hook", hook, + "--mode", mode, + "--variant", variant, + "--rollout", rollout, + "--max-new-tokens", str(max_new_tokens), + "--prompt", prompt, + "--dtype", dtype, + "--out", str(out_path), + ] + proc = subprocess.run( + cmd, capture_output=True, text=True, timeout=timeout, env=env, cwd=REPO_ROOT, + ) + if proc.returncode != 0: + raise RuntimeError( + f"rollout={rollout} variant={variant} hook={hook} failed " + f"(rc={proc.returncode})\n" + f"--- stdout ---\n{proc.stdout}\n" + f"--- stderr (tail) ---\n{proc.stderr[-3000:]}" + ) + + +# --------------------------------------------------------------------------- +# Orchestration +# --------------------------------------------------------------------------- + + +def run_study( + *, + framework: str, + model: str, + mode: str, + hooks: List[str], + out_dir: Path, + variant: str = "p", + dtype: str = "float16", + topk: int = 5, + max_new_tokens: int = 4, + prompt: str = "Hello", + threshold: Optional[float] = None, +) -> StudyResult: + """Run the full per-hook study and return a populated :class:`StudyResult`. + + Captures the unhooked baseline once, then one monitored rollout per hook, + comparing each against the baseline and applying the §9 alert policy. For + ``variant='compare'`` each hooked rollout runs inside the hardened + :func:`tests.isolate_hook.isolated_hook` context manager (plan §6). + """ + import torch + + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + env = _build_subprocess_env() + thr = threshold if threshold is not None else cuda_graph_threshold(model, dtype) + + result = StudyResult( + framework=framework, model=model, mode=mode, variant=variant, + dtype=dtype, standard=standard_for_mode(mode), threshold=thr, topk=topk, + ) + + # vLLM saves top-k logprobs; restrict diff metrics to reported entries only. + sparse_floor: Optional[float] = -1e30 if framework == "vllm" else None + + # 1. Baseline (unhooked) captured once; hook arg is unused by the runner. + baseline_path = out_dir / "baseline.pt" + _run_rollout( + framework=framework, model_key=model, hook=hooks[0] if hooks else "q", + mode=mode, variant="p", rollout="baseline", out_path=baseline_path, + max_new_tokens=max_new_tokens, prompt=prompt, env=env, dtype=dtype, + ) + baseline = torch.load(baseline_path, map_location="cpu") + + # 2-3. One hook at a time; compare; apply alert policy. + for hook in hooks: + drift = HookDrift(hook=hook) + try: + hooked_path = out_dir / f"hooked_{hook}.pt" + if variant == "compare": + from tests.isolate_hook import isolated_hook + + with isolated_hook(framework, model, hook): + _run_rollout( + framework=framework, model_key=model, hook=hook, mode=mode, + variant="compare", rollout="hooked", out_path=hooked_path, + max_new_tokens=max_new_tokens, prompt=prompt, env=env, dtype=dtype, + ) + else: + _run_rollout( + framework=framework, model_key=model, hook=hook, mode=mode, + variant="p", rollout="hooked", out_path=hooked_path, + max_new_tokens=max_new_tokens, prompt=prompt, env=env, dtype=dtype, + ) + monitored = torch.load(hooked_path, map_location="cpu") + drift = compute_drift( + hook, baseline, monitored, mode=mode, threshold=thr, topk=topk, + sparse_floor=sparse_floor, + ) + except Exception as exc: # subprocess crash / OOM / dirty restore -> alert + drift.error = f"{type(exc).__name__}: {exc}" + _finalize_alert(drift, mode=mode) + result.hooks.append(drift) + + return result + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def main(argv: Optional[List[str]] = None) -> int: + ap = argparse.ArgumentParser( + description="Per-hook numeric-difference study vs the unhooked baseline." + ) + ap.add_argument("--framework", choices=["hf", "vllm"], default="hf") + ap.add_argument("--model", default="qwen3") + ap.add_argument("--mode", choices=["eager", "cuda_graph"], default="eager") + ap.add_argument("--variant", choices=["p", "compare"], default="p", + help="p: production _p+hook_selection; compare: _compare under isolated_hook") + ap.add_argument("--hooks", default=",".join(_DEFAULT_HOOKS), + help="Comma-separated hook short-names (e.g. q,k,resid_pre,final_logits).") + ap.add_argument("--dtype", default="float16") + ap.add_argument("--topk", type=int, default=5) + ap.add_argument("--max-new-tokens", type=int, default=4) + ap.add_argument("--prompt", default="Hello") + ap.add_argument("--threshold", type=float, default=None, + help="Override the CUDA-graph max-abs-diff alert threshold.") + ap.add_argument("--work-dir", default=None, + help="Scratch dir for per-rollout .pt files (default: temp dir).") + ap.add_argument("--out", default=None, + help="Write the JSON artifact here (default: stdout table only).") + args = ap.parse_args(argv) + + hooks = [h.strip() for h in args.hooks.split(",") if h.strip()] + if not hooks: + ap.error("no hooks selected") + + import tempfile + + tmp_ctx = None + if args.work_dir: + work_dir = Path(args.work_dir) + else: + tmp_ctx = tempfile.TemporaryDirectory(prefix="numeric_study_") + work_dir = Path(tmp_ctx.name) + + try: + result = run_study( + framework=args.framework, model=args.model, mode=args.mode, + hooks=hooks, out_dir=work_dir, variant=args.variant, dtype=args.dtype, + topk=args.topk, max_new_tokens=args.max_new_tokens, prompt=args.prompt, + threshold=args.threshold, + ) + finally: + if tmp_ctx is not None: + tmp_ctx.cleanup() + + print(format_table(result)) + if args.out: + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text(json.dumps(result.to_dict(), indent=2), encoding="utf-8") + print(f"\nwrote JSON artifact -> {out_path}") + + # Non-zero exit if any hook alerted, so CI / wrappers can gate on it. + return 1 if result.any_alert else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_e2e_correctness_vs_hf.py b/tests/test_e2e_correctness_vs_hf.py index 2ae462a94..abcb6a761 100644 --- a/tests/test_e2e_correctness_vs_hf.py +++ b/tests/test_e2e_correctness_vs_hf.py @@ -1,80 +1,28 @@ -# tests/test_e2e_correctness_hf.py -# PYTHONPATH=./:./monitoring:$PYTHONPATH E2E_PRINT_TEXT=1 E2E_HF_DROP_LAST_TOKEN=1 E2E_PRINT_TOPK_LOGITS=1 pytest -q -s tests/test_e2e_correctness_vs_hf.py -"""E2E correctness test: monitoring DB vs HuggingFace Transformers (HF-driven ground truth). - -This test runs the repo monitoring pipeline end-to-end (native backend + host engine + ClickHouse), -then uses HuggingFace Transformers as the reference implementation (no TransformerLens). - -IMPORTANT: "prompt" in this test means FULL TOKEN SEQUENCE = prefill + decode ------------------------------------------------------------------------ -We treat DB `token_ids` as the ground-truth sequence of tokens for each request. That sequence -includes the initial prompt tokens (prefill) plus any decode tokens that were appended. - -HF reference modes ------------------------------------------------------------------------ - - ROL (manual rollout, batched): - * Incremental greedy KV-cache rollout over the full padded batch, using generate-style - position_ids derived from attention_mask (so left-padding doesn't shift positions). - * We strip left-pad per row and stop per row at EOS (no trailing padded steps), so outputs - align to DB token_ids. - - - GEN (HF generate(), batched): - * Run hf_model.generate() ONCE on the full padded batch (same input_ids/attention_mask as monitoring), - with output_scores=True. - * Strip left-pad per row and trim at EOS so sequences align to DB token_ids. - -For logits we effectively compare THREE sources: - - DB final_logits (from ClickHouse) - - ROL final_logits (manual rollout; [T, vocab]) - - GEN scores (from generate(); available only for positions t in [prompt_len-1, T-2]) - -Request-id convention (from MonitoringEngine._register_db_step in engine.py) ------------------------------------------------------------------------ -When a new batch is (re)initialized: - - gid = self._auto_batch_group_id - self._auto_batch_group_id += 1 - self._active_batch_request_ids = [f"{gid}:{i}" for i in range(batch_size)] - -So request_id == ":" where local_index is the batch row index. - -We map DB requests to HF batch rows using local_index, and add safety checks. - -DB tensor shapes (IMPORTANT) ------------------------------------------------------------------------ -DB offloaded tensors DO NOT have batch dim now: - - token_ids: [T] - - hook_embed/pos/final_ln: [T, d_model] - - resid_pre/post: [T, d_model] - - attn pattern/scores: [n_heads, Tq, Tk] - - final_logits: [R, vocab] or [vocab] (R is often full sequence length) - -Env vars ------------------------------------------------------------------------ - - E2E_BATCH_SIZE (default 4) - - E2E_MAX_NEW_TOKENS (default 8) - - E2E_MODEL (default "gpt2"; "qwen3" alias supported) - - E2E_CHUNK_BYTES (default 262144) - - - E2E_PRINT_TEXT (default 0): if 1, print decoded text from DB token_ids and from HF rollout + HF generate(). - - E2E_HF_DROP_LAST_TOKEN (default 0): if 1, drop the last token (and aligned tensors) from HF refs before compares. - - E2E_PRINT_TOPK_LOGITS (default 0): if 1, print top-k logits at every position for DB vs ROL vs GEN. - - E2E_PRINT_TOPK_LOGITS_K (default 5): top-k to print per position. - -ClickHouse ------------------------------------------------------------------------ - - DMX_DB_HOST, DMX_DB_PORT, DMX_DB_USER, DMX_DB_PASSWORD, DMX_DB_DATABASE, DMX_DB_TABLE +"""HF E2E correctness — thin wrappers over the configurable matrix (plan §5). + +This file used to carry ~1.6k lines of in-process HF rollout + tensor +comparison logic, three tests (one permanently disabled via +``@skipif(True)``), 16 skip sites, and two ``_legacy`` bodies kept "for +reference". All of that comparison logic now lives in :mod:`tests.lib` and +the dispatch in :mod:`tests.e2e_matrix`; these wrappers just drive the +matrix for the equivalent HF cell and assert on its checks. + +The test *names* are preserved because ``tests/tools/verify_hf.sh`` invokes +them by node id (``::test_e2e_correctness_hf`` / +``::test_e2e_cuda_graphs_vs_eager_hf``) and threads the ring-size / model +env vars the matrix wrapper reads. + +The removed ``test_e2e_correctness_hf_cuda_graphs`` was permanently disabled +(its compiled-rollout reference could not replicate generate()'s internal +StaticCache handling) and explicitly superseded by +``test_e2e_cuda_graphs_vs_eager_hf`` -- no still-passing assertion was lost. """ - from __future__ import annotations -import os -import sys -import uuid -from typing import Dict, List, Tuple - import pytest -import torch + +from tests._requirements import require_cuda, require_clickhouse +from tests.e2e_matrix import matrix_argv_from_env, run_single pytestmark = [ pytest.mark.gpu, @@ -206,1085 +154,18 @@ def _make_host_cfg(db_cfg_native): @pytest.mark.skipif(not torch.backends.cuda.is_built(), reason="CUDA not built") def test_e2e_correctness_hf(subtests) -> None: - """E2E correctness: compare HOOKED model (ring transport -> ClickHouse) - against ORIGINAL model (HF output_hidden_states=True). + """HF eager: hooked model (ring -> ClickHouse) vs original model. - Three subprocesses — parent process never touches CUDA: - 1. Reference: original model -> tensors on disk - 2. Monitored: hooked model + ring transport -> ClickHouse - 3. Comparator: reads both, compares, writes result.json + Equivalent matrix cell: ``--backend hf --mode eager --standard allclose`` + (HF dispatches to hf_comparator, which does the value comparison with + ``E2E_TOLERANCE``; default 0.01 for eager). """ - import json - import subprocess - import tempfile - import shutil - - run_dir = tempfile.mkdtemp(prefix="hf_e2e_") - ref_dir = os.path.join(run_dir, "ref") - mon_dir = os.path.join(run_dir, "mon") - result_file = os.path.join(run_dir, "result.json") - project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - - try: - # Step 1: Reference run (original model, no CUDA in parent) - print("\n [1/3] Reference run (original model)...", flush=True) - r1 = subprocess.run( - [sys.executable, "-m", "tests.hf_reference_runner", - "--output-dir", ref_dir], - env=os.environ, capture_output=True, text=True, cwd=project_root, - ) - if r1.returncode != 0: - pytest.fail(f"Reference runner failed:\n{r1.stderr[-2000:]}") - - # Step 2: Monitored run (hooked model + ring transport) - print(" [2/3] Monitored run (hooked model + ring)...", flush=True) - r2 = subprocess.run( - [sys.executable, "-m", "tests.hf_monitored_runner", - "--output-dir", mon_dir], - env=os.environ, capture_output=True, text=True, cwd=project_root, - ) - if r2.returncode != 0: - pytest.fail(f"Monitored runner failed:\n{r2.stderr[-2000:]}") - - # Step 3: Comparator (CPU only, reads disk + ClickHouse) - print(" [3/3] Comparing...", flush=True) - r3 = subprocess.run( - [sys.executable, "-m", "tests.hf_comparator", - "--ref-dir", ref_dir, - "--mon-dir", mon_dir, - "--result-file", result_file], - env=os.environ, capture_output=True, text=True, cwd=project_root, - ) - if r3.returncode != 0: - pytest.fail(f"Comparator failed:\n{r3.stderr[-2000:]}") - - # Read results - with open(result_file) as f: - results = json.load(f) - - # Report via subtests - for test in results["tests"]: - with subtests.test(test["name"]): - assert test["passed"], test.get("detail", "") - - finally: - shutil.rmtree(run_dir, ignore_errors=True) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA + native backend required") -def _test_e2e_correctness_hf_legacy(subtests) -> None: - """Legacy version kept for reference. Not called by verify_hf.sh.""" - try: - import clickhouse_driver # noqa: F401 - except Exception: - pytest.skip("clickhouse-driver is required") - - try: - from monitoring import ( # type: ignore - MonitoringConfig, - MonitoringEngine, - ) - from monitoring._native_engine import ClickHouseClientConfig # type: ignore - from monitoring.config import CaptureSchedule # type: ignore - from integration.hf_adapter import generate_with_monitoring # type: ignore - except Exception as exc: - pytest.skip(f"monitoring native extension not available: {exc}") - - try: - from transformers import AutoModelForCausalLM, AutoTokenizer - from transformers.models.gpt2_p.modeling_gpt2 import HookedGPT2LMHeadModel # type: ignore - from transformers.models.qwen3_p.modeling_qwen3 import HookedQwen3ForCausalLM # type: ignore - from transformers.models.llama_p.modeling_llama import HookedLlamaForCausalLM # type: ignore - except Exception as exc: - pytest.skip(f"transformers or repo Hooked* classes not available: {exc}") - - # ----------------------------------------------------------------------- - # Configuration - # ----------------------------------------------------------------------- - - batch_size = int(os.environ.get("E2E_BATCH_SIZE", "4")) - if batch_size < 1: - raise ValueError("E2E_BATCH_SIZE must be >= 1") - - max_new_tokens = int(os.environ.get("E2E_MAX_NEW_TOKENS", "8")) - hf_model_id = _resolve_model_id(os.environ.get("E2E_MODEL", "gpt2")) - chunk_bytes = int(os.environ.get("E2E_CHUNK_BYTES", str(256 * 1024))) - - print_text = int(os.environ.get("E2E_PRINT_TEXT", "0")) == 1 - # E2E_HF_DROP_LAST_TOKEN is no longer needed: both monitored and HF reference - # use generate(), so they produce the same number of tokens/hidden states. - print_topk_logits = int(os.environ.get("E2E_PRINT_TOPK_LOGITS", "0")) == 1 - topk_k = int(os.environ.get("E2E_PRINT_TOPK_LOGITS_K", "5")) - if topk_k < 1: - raise ValueError("E2E_PRINT_TOPK_LOGITS_K must be >= 1") - - device = torch.device("cuda") - - # ----------------------------------------------------------------------- - # Tokenizer + prompts - # ----------------------------------------------------------------------- - - tokenizer = AutoTokenizer.from_pretrained(hf_model_id) - if tokenizer.pad_token_id is None: - tokenizer.pad_token_id = tokenizer.eos_token_id - tokenizer.padding_side = "left" - eos_id = int(tokenizer.eos_token_id) - pad_id = int(tokenizer.pad_token_id) - - prompts = [("Hello " * (i + 1)).strip() for i in range(batch_size)] - encoded = tokenizer(prompts, return_tensors="pt", padding=True) - input_ids = encoded["input_ids"].to(device) - attention_mask = encoded["attention_mask"].to(device) - - # Initial prompt tokens for safety prefix checks - hf_initial_prompt_tokens: List[torch.Tensor] = [] - for j in range(batch_size): - hf_initial_prompt_tokens.append( - _strip_left_pad( - input_ids[j].detach().cpu(), - attention_mask[j].detach().cpu(), - ).to(torch.long) - ) - - # ----------------------------------------------------------------------- - # Monitoring config (config-driven; no env var toggles) - # ----------------------------------------------------------------------- - - mon_cfg = MonitoringConfig( - schedule=CaptureSchedule(capture_prefill=True, capture_decode=True), - ) - - # ----------------------------------------------------------------------- - # ClickHouse config (for the monitored run) - # ----------------------------------------------------------------------- - - db_cfg_native = ClickHouseClientConfig() - db_cfg_native.host = os.environ.get("DMX_DB_HOST", "localhost") - db_cfg_native.port = int(os.environ.get("DMX_DB_PORT", "9000")) - db_cfg_native.username = os.environ.get("DMX_DB_USER", "default") - db_cfg_native.password = os.environ.get("DMX_DB_PASSWORD", "") - db_cfg_native.database = os.environ.get("DMX_DB_DATABASE", "default") - db_cfg_native.table = os.environ.get("DMX_DB_TABLE", "offload") - db_cfg_native.secure = False - db_cfg_native.client_side_compress = "none" - db_cfg_native.client_settings = None - db_cfg_native.create_database_if_missing = True - db_cfg_native.drop_existing_database = True - db_cfg_native.index_granularity = 8192 - - host_cfg = _make_host_cfg(db_cfg_native) - ring_cfg = _make_ring_cfg() - - # ----------------------------------------------------------------------- - # Monitored run - # ----------------------------------------------------------------------- - - unique_run_model_id = f"e2e_correctness_hf::{uuid.uuid4().hex}"[:120] - engine = MonitoringEngine( - config=mon_cfg, model_id=unique_run_model_id, db_config=host_cfg - ) - engine.enable_ring_transport(ring_cfg) - - if "qwen3" in hf_model_id.lower(): - model_cls = HookedQwen3ForCausalLM - elif "llama" in hf_model_id.lower(): - model_cls = HookedLlamaForCausalLM - else: - model_cls = HookedGPT2LMHeadModel - mon_model = model_cls.from_pretrained(hf_model_id, attn_implementation="eager", torch_dtype=torch.float16) - mon_model.to(device).eval() - mon_model.monitoring_engine = engine - - try: - with torch.no_grad(): - _ = generate_with_monitoring( - mon_model, - input_ids=input_ids, - attention_mask=attention_mask, - max_new_tokens=max_new_tokens, - do_sample=False, - pad_token_id=pad_id, - eos_token_id=eos_id, - logits_to_keep=0, - ) - finally: - engine.close() - - # ----------------------------------------------------------------------- - # Read DB (monitoring.clickhouse_reader + monitoring.segment_merger) - # ----------------------------------------------------------------------- - - ch = CHClickhouseDriverReadOnly( - host=str(db_cfg_native.host), - port=int(db_cfg_native.port), - username=str(db_cfg_native.username), - password=str(db_cfg_native.password), - database=str(db_cfg_native.database), - table=str(db_cfg_native.table), - secure=bool(getattr(db_cfg_native, "secure", False)), - client_settings=getattr(db_cfg_native, "client_settings", None), - decode_strings=True, - ) - try: - rows = ch.prefix_get((unique_run_model_id, ), return_full_key_tuple=True) - finally: - ch.close() - - if not rows: - pytest.fail(f"No rows found in ClickHouse for model_id={unique_run_model_id}") - - # If multiple shard_ranks are present, pick rank 0 if available, else the minimum. - shard_ranks = sorted({int(key[4]) for key, _t in rows}) - chosen_shard_rank = 0 if 0 in shard_ranks else (shard_ranks[0] if shard_ranks else 0) - rows = [(k, t) for (k, t) in rows if int(k[4]) == chosen_shard_rank] - - grouped: Dict[str, Dict[Tuple[int, str], List[Tuple[int, int, torch.Tensor]]]] = {} - for full_key, t_raw in rows: - # full_key = (model_id, request_id, act_name, layer_no, shard_rank, start_token_idx, end_token_idx) - _model_id, req_id, act_name_raw, layer_no_raw, _shard_rank, s, e = full_key - - layer_no, act_name = _canon_layer_and_act(str(act_name_raw), int(layer_no_raw)) - - t = t_raw.detach().cpu() - grouped.setdefault(str(req_id), {}).setdefault((layer_no, act_name), []).append( - (int(s), int(e), t) - ) - - request_ids = sorted(grouped.keys(), key=_parse_request_id) - - def _sort_chunks(chunks: List[Tuple[int, int, torch.Tensor]]) -> List[Tuple[int, int, torch.Tensor]]: - return sorted(chunks, key=lambda x: (x[0], x[1])) - - def _validate_contiguous( - chunks_sorted: List[Tuple[int, int, torch.Tensor]], expected_end: int, ctx: str - ) -> None: - if not chunks_sorted: - raise AssertionError(f"{ctx}: no chunks") - if chunks_sorted[0][0] != 0: - raise AssertionError(f"{ctx}: first chunk start={chunks_sorted[0][0]} expected 0") - prev_end = chunks_sorted[0][1] - for s2, e2, _t in chunks_sorted[1:]: - if s2 != prev_end: - raise AssertionError(f"{ctx}: non-contiguous chunks: start={s2} prev_end={prev_end}") - prev_end = e2 - if prev_end != expected_end: - raise AssertionError(f"{ctx}: coverage end={prev_end} expected_end={expected_end}") - - # Safety: this run should have exactly one group_id (single batch reset) - seen_group_ids: set[int] = set() - - db_token_ids_by_req: Dict[str, torch.Tensor] = {} - local_index_by_req: Dict[str, int] = {} - prompt_len_by_req: Dict[str, int] = {} - - for req_id in request_ids: - group_id, local_i = _parse_request_id(req_id) - seen_group_ids.add(group_id) - if not (0 <= local_i < batch_size): - raise AssertionError(f"{req_id}: local_index={local_i} out of range batch_size={batch_size}") - local_index_by_req[req_id] = local_i - - hooks_map = grouped[req_id] - if (-1, "token_ids") not in hooks_map: - raise AssertionError(f"{req_id}: DB missing token_ids") - - tok_chunks = _sort_chunks(hooks_map[(-1, "token_ids")]) - db_tok_end = int(tok_chunks[-1][1]) - _validate_contiguous(tok_chunks, expected_end=db_tok_end, ctx=f"{req_id} token_ids") - - db_tok = merge_segments([t for _, _, t in tok_chunks], "token_ids").to(torch.long) - if db_tok.ndim != 1: - db_tok = db_tok.view(-1) - - prompt0 = hf_initial_prompt_tokens[local_i] - plen0 = int(prompt0.numel()) - if db_tok.numel() < plen0 or not torch.equal(db_tok[:plen0], prompt0): - raise AssertionError( - f"{req_id}: request_id->row mapping safety check failed (initial prompt prefix). " - f"initial_prompt_len={plen0} db_tok_len={db_tok.numel()}" - ) - - db_token_ids_by_req[req_id] = db_tok.cpu() - prompt_len_by_req[req_id] = plen0 - - if len(seen_group_ids) != 1: - raise AssertionError(f"expected exactly one group_id for this test run, got {sorted(seen_group_ids)}") - - # ----------------------------------------------------------------------- - # Build HF references (no assertions here) - # ----------------------------------------------------------------------- - - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_id, - attn_implementation="eager", - torch_dtype=torch.float16, - ).to(device).eval() - - # Optional modules for GPT2-like reconstructions - wte = getattr(getattr(hf_model, "transformer", None), "wte", None) - wpe = getattr(getattr(hf_model, "transformer", None), "wpe", None) - # Qwen3-like: embed_tokens on model sub-module (no separate pos embed) - embed_tokens = getattr(getattr(hf_model, "model", None), "embed_tokens", None) - - # Compute num_layers once (used in per-layer comparison loop below) - num_layers = get_num_layers_from_config(hf_model) - - req_order = sorted(request_ids, key=_parse_request_id) - - # Run reference on ORIGINAL (non-hooked) model in a subprocess. - # This validates that our hooks don't change the model output, - # and ensures clean GPU memory isolation. - import subprocess, tempfile - ref_dir = tempfile.mkdtemp(prefix="hf_ref_") - ref_env = {**os.environ, "E2E_BATCH_SIZE": str(batch_size), - "E2E_MAX_NEW_TOKENS": str(max_new_tokens), - "E2E_MODEL": os.environ.get("E2E_MODEL", "gpt2")} - ref_result = subprocess.run( - [sys.executable, "-m", "tests.hf_reference_runner", "--output-dir", ref_dir], - env=ref_env, capture_output=True, text=True, cwd=os.path.dirname(os.path.dirname(__file__)), - ) - if ref_result.returncode != 0: - pytest.fail(f"Reference runner failed:\n{ref_result.stderr[-2000:]}") - hf_refs_batch = _load_hf_refs_from_disk(ref_dir) - import shutil - shutil.rmtree(ref_dir, ignore_errors=True) - hf_gens_batch = _hf_generate_collect_scores_batched( - hf_model=hf_model, - input_ids_batch=input_ids, - attention_mask_batch=attention_mask, - max_new_tokens=max_new_tokens, - eos_token_id=eos_id, - pad_token_id=pad_id, - device=device, - ) - if len(hf_refs_batch) != batch_size or len(hf_gens_batch) != batch_size: - raise AssertionError( - f"HF batched refs unexpected batch: rollout={len(hf_refs_batch)} " - f"gen={len(hf_gens_batch)} batch_size={batch_size}" - ) - - hf_ref_by_req: Dict[str, _HFRef] = {} - hf_gen_by_req: Dict[str, _HFGenRef] = {} - - def _decode(ids: torch.Tensor) -> str: - return tokenizer.decode(ids.tolist(), skip_special_tokens=False) - - for req_id in req_order: - i = local_index_by_req[req_id] - plen = int(prompt_len_by_req[req_id]) - - ref = hf_refs_batch[i] - gen_ref = hf_gens_batch[i] - - hf_ref_by_req[req_id] = ref - hf_gen_by_req[req_id] = gen_ref - - if print_text: - db_seq = db_token_ids_by_req[req_id] - db_prompt = db_seq[:plen] - db_gen = db_seq[plen:] - rol_prompt = ref.token_ids[:plen] - rol_gen = ref.token_ids[plen:] - gen_prompt = gen_ref.token_ids[:plen] - gen_gen = gen_ref.token_ids[plen:] - - print(f"\n=== {req_id} (local_index={i}, shard_rank={chosen_shard_rank}) ===") - print(f"DB: prompt_tokens={plen} generated_tokens={int(db_gen.numel())} total_tokens={int(db_seq.numel())}") - print(f"DB PROMPT: {_decode(db_prompt)!r}") - print(f"DB GENERATED: {_decode(db_gen)!r}") - print(f"DB FULL: {_decode(db_seq)!r}") - print(f"ROL: prompt_tokens={plen} generated_tokens={int(rol_gen.numel())} total_tokens={int(ref.token_ids.numel())}") - print(f"ROL PROMPT: {_decode(rol_prompt)!r}") - print(f"ROL GENERATED:{_decode(rol_gen)!r}") - print(f"ROL FULL: {_decode(ref.token_ids)!r}") - print(f"GEN: prompt_tokens={plen} generated_tokens={int(gen_gen.numel())} total_tokens={int(gen_ref.token_ids.numel())}") - print(f"GEN PROMPT: {_decode(gen_prompt)!r}") - print(f"GEN GENERATED:{_decode(gen_gen)!r}") - print(f"GEN FULL: {_decode(gen_ref.token_ids)!r}") - print("TOKENS MATCH?: YES (DB==ROL==GEN)") - - # Helpers for top-k logit printing - def _tok_piece(tok_id: int) -> str: - try: - return tokenizer.decode([tok_id], skip_special_tokens=False) - except Exception: - return f"" - - def _fmt_topk(ids_row: torch.Tensor, vals_row: torch.Tensor) -> str: - parts: List[str] = [] - for tid, v in zip(ids_row.tolist(), vals_row.tolist()): - parts.append(f"{int(tid)}:{_tok_piece(int(tid))!r}:{float(v):.6g}") - return " | ".join(parts) - - # ----------------------------------------------------------------------- - # Subtests: one per assertion - # ----------------------------------------------------------------------- - - for req_id in req_order: - i = local_index_by_req[req_id] - hooks_map = grouped[req_id] - - seq = db_token_ids_by_req[req_id].to(torch.long) # [T] - seq_len = int(seq.numel()) - ref = hf_ref_by_req[req_id] - gen_ref = hf_gen_by_req[req_id] - prompt_len = int(prompt_len_by_req[req_id]) - gen_base_pos = prompt_len - 1 # scores[0] corresponds to logits at pos (prompt_len-1) - - # --- token_ids --- - with subtests.test(msg=f"{req_id}/token_ids_rol"): - assert bitwise_equal(ref.token_ids, seq), ( - f"HF rollout tokens != DB token_ids " - f"(hf_len={int(ref.token_ids.numel())} db_len={seq_len})" - ) - - with subtests.test(msg=f"{req_id}/token_ids_gen"): - assert bitwise_equal(gen_ref.token_ids, seq), ( - f"HF generate tokens != DB token_ids " - f"(hf_len={int(gen_ref.token_ids.numel())} db_len={seq_len})" - ) - - # --- final_logits --- - logits_chunks_raw = hooks_map.get((-1, "final_logits"), []) - if logits_chunks_raw: - lchunks = sorted(logits_chunks_raw, key=lambda x: (x[0], x[1])) - db_logits_full = merge_segments([t for _, _, t in lchunks], "final_logits") - if db_logits_full.ndim == 1: - db_logits_full = db_logits_full.unsqueeze(0) - # ref.final_logits is decode-only scores from generate(): - # ref[s] = logits at position (prompt_len - 1 + s). - # Align with DB logits by position. - n_ref = int(ref.final_logits.shape[0]) - start = prompt_len - 1 - end = min(start + n_ref, int(db_logits_full.shape[0])) - n = end - start - db_slice = db_logits_full[start:end, :] - rol_slice = ref.final_logits[:n, :] - vocab_db = int(db_slice.shape[1]) - - if print_topk_logits: - print(f"\n=== TOP{topk_k} LOGITS {req_id} (local_index={i}, shard_rank={chosen_shard_rank}) ===") - print(f"seq_len={seq_len} vocab={vocab_db}") - db_topv, db_topi = torch.topk(db_slice.float(), k=topk_k, dim=-1) - rol_topv, rol_topi = torch.topk(rol_slice.float(), k=topk_k, dim=-1) - for tpos in range(seq_len): - cur_id = int(seq[tpos].item()) - cur_piece = _tok_piece(cur_id) - if tpos + 1 < seq_len: - nxt_id = int(seq[tpos + 1].item()) - nxt_piece = _tok_piece(nxt_id) - label_str = f" next={nxt_id}:{nxt_piece!r}" - else: - label_str = " next=" - print(f"\npos={tpos} tok={cur_id}:{cur_piece!r}{label_str}") - print(f" DB: {_fmt_topk(db_topi[tpos], db_topv[tpos])}") - print(f" ROL: {_fmt_topk(rol_topi[tpos], rol_topv[tpos])}") - if gen_base_pos >= 0 and gen_base_pos <= tpos <= (gen_base_pos + len(gen_ref.scores) - 1): - sidx = tpos - gen_base_pos - gs = gen_ref.scores[sidx] - g_topv, g_topi = torch.topk(gs.float(), k=topk_k, dim=-1) - print(f" GEN: {_fmt_topk(g_topi, g_topv)}") - else: - print(" GEN: ") - - with subtests.test(msg=f"{req_id}/final_logits"): - if not bitwise_equal(db_slice, rol_slice): - diff = (db_slice.float() - rol_slice.float()).abs() - max_abs = float(diff.max().item()) - flat_idx = int(diff.view(-1).argmax().item()) - r = flat_idx // vocab_db - c = flat_idx % vocab_db - pytest.fail(f"final_logits mismatch (max_abs={max_abs}) at row={r} vocab_idx={c}") - - # --- hook_embed / hook_pos_embed (GPT2-like: both wte and wpe) --- - if ( - wte is not None - and wpe is not None - and (-1, "hook_embed") in hooks_map - and (-1, "hook_pos_embed") in hooks_map - ): - ids = seq.to(device) - pos = _positions_for_unpadded(seq_len, device=device) - emb = wte(ids).detach().cpu() # [T, d] - pos_emb = wpe(pos).detach().cpu() # [T, d] - - chunks = sorted(hooks_map[(-1, "hook_embed")], key=lambda x: (x[0], x[1])) - _validate_contiguous(chunks, expected_end=seq_len, ctx=f"{req_id} hook_embed") - db_t = merge_segments([t for _, _, t in chunks], "hook_embed") - with subtests.test(msg=f"{req_id}/hook_embed"): - assert tuple(db_t.shape) == tuple(emb.shape), ( - f"hook_embed shape mismatch db={tuple(db_t.shape)} hf={tuple(emb.shape)}" - ) - assert bitwise_equal(db_t, emb), ( - f"hook_embed mismatch (max_abs={float((db_t.float() - emb.float()).abs().max().item())})" - ) - - chunks = sorted(hooks_map[(-1, "hook_pos_embed")], key=lambda x: (x[0], x[1])) - _validate_contiguous(chunks, expected_end=seq_len, ctx=f"{req_id} hook_pos_embed") - db_t = merge_segments([t for _, _, t in chunks], "hook_pos_embed") - with subtests.test(msg=f"{req_id}/hook_pos_embed"): - assert tuple(db_t.shape) == tuple(pos_emb.shape), ( - f"hook_pos_embed shape mismatch db={tuple(db_t.shape)} hf={tuple(pos_emb.shape)}" - ) - assert bitwise_equal(db_t, pos_emb), ( - f"hook_pos_embed mismatch (max_abs={float((db_t.float() - pos_emb.float()).abs().max().item())})" - ) - - # --- hook_embed only (Qwen3-like: RoPE, no separate pos embed) --- - elif ( - embed_tokens is not None - and wpe is None - and (-1, "hook_embed") in hooks_map - ): - emb = embed_tokens(seq.to(device)).detach().cpu() # [T, d] - chunks = sorted(hooks_map[(-1, "hook_embed")], key=lambda x: (x[0], x[1])) - _validate_contiguous(chunks, expected_end=seq_len, ctx=f"{req_id} hook_embed") - db_t = merge_segments([t for _, _, t in chunks], "hook_embed") - with subtests.test(msg=f"{req_id}/hook_embed"): - assert tuple(db_t.shape) == tuple(emb.shape), ( - f"hook_embed shape mismatch db={tuple(db_t.shape)} hf={tuple(emb.shape)}" - ) - assert bitwise_equal(db_t, emb), ( - f"hook_embed mismatch (max_abs={float((db_t.float() - emb.float()).abs().max().item())})" - ) - - # --- hook_final_ln --- - # TODO: hook_final_ln comparison skipped. - - # --- per-layer: attention pattern + resid_pre --- - # Support both GPT-2 naming (blocks.attn.hook_pattern, blocks.hook_resid_*) - # and Qwen3 naming (layers.self_attn.hook_pattern, layers.hook_resid_*) - _ATTN_PATTERN_KEYS = ("blocks.attn.hook_pattern", "layers.self_attn.hook_pattern") - _RESID_PRE_KEYS = ("blocks.hook_resid_pre", "layers.hook_resid_pre") - - n_layers = len(ref.attn_pattern) if ref.attn_pattern else 0 - assert n_layers == num_layers, ( - f"{req_id}: attn_pattern layer count mismatch: rollout={n_layers} config={num_layers}" - ) - - for layer_no in range(n_layers): - # attn_pattern: compare per-chunk (can't merge because kv_dim - # differs between prefill [H, plen, plen] and decode [H, 1, plen+i]) - key = next(((layer_no, k) for k in _ATTN_PATTERN_KEYS if (layer_no, k) in hooks_map), None) - if key is not None: - pat = ref.attn_pattern[layer_no] # [H, T, T] - chunks = sorted(hooks_map[key], key=lambda x: (x[0], x[1])) - all_ok = True - fail_msg = "" - for start, end, t_chunk in chunks: - q_len = end - start - db_c = t_chunk - if db_c.ndim == 4 and db_c.shape[0] == 1: - db_c = db_c.squeeze(0) - # kv_dim for these rows: causal, valid up to position 'end' - kv_valid = end - db_c = db_c[:, :q_len, :kv_valid] - ref_c = pat[:, start:end, :kv_valid] - if db_c.shape != ref_c.shape: - all_ok = False - fail_msg = (f"shape mismatch at [{start}:{end}] " - f"db={db_c.shape} ref={ref_c.shape}") - break - if not bitwise_equal(db_c, ref_c): - max_abs = float((db_c.float() - ref_c.float()).abs().max().item()) - all_ok = False - fail_msg = (f"value mismatch at [{start}:{end}] " - f"max_abs={max_abs:.6f}") - break - with subtests.test(msg=f"{req_id}/layer{layer_no}/attn_pattern"): - assert all_ok, ( - f"pattern layer={layer_no}: {fail_msg}" - ) - - key = next(((layer_no, k) for k in _RESID_PRE_KEYS if (layer_no, k) in hooks_map), None) - if key is not None and ref.hidden_states and layer_no < len(ref.hidden_states): - hs = ref.hidden_states[layer_no] # [T, d] - chunks = sorted(hooks_map[key], key=lambda x: (x[0], x[1])) - _validate_contiguous(chunks, expected_end=seq_len, ctx=f"{req_id} layer{layer_no} resid_pre") - db_t = merge_segments([t for _, _, t in chunks], key[1]) - with subtests.test(msg=f"{req_id}/layer{layer_no}/resid_pre"): - assert tuple(db_t.shape) == tuple(hs.shape), ( - f"resid_pre shape mismatch layer={layer_no} db={tuple(db_t.shape)} hf={tuple(hs.shape)}" - ) - assert bitwise_equal(db_t, hs), ( - f"resid_pre mismatch layer={layer_no} " - f"(max_abs={float((db_t.float() - hs.float()).abs().max().item())})" - ) - - # --- resid_final (global: last layer's pre-norm residual) --- - # HF's output_hidden_states[-1] is POST-final-norm (after ln_f), - # not pre-norm. resid_final captures pre-norm. No direct HF - # reference available, so we only check shape and presence. - key = (-1, "hook_resid_final") - if key in hooks_map: - chunks = sorted(hooks_map[key], key=lambda x: (x[0], x[1])) - _validate_contiguous(chunks, expected_end=seq_len, ctx=f"{req_id} resid_final") - db_t = merge_segments([t for _, _, t in chunks], key[1]) - with subtests.test(msg=f"{req_id}/resid_final"): - assert db_t.shape[-1] == ref.hidden_states[0].shape[-1] if ref.hidden_states else True, ( - f"resid_final hidden_dim mismatch" - ) - assert db_t.shape[0] == seq_len, ( - f"resid_final token count mismatch db={db_t.shape[0]} expected={seq_len}" - ) - - -# --------------------------------------------------------------------------- -# CUDA-graph correctness test -# --------------------------------------------------------------------------- - -@pytest.mark.skipif(True, reason=( - "DISABLED: compiled rollout reference (StaticCache + torch.compile) cannot " - "replicate generate()'s internal StaticCache attention mask / position handling. " - "The manual rollout produces different hidden states than generate() even for " - "identical inputs — this is a fundamental mismatch in how HF handles StaticCache " - "internally vs externally. Use test_e2e_cuda_graphs_vs_eager_hf instead, which " - "compares CUDA-graph DB against an uncompiled eager reference with relaxed tolerance." -)) -def test_e2e_correctness_hf_cuda_graphs(subtests) -> None: - """Same as test_e2e_correctness_hf but with torch.compile + static KV cache (CUDA graphs). - - NOTE: This test is currently DISABLED. The compiled rollout reference uses - StaticCache + torch.compile on a manual decode loop, but this produces - different numerical results from HF generate(cache_implementation="static") - because generate() handles attention masks and position_ids differently - internally. CUDA graphs also prevent reading hidden states from generate() - (Bug 11 in debug.log). See test_e2e_cuda_graphs_vs_eager_hf for the - working alternative. - - Run with: - CUDA_MODULE_LOADING=EAGER pytest -q -s tests/test_e2e_correctness_vs_hf.py::test_e2e_correctness_hf_cuda_graphs - """ - try: - import clickhouse_driver # noqa: F401 - except Exception: - pytest.skip("clickhouse-driver is required") - - try: - from monitoring import ( # type: ignore - MonitoringConfig, - MonitoringEngine, - ) - from monitoring._native_engine import ClickHouseClientConfig # type: ignore - from monitoring.config import CaptureSchedule # type: ignore - from integration.hf_adapter import generate_with_monitoring # type: ignore - except Exception as exc: - pytest.skip(f"monitoring native extension not available: {exc}") - - try: - from transformers import AutoModelForCausalLM, AutoTokenizer - from transformers.models.gpt2_p.modeling_gpt2 import HookedGPT2LMHeadModel # type: ignore - from transformers.models.qwen3_p.modeling_qwen3 import HookedQwen3ForCausalLM # type: ignore - from transformers.models.llama_p.modeling_llama import HookedLlamaForCausalLM # type: ignore - except Exception as exc: - pytest.skip(f"transformers or Hooked* classes not available: {exc}") - - # ----------------------------------------------------------------------- - # Config — fixed small values to keep the test fast - # ----------------------------------------------------------------------- - batch_size = int(os.environ.get("E2E_BATCH_SIZE", "4")) - max_new_tokens = int(os.environ.get("E2E_MAX_NEW_TOKENS", "8")) - hf_model_id = os.environ.get("E2E_MODEL", "gpt2") - hf_model_id = _MODEL_ALIASES.get(hf_model_id.lower(), hf_model_id) - chunk_bytes = int(os.environ.get("E2E_CHUNK_BYTES", str(256 * 1024))) - # E2E_HF_DROP_LAST_TOKEN is no longer needed: both monitored and HF reference - # use generate(), so they produce the same number of tokens/hidden states. - - device = torch.device("cuda") - - # ----------------------------------------------------------------------- - # Tokenizer + prompts - # ----------------------------------------------------------------------- - tokenizer = AutoTokenizer.from_pretrained(hf_model_id) - if tokenizer.pad_token_id is None: - tokenizer.pad_token_id = tokenizer.eos_token_id - tokenizer.padding_side = "left" - eos_id = int(tokenizer.eos_token_id) - pad_id = int(tokenizer.pad_token_id) - - prompts = [("Hello " * (i + 1)).strip() for i in range(batch_size)] - encoded = tokenizer(prompts, return_tensors="pt", padding=True) - input_ids = encoded["input_ids"].to(device) - attention_mask = encoded["attention_mask"].to(device) - - hf_initial_prompt_tokens: List[torch.Tensor] = [] - for j in range(batch_size): - hf_initial_prompt_tokens.append( - _strip_left_pad( - input_ids[j].detach().cpu(), - attention_mask[j].detach().cpu(), - ).to(torch.long) - ) - - # ----------------------------------------------------------------------- - # Monitoring + DB config - # ----------------------------------------------------------------------- - mon_cfg = MonitoringConfig( - schedule=CaptureSchedule(capture_prefill=True, capture_decode=True), - ) - - db_cfg_native = ClickHouseClientConfig() - db_cfg_native.host = os.environ.get("DMX_DB_HOST", "localhost") - db_cfg_native.port = int(os.environ.get("DMX_DB_PORT", "9000")) - db_cfg_native.username = os.environ.get("DMX_DB_USER", "default") - db_cfg_native.password = os.environ.get("DMX_DB_PASSWORD", "") - db_cfg_native.database = os.environ.get("DMX_DB_DATABASE", "default") - db_cfg_native.table = os.environ.get("DMX_DB_TABLE", "offload") - db_cfg_native.secure = False - db_cfg_native.client_side_compress = "none" - db_cfg_native.client_settings = None - db_cfg_native.create_database_if_missing = True - db_cfg_native.drop_existing_database = True - db_cfg_native.index_granularity = 8192 - - host_cfg = _make_host_cfg(db_cfg_native) - ring_cfg = _make_ring_cfg() - - # ----------------------------------------------------------------------- - # Monitored model — compiled with torch.compile + static cache (CUDA graphs) - # ----------------------------------------------------------------------- - unique_run_model_id = f"e2e_cuda_graphs::{uuid.uuid4().hex}"[:120] - engine = MonitoringEngine( - config=mon_cfg, - model_id=unique_run_model_id, db_config=host_cfg, - ) - engine.enable_ring_transport(ring_cfg) - - if "qwen3" in hf_model_id.lower(): - model_cls = HookedQwen3ForCausalLM - elif "llama" in hf_model_id.lower(): - model_cls = HookedLlamaForCausalLM - else: - model_cls = HookedGPT2LMHeadModel - mon_model = model_cls.from_pretrained( - hf_model_id, attn_implementation="eager", torch_dtype=torch.float16, - ) - mon_model.to(device).eval() - - mon_model.monitoring_engine = engine - - try: - from transformers import CompileConfig - with torch.no_grad(): - gen_out = generate_with_monitoring( - mon_model, - input_ids=input_ids, - attention_mask=attention_mask, - max_new_tokens=max_new_tokens, - do_sample=False, - pad_token_id=pad_id, - eos_token_id=eos_id, - cache_implementation="static", - compile_config=CompileConfig(mode="reduce-overhead", fullgraph=False), - ) - finally: - engine.close() - - # Build per-request reference sequences from the generate() output. - # We use the compiled model's own output as the reference — this avoids - # any comparison against a different model that may compute different - # values under static cache or torch.compile. - gen_out_cpu = gen_out.detach().cpu().long() # [batch, total_len] - ref_seqs: List[torch.Tensor] = [] - for j in range(batch_size): - seq = _strip_left_pad(gen_out_cpu[j], (gen_out_cpu[j] != pad_id).long()) - ref_seqs.append(seq) - - # ----------------------------------------------------------------------- - # Read DB - # ----------------------------------------------------------------------- - from monitoring.clickhouse_reader import CHClickhouseDriverReadOnly - from monitoring.segment_merger import merge_segments, parse_internal_id - - ch = CHClickhouseDriverReadOnly( - host=str(db_cfg_native.host), - port=int(db_cfg_native.port), - username=str(db_cfg_native.username), - password=str(db_cfg_native.password), - database=str(db_cfg_native.database), - table=str(db_cfg_native.table), - secure=bool(getattr(db_cfg_native, "secure", False)), - client_settings=getattr(db_cfg_native, "client_settings", None), - decode_strings=True, - ) - try: - rows = ch.prefix_get((unique_run_model_id,), return_full_key_tuple=True) - finally: - ch.close() - - print(f"\n[DEBUG] Total DB rows: {len(rows)}") - if not rows: - pytest.fail(f"No rows found in ClickHouse for model_id={unique_run_model_id!r}. " - "This means monitoring produced no output at all under CUDA graphs.") - - shard_ranks = sorted({int(key[4]) for key, _t in rows}) - chosen_shard = 0 if 0 in shard_ranks else shard_ranks[0] - rows = [(k, t) for (k, t) in rows if int(k[4]) == chosen_shard] - - grouped: Dict[str, Dict[Tuple[int, str], List[Tuple[int, int, torch.Tensor]]]] = {} - for full_key, t_raw in rows: - _model_id, req_id, act_name_raw, layer_no_raw, _shard, s, e = full_key - layer_no, act_name = _canon_layer_and_act(str(act_name_raw), int(layer_no_raw)) - grouped.setdefault(str(req_id), {}).setdefault( - (layer_no, act_name), [] - ).append((int(s), int(e), t_raw.detach().cpu())) - - request_ids = sorted(grouped.keys(), key=_parse_request_id) - - # DEBUG: show per-request hook chunk counts - from collections import Counter - hook_totals = Counter() - for rid in request_ids: - hooks_map = grouped[rid] - for (layer, hname), chunks in hooks_map.items(): - hook_totals[hname] += len(chunks) - print(f"[DEBUG] Hook totals across all requests:") - for hname in ['token_ids', 'hook_embed', 'hook_pos_embed', 'blocks.0.hook_resid_pre', 'hook_resid_final', 'hook_final_ln', 'final_logits']: - print(f" {hname}: {hook_totals.get(hname, 0)} chunks") - rid = request_ids[0] - hooks_map = grouped[rid] - print(f"[DEBUG] All hooks for {rid}:") - for (layer, hname) in sorted(hooks_map.keys()): - chunks = hooks_map[(layer, hname)] - print(f" ({layer},{hname}): {len(chunks)} chunks") - - # ----------------------------------------------------------------------- - # HF reference — compiled manual rollout with StaticCache. - # Uses torch.compile(mode="reduce-overhead", fullgraph=False) on the - # decode step + cudagraph_mark_step_begin() + immediate .detach().cpu() - # to clone hidden states before CUDA graph buffers are overwritten. - # (generate() can't do this — see Bug 11 in debug.log) - # ----------------------------------------------------------------------- - hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_id, attn_implementation="eager", torch_dtype=torch.float16, - ).to(device).eval() - - wte = getattr(getattr(hf_model, "transformer", None), "wte", None) - wpe = getattr(getattr(hf_model, "transformer", None), "wpe", None) - embed_tokens = getattr(getattr(hf_model, "model", None), "embed_tokens", None) - num_layers = get_num_layers_from_config(hf_model) - - hf_refs_batch = _hf_greedy_rollout_collect_all_batched( - hf_model=hf_model, - input_ids_batch=input_ids, - attention_mask_batch=attention_mask, - max_new_tokens=max_new_tokens, - eos_token_id=eos_id, - pad_token_id=pad_id, - device=device, - want_hidden_states=True, - want_attentions=True, - compiled=True, - ) - - # ----------------------------------------------------------------------- - # Pre-loop: build per-request dicts and safety-check token_ids - # ----------------------------------------------------------------------- - local_index_by_req: Dict[str, int] = {} - prompt_len_by_req: Dict[str, int] = {} - db_token_ids_by_req: Dict[str, torch.Tensor] = {} - - for req_id in request_ids: - _gid, local_i = _parse_request_id(req_id) - local_index_by_req[req_id] = local_i - - req_order = sorted(request_ids, key=_parse_request_id) - - for req_id in req_order: - local_i = local_index_by_req[req_id] - hooks_map = grouped[req_id] - prompt0 = hf_initial_prompt_tokens[local_i] - plen = int(prompt0.numel()) - prompt_len_by_req[req_id] = plen - - if (-1, "token_ids") not in hooks_map: - raise AssertionError(f"{req_id}: DB missing token_ids under CUDA graphs") - - tok_chunks = sorted(hooks_map[(-1, "token_ids")], key=lambda x: (x[0], x[1])) - db_tok = merge_segments([t for _, _, t in tok_chunks], "token_ids").to(torch.long) - if db_tok.ndim != 1: - db_tok = db_tok.view(-1) - - if db_tok.numel() < plen or not torch.equal(db_tok[:plen], prompt0): - raise AssertionError( - f"{req_id}: DB token_ids prompt prefix mismatch under CUDA graphs " - f"(plen={plen} db_len={db_tok.numel()})" - ) - db_token_ids_by_req[req_id] = db_tok.cpu() - - def _sort_chunks(chunks): - return sorted(chunks, key=lambda x: (x[0], x[1])) - - def _validate_contiguous(chunks_sorted, expected_end, ctx): - if not chunks_sorted: - raise AssertionError(f"{ctx}: no chunks") - if chunks_sorted[0][0] != 0: - raise AssertionError(f"{ctx}: first chunk start={chunks_sorted[0][0]} expected 0") - prev_end = chunks_sorted[0][1] - for s2, e2, _t in chunks_sorted[1:]: - if s2 != prev_end: - raise AssertionError(f"{ctx}: non-contiguous chunks start={s2} prev_end={prev_end}") - prev_end = e2 - if prev_end != expected_end: - raise AssertionError(f"{ctx}: coverage end={prev_end} expected_end={expected_end}") - - # ----------------------------------------------------------------------- - # Per-request assertions (full verification) - # ----------------------------------------------------------------------- - _RESID_PRE_KEYS = ("blocks.hook_resid_pre", "layers.hook_resid_pre") - - # HF reference is uncompiled; monitored model is compiled (reduce-overhead). - # Fall back to allclose if not bitwise equal. - _COMPILED_ATOL = 0.5 # safety net: real transport errors are >> 1 - - def _assert_close_or_bitwise(db_t, ref_t, label): - if bitwise_equal(db_t, ref_t): - return - diff = (db_t.float() - ref_t.float()).abs() - max_abs = float(diff.max().item()) - if torch.allclose(db_t.float(), ref_t.float(), atol=_COMPILED_ATOL, rtol=0.0): - import warnings - warnings.warn( - f"[NOT BITWISE] {label}: max_abs_diff={max_abs:.6f} " - f"(within atol={_COMPILED_ATOL}, but not bitwise equal)" - ) - return - pytest.fail( - f"{label}: max_abs_diff={max_abs:.6f} > atol={_COMPILED_ATOL}" - ) - - for req_id in req_order: - local_i = local_index_by_req[req_id] - hooks_map = grouped[req_id] - plen = prompt_len_by_req[req_id] - prompt0 = hf_initial_prompt_tokens[local_i] - - db_tok = db_token_ids_by_req[req_id] - seq_len = int(db_tok.numel()) - - ref = hf_refs_batch[local_i] - - # --- token_ids --- - with subtests.test(msg=f"{req_id}/cuda_graph/token_ids_present"): - assert (-1, "token_ids") in hooks_map, f"{req_id}: DB missing token_ids" - - with subtests.test(msg=f"{req_id}/cuda_graph/prompt_prefix"): - assert db_tok.numel() >= plen and torch.equal(db_tok[:plen], prompt0), ( - f"{req_id}: prompt prefix mismatch" - ) - - with subtests.test(msg=f"{req_id}/cuda_graph/token_ids_match_hf"): - assert bitwise_equal(db_tok, ref.token_ids), ( - f"{req_id}: DB token_ids do not match HF generate() under CUDA graphs. " - f"db_len={db_tok.numel()} hf_len={ref.token_ids.numel()} " - f"(if db_len << hf_len the CUDA-graph monitoring bug is present)" - ) - - # --- final_logits --- - logits_chunks_raw = hooks_map.get((-1, "final_logits"), []) - if logits_chunks_raw: - lchunks = _sort_chunks(logits_chunks_raw) - db_logits = merge_segments([t for _, _, t in lchunks], "final_logits") - if db_logits.ndim == 1: - db_logits = db_logits.unsqueeze(0) - n_ref = int(ref.final_logits.shape[0]) - start = plen - 1 - end = min(start + n_ref, int(db_logits.shape[0])) - n = end - start - db_slice = db_logits[start:end, :] - rol_slice = ref.final_logits[:n, :] - with subtests.test(msg=f"{req_id}/cuda_graph/final_logits"): - assert n > 0, f"final_logits: no overlapping rows" - assert db_slice.shape[1] == rol_slice.shape[1], ( - f"final_logits vocab mismatch db={db_slice.shape[1]} hf={rol_slice.shape[1]}" - ) - _assert_close_or_bitwise(db_slice, rol_slice, f"{req_id} final_logits") - - # --- hook_embed --- - seq = db_tok.to(device) - if wte is not None and wpe is not None and (-1, "hook_embed") in hooks_map: - emb = wte(seq).detach().cpu() - chunks = _sort_chunks(hooks_map[(-1, "hook_embed")]) - _validate_contiguous(chunks, seq_len, f"{req_id} hook_embed") - db_t = merge_segments([t for _, _, t in chunks], "hook_embed") - with subtests.test(msg=f"{req_id}/cuda_graph/hook_embed"): - assert tuple(db_t.shape) == tuple(emb.shape), ( - f"hook_embed shape mismatch db={tuple(db_t.shape)} hf={tuple(emb.shape)}" - ) - _assert_close_or_bitwise(db_t, emb, f"{req_id} hook_embed") - elif embed_tokens is not None and wpe is None and (-1, "hook_embed") in hooks_map: - emb = embed_tokens(seq).detach().cpu() - chunks = _sort_chunks(hooks_map[(-1, "hook_embed")]) - _validate_contiguous(chunks, seq_len, f"{req_id} hook_embed") - db_t = merge_segments([t for _, _, t in chunks], "hook_embed") - with subtests.test(msg=f"{req_id}/cuda_graph/hook_embed"): - assert tuple(db_t.shape) == tuple(emb.shape), ( - f"hook_embed shape mismatch db={tuple(db_t.shape)} hf={tuple(emb.shape)}" - ) - _assert_close_or_bitwise(db_t, emb, f"{req_id} hook_embed") - - # --- hook_pos_embed (GPT2 only) --- - if wpe is not None and (-1, "hook_pos_embed") in hooks_map: - pos = _positions_for_unpadded(seq_len, device=device) - pos_emb = wpe(pos).detach().cpu() - chunks = _sort_chunks(hooks_map[(-1, "hook_pos_embed")]) - _validate_contiguous(chunks, seq_len, f"{req_id} hook_pos_embed") - db_t = merge_segments([t for _, _, t in chunks], "hook_pos_embed") - with subtests.test(msg=f"{req_id}/cuda_graph/hook_pos_embed"): - assert tuple(db_t.shape) == tuple(pos_emb.shape), ( - f"hook_pos_embed shape mismatch db={tuple(db_t.shape)} hf={tuple(pos_emb.shape)}" - ) - _assert_close_or_bitwise(db_t, pos_emb, f"{req_id} hook_pos_embed") - - # --- per-layer: resid_pre + attn_pattern --- - _ATTN_PATTERN_KEYS = ("blocks.attn.hook_pattern", "layers.self_attn.hook_pattern") - for layer_no in range(num_layers): - key = next(((layer_no, k) for k in _RESID_PRE_KEYS if (layer_no, k) in hooks_map), None) - if key is not None and ref.hidden_states and layer_no < len(ref.hidden_states): - hs = ref.hidden_states[layer_no] - chunks = _sort_chunks(hooks_map[key]) - _validate_contiguous(chunks, seq_len, f"{req_id} layer{layer_no} resid_pre") - db_t = merge_segments([t for _, _, t in chunks], key[1]) - with subtests.test(msg=f"{req_id}/cuda_graph/layer{layer_no}/resid_pre"): - assert tuple(db_t.shape) == tuple(hs.shape), ( - f"resid_pre shape mismatch layer={layer_no} db={tuple(db_t.shape)} hf={tuple(hs.shape)}" - ) - _assert_close_or_bitwise(db_t, hs, f"{req_id} layer{layer_no} resid_pre") - - # attn_pattern: both DB and ref use static cache, same padding - key = next(((layer_no, k) for k in _ATTN_PATTERN_KEYS if (layer_no, k) in hooks_map), None) - if key is not None and ref.attn_pattern and layer_no < len(ref.attn_pattern): - pat = ref.attn_pattern[layer_no] - chunks = _sort_chunks(hooks_map[key]) - db_pat = merge_segments([t for _, _, t in chunks], key[1]) - if db_pat.ndim == 4 and db_pat.shape[0] == 1: - db_pat = db_pat.squeeze(0) - with subtests.test(msg=f"{req_id}/cuda_graph/layer{layer_no}/attn_pattern"): - assert tuple(db_pat.shape) == tuple(pat.shape), ( - f"pattern shape mismatch layer={layer_no} db={tuple(db_pat.shape)} hf={tuple(pat.shape)}" - ) - _assert_close_or_bitwise(db_pat, pat, f"{req_id} layer{layer_no} attn_pattern") - - # --- resid_final --- - # HF's output_hidden_states[-1] is post-final-norm, not pre-norm. - # Only check shape and presence. - key = (-1, "hook_resid_final") - if key in hooks_map: - chunks = _sort_chunks(hooks_map[key]) - _validate_contiguous(chunks, seq_len, f"{req_id} resid_final") - db_t = merge_segments([t for _, _, t in chunks], key[1]) - with subtests.test(msg=f"{req_id}/cuda_graph/resid_final"): - assert db_t.shape[0] == seq_len, ( - f"resid_final token count mismatch db={db_t.shape[0]} expected={seq_len}" - ) - - -# --------------------------------------------------------------------------- -# Test: CUDA-graph monitored DB vs uncompiled eager HF reference -# --------------------------------------------------------------------------- + cr = run_single(matrix_argv_from_env("hf", "allclose", mode="eager")) + _assert_cell(subtests, cr) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA + native backend required") +@require_cuda() +@require_clickhouse() def test_e2e_cuda_graphs_vs_eager_hf(subtests) -> None: """Compare CUDA-graph monitored run against original eager model. diff --git a/tests/test_e2e_lib.py b/tests/test_e2e_lib.py new file mode 100644 index 000000000..acced9b0c --- /dev/null +++ b/tests/test_e2e_lib.py @@ -0,0 +1,369 @@ +"""CPU unit coverage for the shared E2E lib + matrix expansion (plan §7, §8). + +These are pure-CPU (torch-only) — no CUDA / ClickHouse / vLLM / weights — and +guard the de-duplicated comparison/align/report logic plus the matrix's cell +expansion and env translation (exercised via ``--dry-run`` internals). +""" +from __future__ import annotations + +import json + +import pytest +import torch + +from tests.lib import align, compare +from tests.lib.compare import Check +from tests.lib.report import ( + CellResult, + checks_from_legacy_result, + human_table, + read_jsonl, + write_jsonl, +) + +pytestmark = pytest.mark.cpu + + +# --------------------------------------------------------------------------- +# compare.py — the four standards +# --------------------------------------------------------------------------- + + +class TestCompareStandards: + def test_bitwise_equal_records_zero_drift(self): + a = torch.arange(12, dtype=torch.float32).reshape(3, 4) + c = compare.bitwise(a, a.clone(), "x") + assert c.passed and c.max_abs == 0.0 and c.mean_abs == 0.0 + assert c.first_diff_pos is None + + def test_bitwise_detects_diff_and_reports_first_pos(self): + a = torch.arange(6, dtype=torch.float32) + b = a.clone() + b[2] = 99.0 + c = compare.bitwise(a, b, "x") + assert not c.passed + assert c.first_diff_pos == 2 + assert c.max_abs == pytest.approx(99.0 - 2.0) + + def test_bitwise_shape_mismatch_fails_closed(self): + c = compare.bitwise(torch.zeros(4), torch.zeros(5), "x") + assert not c.passed and "shape mismatch" in c.detail + + def test_bitwise_dtype_mismatch_fails_closed(self): + c = compare.bitwise(torch.zeros(4, dtype=torch.float32), + torch.zeros(4, dtype=torch.float16), "x") + assert not c.passed and "dtype mismatch" in c.detail + + def test_allclose_passes_within_tol_but_records_drift(self): + a = torch.zeros(8) + b = a + 1e-4 + c = compare.allclose(a, b, "x", atol=1e-3) + assert c.passed + # drift recorded even on a pass + assert c.max_abs == pytest.approx(1e-4, abs=1e-9) + assert c.mean_abs == pytest.approx(1e-4, abs=1e-9) + + def test_allclose_fails_outside_tol(self): + a = torch.zeros(8) + b = a + 1.0 + c = compare.allclose(a, b, "x", atol=1e-3) + assert not c.passed and c.first_diff_pos == 0 + + def test_transport_bitwise_is_exact(self): + a = torch.randn(4, 4) + assert compare.transport_bitwise(a, a.clone()).passed + assert not compare.transport_bitwise(a, a + 1e-6).passed + + def test_row_count_ok(self): + counts = {f"blocks.hook_{i}": 5 for i in range(12)} + counts["final_logits"] = 3 + c = compare.row_count(counts) + assert c.passed, c.detail + + def test_row_count_uneven_fails(self): + counts = {f"blocks.hook_{i}": 5 for i in range(12)} + counts["blocks.hook_0"] = 4 + counts["final_logits"] = 3 + c = compare.row_count(counts) + assert not c.passed and "uneven" in c.detail + + def test_row_count_missing_final_logits_fails(self): + counts = {f"blocks.hook_{i}": 5 for i in range(12)} + c = compare.row_count(counts) + assert not c.passed and "final_logits" in c.detail + + def test_row_count_too_few_types_fails(self): + c = compare.row_count({"blocks.hook_0": 1, "final_logits": 1}) + assert not c.passed + + def test_compare_tensors_dispatch_and_unknown(self): + a = torch.zeros(3) + assert compare.compare_tensors(a, a.clone(), "bitwise").passed + with pytest.raises(ValueError): + compare.compare_tensors(a, a, "row_count") + + def test_bytes_identical_neg_zero(self): + # -0.0 and 0.0 are equal numerically but differ bitwise. + a = torch.tensor([0.0]) + b = torch.tensor([-0.0]) + assert torch.equal(a, b) # numerically equal + assert not compare.bytes_identical(a, b) # but not byte-identical + + +# --------------------------------------------------------------------------- +# align.py +# --------------------------------------------------------------------------- + + +class TestAlign: + def test_parse_request_id(self): + assert align.parse_request_id("3:7") == (3, 7) + with pytest.raises(ValueError): + align.parse_request_id("nope") + + def test_normalize_request_id_strips_vllm_suffix(self): + assert align.normalize_request_id("12:0-deadbeef") == "12:0" + assert align.normalize_request_id("12:0") == "12:0" + + def test_strip_left_pad(self): + ids = torch.tensor([0, 0, 5, 6, 7]) + attn = torch.tensor([0, 0, 1, 1, 1]) + assert torch.equal(align.strip_left_pad(ids, attn), torch.tensor([5, 6, 7])) + + def test_strip_left_pad_all_padding(self): + ids = torch.tensor([0, 0]) + attn = torch.tensor([0, 0]) + assert align.strip_left_pad(ids, attn).numel() == 0 + + def test_trim_eos_drops_and_keeps(self): + ids = torch.tensor([1, 2, 9, 3]) + assert torch.equal(align.trim_eos(ids, 9), torch.tensor([1, 2])) + assert torch.equal(align.trim_eos(ids, 9, keep_eos=True), torch.tensor([1, 2, 9])) + + def test_trim_eos_absent_returns_all(self): + ids = torch.tensor([1, 2, 3]) + assert torch.equal(align.trim_eos(ids, 9), ids) + + def test_align_to_min_len(self): + a = torch.arange(5) + b = torch.arange(3) + ra, rb = align.align_to_min_len(a, b) + assert ra.numel() == rb.numel() == 3 + + def test_logits_align_skip(self): + # skip = max(0, db_len - ref_len - 1): the DB keeps every position + # while generate() yields gen-1 rows, so the head offset is the + # prefill span. Never negative when ref is longer than db. + assert align.logits_align_skip(db_len=9, ref_len=4) == 4 + assert align.logits_align_skip(db_len=2, ref_len=5) == 0 + + +# --------------------------------------------------------------------------- +# report.py +# --------------------------------------------------------------------------- + + +class TestReport: + def test_cellresult_finalize_and_record(self): + cr = CellResult("vllm", "qwen3", "eager", "transport_bitwise", "vllm-full", tp=1) + cr.checks = [Check("token_ids", True), Check("layer.0.q", True, max_abs=0.0)] + cr.finalize() + rec = cr.to_record() + assert rec["passed"] is True + assert rec["backend"] == "vllm" and rec["hook_selection"] == "vllm-full" + assert len(rec["checks"]) == 2 + assert rec["checks"][1]["max_abs"] == 0.0 + + def test_cellresult_fails_when_any_check_fails(self): + cr = CellResult("hf", "gpt2", "eager", "bitwise", "vllm-full") + cr.checks = [Check("a", True), Check("b", False, detail="boom")] + cr.finalize() + assert cr.passed is False + + def test_cellresult_no_checks_is_not_passed(self): + cr = CellResult("hf", "gpt2", "eager", "bitwise", "vllm-full").finalize() + assert cr.passed is False + + def test_cellresult_error_sets_failed(self): + cr = CellResult("hf", "gpt2", "eager", "bitwise", "vllm-full") + cr.error = "runner crashed" + cr.checks = [Check("a", True)] + cr.finalize() + assert cr.passed is False + assert cr.to_record()["error"] == "runner crashed" + + def test_legacy_result_adapter(self): + legacy = {"tests": [ + {"name": "rows_found", "passed": True, "detail": "10 rows"}, + {"name": "x", "passed": False, "detail": "max_abs=1e-3"}, + ]} + checks = checks_from_legacy_result(legacy) + assert [c.name for c in checks] == ["rows_found", "x"] + assert checks[0].passed and not checks[1].passed + + def test_jsonl_roundtrip(self, tmp_path): + crs = [] + for passed in (True, False): + cr = CellResult("vllm", "gpt2", "eager", "row_count", "vllm-full") + cr.checks = [Check("c", passed)] + crs.append(cr.finalize()) + out = tmp_path / "e2e.jsonl" + write_jsonl(crs, str(out)) + recs = read_jsonl(str(out)) + assert len(recs) == 2 + assert recs[0]["passed"] is True and recs[1]["passed"] is False + # each line is valid standalone JSON + for line in out.read_text().splitlines(): + json.loads(line) + + def test_human_table_renders_counts(self): + crs = [ + CellResult("vllm", "gpt2", "eager", "bitwise", "vllm-full").finalize(), + ] + crs[0].checks = [Check("c", True)] + crs[0].finalize() + table = human_table(crs) + assert "backend" in table and "1/1 cells passed" in table + + +# --------------------------------------------------------------------------- +# e2e_matrix — cell expansion + env translation (the dry-run surface) +# --------------------------------------------------------------------------- + + +class TestMatrixExpansion: + def _args(self, **over): + from tests.e2e_matrix import build_parser + argv = [] + for k, v in over.items(): + argv += [f"--{k.replace('_', '-')}", str(v)] + return build_parser().parse_args(argv) + + def test_cartesian_product_count(self): + from tests.e2e_matrix import build_cells + args = self._args(backend="hf,vllm", model="gpt2,qwen3", + mode="eager,cuda_graph", standard="row_count") + cells = build_cells(args) + assert len(cells) == 2 * 2 * 2 * 1 + + def test_env_translates_hook_selection(self): + from tests.e2e_matrix import build_cells, cell_env + args = self._args(backend="vllm", model="qwen3", mode="cuda_graph", + standard="transport_bitwise", hooks="q") + cell = build_cells(args)[0] + env = cell_env(cell, args, base={}) + # public -> internal contract (plan §2) + assert env["E2E_HOOK_SELECTION"] == "q" + assert env["DMX_HOOK_SELECTION"] == "q" + # cuda_graph -> not eager + assert env["E2E_ENFORCE_EAGER"] == "0" + assert env["E2E_CUDA_GRAPHS"] == "1" + assert env["E2E_MODEL"] == "qwen3" + + def test_eager_sets_enforce_eager(self): + from tests.e2e_matrix import build_cells, cell_env + args = self._args(backend="hf", model="gpt2", mode="eager", standard="bitwise") + env = cell_env(build_cells(args)[0], args, base={}) + assert env["E2E_ENFORCE_EAGER"] == "1" + assert env["E2E_CUDA_GRAPHS"] == "0" + + def test_plan_vllm_identical_dispatch(self): + from tests.e2e_matrix import build_cells, plan_cell + args = self._args(backend="vllm", model="qwen3", mode="eager", + standard="transport_bitwise", hooks="vllm-full") + steps, comparator, _result = plan_cell(build_cells(args)[0], "/run") + labels = [s.label for s in steps] + assert labels == ["enable_ref_hooks", "vllm_ref", "vllm_monitored", "compare"] + assert comparator == "tests.vllm_identical_comparator" + + def test_plan_vllm_rowcount_dispatch(self): + from tests.e2e_matrix import build_cells, plan_cell + args = self._args(backend="vllm", model="gpt2", mode="eager", + standard="row_count") + steps, comparator, _ = plan_cell(build_cells(args)[0], "/run") + assert [s.label for s in steps] == ["vllm_monitored", "compare"] + assert comparator == "tests.vllm_rowcnt_comparator" + + def test_plan_hf_dispatch(self): + from tests.e2e_matrix import build_cells, plan_cell + args = self._args(backend="hf", model="gpt2", mode="eager", standard="allclose") + steps, comparator, _ = plan_cell(build_cells(args)[0], "/run") + assert [s.label for s in steps] == ["hf_ref", "hf_monitored", "compare"] + assert comparator == "tests.hf_comparator" + + def test_unknown_backend_raises(self): + from tests.e2e_matrix import Cell, plan_cell + with pytest.raises(ValueError): + plan_cell(Cell("nope", "gpt2", "eager", "bitwise", "vllm-full"), "/run") + + def test_main_dry_run_no_side_effects(self, capsys): + # hf only supports allclose; expand over two models to get 2 cells. + from tests.e2e_matrix import main + rc = main(["--backend", "hf", "--model", "gpt2,qwen3", + "--standard", "allclose", "--dry-run"]) + assert rc == 0 + out = capsys.readouterr().out + assert "2 cell(s) planned" in out + + def test_main_empty_axis_returns_2(self): + from tests.e2e_matrix import main + assert main(["--backend", "", "--dry-run"]) == 2 + + +class TestWrapperTranslation: + """matrix_argv_from_env / run_single — the thin-wrapper surface (plan §5).""" + + def test_env_to_argv_defaults(self): + from tests.e2e_matrix import matrix_argv_from_env, build_parser, build_cells + argv = matrix_argv_from_env("vllm", "bitwise", env={}) + cell = build_cells(build_parser().parse_args(argv))[0] + assert cell.backend == "vllm" and cell.standard == "bitwise" + assert cell.model == "gpt2" and cell.mode == "eager" + assert cell.hooks == "vllm-full" and cell.ring_mb == 256 + + def test_env_enforce_eager_maps_mode(self): + from tests.e2e_matrix import matrix_argv_from_env, build_parser, build_cells + env = {"E2E_ENFORCE_EAGER": "0", "E2E_MODEL": "qwen3"} + cell = build_cells(build_parser().parse_args( + matrix_argv_from_env("vllm", "row_count", env=env)))[0] + assert cell.mode == "cuda_graph" and cell.model == "qwen3" + + def test_explicit_mode_overrides_enforce_eager(self): + from tests.e2e_matrix import matrix_argv_from_env, build_parser, build_cells + # HF cuda-graph wrapper forces mode even though E2E_ENFORCE_EAGER=1. + env = {"E2E_ENFORCE_EAGER": "1"} + cell = build_cells(build_parser().parse_args( + matrix_argv_from_env("hf", "allclose", mode="cuda_graph", env=env)))[0] + assert cell.mode == "cuda_graph" + + def test_hook_selection_precedence(self): + from tests.e2e_matrix import matrix_argv_from_env + # public E2E_HOOK_SELECTION wins over internal DMX_HOOK_SELECTION + argv = matrix_argv_from_env("vllm", "bitwise", env={ + "E2E_HOOK_SELECTION": "q", "DMX_HOOK_SELECTION": "k"}) + assert argv[argv.index("--hooks") + 1] == "q" + # falls back to DMX_HOOK_SELECTION when public unset + argv = matrix_argv_from_env("vllm", "bitwise", env={"DMX_HOOK_SELECTION": "k"}) + assert argv[argv.index("--hooks") + 1] == "k" + + def test_default_tolerance_passthrough(self): + from tests.e2e_matrix import matrix_argv_from_env + argv = matrix_argv_from_env("hf", "allclose", mode="cuda_graph", + default_tolerance="0.5", env={}) + assert argv[argv.index("--tolerance") + 1] == "0.5" + # explicit env overrides the default + argv = matrix_argv_from_env("hf", "allclose", default_tolerance="0.5", + env={"E2E_TOLERANCE": "0.01"}) + assert argv[argv.index("--tolerance") + 1] == "0.01" + + def test_run_single_rejects_multi_cell(self, monkeypatch): + # matrix_argv_from_env always yields one cell; guard the invariant. + from tests import e2e_matrix + args = e2e_matrix.build_parser().parse_args( + ["--backend", "hf,vllm", "--standard", "row_count"]) + monkeypatch.setattr(e2e_matrix, "run_cell", lambda *a, **k: None) + with pytest.raises(ValueError, match="exactly 1 cell"): + # build_cells gives 2 -> run_single must refuse + cells = e2e_matrix.build_cells(args) + assert len(cells) == 2 + e2e_matrix.run_single(["--backend", "hf,vllm", "--standard", "row_count"]) diff --git a/tests/test_e2e_tp2.py b/tests/test_e2e_tp2.py new file mode 100644 index 000000000..3b52e3369 --- /dev/null +++ b/tests/test_e2e_tp2.py @@ -0,0 +1,78 @@ +"""Multi-GPU (TP=2) E2E smoke -- the ``multi_gpu`` suite's TP coverage. + +The configurable matrix (:mod:`tests.e2e_matrix`) treats tensor-parallel size +as a first-class axis (``--tp`` / ``E2E_TP_SIZE``), but the single-GPU wrappers +(``test_vllm_identical`` / ``test_e2e_correctness_vs_hf``) all drive ``tp=1``. +This module drives vLLM matrix cells at ``tp=2`` so the documented +``-m multi_gpu`` suite (docs/testing.md) actually exercises TP sharding rather +than collecting nothing. + +TP=2 is "where meaningful" for the sharded model: with two ranks the +attention/expert projections are split across GPUs, so the reference-vs-monitored +comparison validates that the ring transport reassembles per-rank shards +correctly. ``qwen3`` is the default model (GQA + a non-trivial hidden size makes +the sharding observable); ``gpt2`` is too small for TP to be interesting. The +model is still overridable via ``E2E_MODEL``. + +HF TP=2 is not covered here: ``hf_reference_runner`` / ``hf_monitored_runner`` +do not read ``E2E_TP_SIZE``, do not launch under ``torchrun``, and do not pass +``tp_plan="auto"``. A test claiming HF TP=2 coverage would require 2 GPUs +without actually exercising tensor parallelism, which would be misleading. + +Skip-guarded (``tests/_requirements``) so a runner with <2 GPUs, no vLLM, or no +ClickHouse skips with a reason instead of failing the job. +""" +from __future__ import annotations + +import os + +import pytest + +from tests._requirements import ( + require_clickhouse, + require_gpus, + require_vllm, +) +from tests.e2e_matrix import matrix_argv_from_env, run_single + +pytestmark = [ + pytest.mark.multi_gpu, + pytest.mark.gpu, + pytest.mark.e2e, + pytest.mark.clickhouse, +] + + +def _tp2_env(default_model: str = "qwen3") -> dict: + """os.environ with TP forced to 2 (model still overridable via E2E_MODEL).""" + env = dict(os.environ) + env["E2E_TP_SIZE"] = "2" + env.setdefault("E2E_MODEL", default_model) + return env + + +def _assert_cell(subtests, cr) -> None: + """Fail on a dispatch error; report each matrix check as a subtest.""" + if cr.error: + pytest.fail(f"matrix cell errored: {cr.error}") + assert cr.checks, "matrix produced no checks" + for chk in cr.checks: + with subtests.test(chk.name): + assert chk.passed, chk.detail + + +@pytest.mark.vllm +@require_gpus(2) +@require_vllm() +@require_clickhouse() +def test_vllm_identical_tp2(subtests) -> None: + """vLLM TP=2 transport-bitwise: reference D2D buffers vs ring -> ClickHouse. + + Equivalent matrix cell: ``--backend vllm --standard transport_bitwise --tp 2``. + The bitwise standard stays exact under TP (sharding is a layout change, not a + numeric one), so any per-rank reassembly bug surfaces as a non-zero max_abs. + """ + argv = matrix_argv_from_env("vllm", "transport_bitwise", env=_tp2_env()) + _assert_cell(subtests, run_single(argv)) + + diff --git a/tests/test_numeric_study.py b/tests/test_numeric_study.py new file mode 100644 index 000000000..0a9c3ff07 --- /dev/null +++ b/tests/test_numeric_study.py @@ -0,0 +1,217 @@ +"""Tests for the per-hook numeric-difference study (plan §9 / Phase 5). + +Two layers, mirroring ``test_per_hook_isolation.py``: + +1. **CPU unit tests** (always run): exercise the pure comparison / alert / + report core (``compute_drift``, ``format_table``, serialization, threshold + + standard selection) on small hand-built CPU tensors. These guard the + study's verdict logic -- including its use of the shared ``tests.lib.compare`` + standards -- independently of any GPU rollout. + +2. **GPU smoke** (marked ``numeric`` + ``gpu`` + ``slow``, opt-in): run the + real study for a couple of observational hooks and assert the machinery + produces serializable per-hook records and a self-consistent eager verdict. +""" +from __future__ import annotations + +import json + +import pytest + +from tests._requirements import require_cuda +from tests.lib.compare import Check +from tests.numeric_study import ( + HookDrift, + StudyResult, + VocabDiff, + compute_drift, + cuda_graph_threshold, + format_table, + run_study, + standard_for_mode, +) + + +# --------------------------------------------------------------------------- +# CPU unit tests for the comparison core +# --------------------------------------------------------------------------- + + +def _cap(logits): + import torch + + t = torch.tensor(logits, dtype=torch.float32) + return {"logits": t, "token_ids": t.argmax(dim=-1).to(torch.int64)} + + +@pytest.mark.cpu +class TestComputeDrift: + def test_identical_eager_passes_bitwise(self): + base = _cap([[1.0, 2.0, 0.5], [0.1, 0.2, 3.0]]) + d = compute_drift("q", base, base, mode="eager", threshold=0.15, topk=3) + assert d.check is not None and d.check.name == "logits_bitwise" + assert d.check.passed is True + assert d.check.max_abs == 0.0 + assert d.first_diff_pos == -1 + assert d.token_ids_diverged is False + assert d.topk_vocab_diffs == [] + assert d.alert is False + assert d.n_positions == 2 and d.vocab_size == 3 + + def test_eager_any_drift_alerts(self): + base = _cap([[5.0, 1.0, 0.0], [0.0, 4.0, 1.0]]) + mon = _cap([[5.0, 1.0, 0.0], [0.0, 4.0, 1.25]]) # tiny perturb, no argmax flip + d = compute_drift("k", base, mon, mode="eager", threshold=0.15, topk=2) + assert d.check.passed is False + assert d.first_diff_pos == 1 + assert d.check.max_abs == pytest.approx(0.25, abs=1e-6) + assert d.token_ids_diverged is False + assert d.alert is True + assert any("eager" in r for r in d.alert_reasons) + # Top-k drift at the first differing position points at the perturbed id. + assert d.topk_vocab_diffs[0].token_id == 2 + assert d.topk_vocab_diffs[0].abs_diff == pytest.approx(0.25, abs=1e-6) + + def test_cuda_graph_below_threshold_no_alert(self): + base = _cap([[5.0, 1.0, 0.0]]) + mon = _cap([[5.0, 1.05, 0.0]]) # max_abs 0.05 < 0.15 + d = compute_drift("q", base, mon, mode="cuda_graph", threshold=0.15, topk=1) + assert d.check.name == "logits_allclose" + assert d.check.passed is True + assert d.alert is False + + def test_cuda_graph_above_threshold_alerts(self): + base = _cap([[5.0, 1.0, 0.0]]) + mon = _cap([[5.0, 1.5, 0.0]]) # max_abs 0.5 > 0.15, no argmax flip + d = compute_drift("q", base, mon, mode="cuda_graph", threshold=0.15, topk=1) + assert d.check.passed is False + assert d.alert is True + assert any("threshold" in r for r in d.alert_reasons) + + def test_argmax_flip_always_alerts_even_under_cuda_graph(self): + base = _cap([[5.0, 1.0, 0.0]]) + mon = _cap([[1.0, 9.0, 0.0]]) # argmax 0 -> 1 + d = compute_drift("resid_pre", base, mon, mode="cuda_graph", threshold=100.0, topk=1) + assert d.token_ids_diverged is True + assert d.n_token_diff == 1 + assert d.alert is True + assert any("token ids diverged" in r for r in d.alert_reasons) + + def test_shape_mismatch_alerts(self): + import torch + + base = _cap([[1.0, 2.0, 3.0]]) + mon = { + "logits": torch.zeros((1, 4), dtype=torch.float32), + "token_ids": torch.zeros((1,), dtype=torch.int64), + } + d = compute_drift("q", base, mon, mode="eager", threshold=0.15, topk=1) + assert d.shape_mismatch is True + assert d.alert is True + assert any("identity swap" in r for r in d.alert_reasons) + + def test_empty_capture_alerts(self): + import torch + + empty = { + "logits": torch.zeros((0, 3), dtype=torch.float32), + "token_ids": torch.zeros((0,), dtype=torch.int64), + } + d = compute_drift("q", empty, empty, mode="eager", threshold=0.15, topk=1) + assert d.error is not None + assert d.alert is True + + def test_missing_logits_alerts(self): + d = compute_drift("q", {"token_ids": None}, {"token_ids": None}, + mode="eager", threshold=0.15, topk=1) + assert d.error is not None + assert d.alert is True + + +@pytest.mark.cpu +class TestSelectorsAndReport: + def test_standard_for_mode(self): + assert standard_for_mode("eager") == "bitwise" + assert standard_for_mode("cuda_graph") == "allclose" + + def test_threshold_lookup_and_fallback(self): + assert cuda_graph_threshold("qwen2_moe", "float16") == 0.25 + assert cuda_graph_threshold("nonesuch", "float64") == pytest.approx(0.15) + + def test_format_table_contains_hooks_and_verdict(self): + result = StudyResult( + framework="hf", model="qwen3", mode="eager", variant="p", + dtype="float16", standard="bitwise", threshold=0.15, topk=2, + ) + ok = HookDrift(hook="resid_pre", check=Check("logits_bitwise", True, max_abs=0.0)) + bad = HookDrift( + hook="q", + check=Check("logits_bitwise", False, max_abs=0.3, mean_abs=0.01, detail="max_abs=3e-01"), + first_diff_pos=1, alert=True, alert_reasons=["non-bitwise drift in eager mode"], + topk_vocab_diffs=[VocabDiff(7, 1.0, 1.3, 0.3)], + ) + result.hooks.extend([ok, bad]) + table = format_table(result) + assert "resid_pre" in table + assert "ALERT" in table + assert "tok 7" in table + assert "1/2 hooks clean" in table + + def test_result_to_dict_is_json_serializable(self): + result = StudyResult( + framework="hf", model="qwen3", mode="cuda_graph", variant="compare", + dtype="float16", standard="allclose", threshold=0.15, topk=1, + ) + result.hooks.append( + HookDrift( + hook="q", + check=Check("logits_allclose", False, max_abs=0.2, mean_abs=0.01), + topk_vocab_diffs=[VocabDiff(1, 0.0, 0.2, 0.2)], + alert=True, alert_reasons=["exceeds cuda-graph threshold"], + ) + ) + payload = result.to_dict() + s = json.dumps(payload) # must not raise + rt = json.loads(s) + assert rt["any_alert"] is True + assert rt["variant"] == "compare" + assert rt["hooks"][0]["check"]["name"] == "logits_allclose" + assert rt["hooks"][0]["topk_vocab_diffs"][0]["token_id"] == 1 + + +# --------------------------------------------------------------------------- +# GPU smoke +# --------------------------------------------------------------------------- + + +@pytest.mark.numeric +@pytest.mark.gpu +@pytest.mark.slow +@require_cuda() +def test_numeric_study_hf_qwen3_eager_smoke(tmp_path): + """Eager observational hooks must not drift the logits vs the unhooked + baseline; the study reports per-hook records and the eager verdict is + self-consistent (alert iff the bitwise check failed).""" + hooks = ["resid_pre", "final_logits"] + result = run_study( + framework="hf", model="qwen3", mode="eager", + hooks=hooks, out_dir=tmp_path, max_new_tokens=4, + ) + + # Always surface the table, even on pass. + print("\n" + format_table(result)) + + assert [h.hook for h in result.hooks] == hooks + json.dumps(result.to_dict()) # serializable + + for h in result.hooks: + assert h.error is None, f"{h.hook} rollout errored: {h.error}" + assert h.check is not None + # Eager verdict must be self-consistent with the bitwise outcome. + assert h.alert == (not h.check.passed), ( + f"{h.hook}: alert={h.alert} but check.passed={h.check.passed}" + ) + assert not result.any_alert, ( + "eager observational hooks drifted vs the unhooked baseline:\n" + + format_table(result) + ) diff --git a/tests/test_per_hook_isolation.py b/tests/test_per_hook_isolation.py index e718d224b..8fcbfa8e6 100644 --- a/tests/test_per_hook_isolation.py +++ b/tests/test_per_hook_isolation.py @@ -126,10 +126,10 @@ def test_indentation_preserved_on_commented_lines(self): assert indent == 8, f"unexpected indent on: {line!r}" -@pytest.mark.cpu class TestPatcherRoundTrip: """Verify the on-disk patch / unpatch context manager.""" + @pytest.mark.framework_fork @pytest.mark.parametrize("framework,model_key", [ ("hf", "gpt2"), ("hf", "qwen3"), ("hf", "llama"), ("vllm", "gpt2"), ("vllm", "qwen3"), ("vllm", "llama"), @@ -148,6 +148,7 @@ def test_round_trip_byte_identical(self, framework, model_key): backup = p.with_suffix(p.suffix + ".copy_isolate_backup") assert not backup.exists() + @pytest.mark.cpu def test_dirty_restore_raises_loudly(self, tmp_path, monkeypatch): """If the file can't be restored byte-identically, exit must raise. @@ -170,6 +171,7 @@ def test_dirty_restore_raises_loudly(self, tmp_path, monkeypatch): backup = target.with_suffix(target.suffix + ".copy_isolate_backup") backup.unlink() + @pytest.mark.cpu def test_stale_backup_raises(self, tmp_path, monkeypatch): """If a previous run crashed mid-patch leaving a backup, refuse.""" from tests import isolate_hook @@ -186,7 +188,7 @@ def test_stale_backup_raises(self, tmp_path, monkeypatch): isolate_hook.patch("test", "fake", "q") -@pytest.mark.cpu +@pytest.mark.framework_fork class TestRealCompareModelsContainAllExpectedHooks: """The smoke cells assume specific hooks have a `.copy_()` line in the real _compare files. If a hook is missing the smoke fails opaquely, @@ -409,6 +411,7 @@ def test_smoke_hooks_present_in_compare_source(self, framework, model_key): model=model_name, max_model_len=128, gpu_memory_utilization=0.5, + max_num_seqs=32, enforce_eager=(args.mode == 'eager'), ) diff --git a/tests/test_vllm_identical.py b/tests/test_vllm_identical.py index 41087a516..1b827d809 100644 --- a/tests/test_vllm_identical.py +++ b/tests/test_vllm_identical.py @@ -1,44 +1,28 @@ -"""vLLM identical check — bitwise tensor comparison between ref model -(GPU buffer D2D capture → disk) and monitored model (ring transport → ClickHouse). +"""vLLM identical check — thin wrapper over the configurable matrix (plan §5). -Four steps, parent never touches CUDA: - 0. Sanity check: original vs ref model logprobs (informational, never fails) - 1. Reference run (RefDiskWorker, D2D capture → disk) - 2. Monitored run (DMXGPUWorker, ring transport → ClickHouse) - 3. Comparator (CPU only, logprob comparison + bitwise tensor check) +Bitwise tensor comparison between the reference model (GPU buffer D2D +capture -> disk) and the monitored model (ring transport -> ClickHouse). +The orchestration (enable_ref_hooks -> vllm_ref_runner -> vllm_monitored_runner +-> vllm_identical_comparator) now lives in :mod:`tests.e2e_matrix`; this +wrapper drives the matrix's vLLM ``bitwise`` cell and asserts on its checks. -Environment variables: - E2E_MODEL "gpt2" (default) or "qwen3" - E2E_NUM_PROMPTS Number of prompts (default 8) - E2E_MAX_NEW_TOKENS Tokens to generate per prompt (default 20) - E2E_ENFORCE_EAGER "1" to disable CUDA graphs (default "1") - E2E_DTYPE Model dtype, e.g. "bfloat16", "float16", "auto" (default "bfloat16") - E2E_REF_MAX_LEN Max first-dim for buffers (default 8192) - E2E_MAX_NUM_BATCHED_TOKENS vLLM scheduler max_num_batched_tokens (default 512) - E2E_RING_PAYLOAD_MB Ring payload size (default 4096) - E2E_RING_PINNED_MB Pinned staging size (default 4096) - E2E_HOOK_SELECTION Public hook selection input (default "vllm-full"); - translated to DMX_HOOK_SELECTION for subprocesses - DMX_DB_HOST ClickHouse host (default "localhost") - DMX_DB_PORT ClickHouse port (default 9000) +The test name is preserved because ``tests/tools/verify_vllm.sh`` and +``tests/tools/identical_vllm.sh`` invoke this file and thread the model / +ring-size / hook-selection env vars the matrix wrapper reads: -Requires: - - ClickHouse running - - VLLM_DISABLE_COMPILE_CACHE=1 (set automatically) + E2E_MODEL, E2E_ENFORCE_EAGER, E2E_DTYPE, E2E_RING_PAYLOAD_MB, + E2E_RING_PINNED_MB, E2E_HOOK_SELECTION (-> internal DMX_HOOK_SELECTION), + E2E_REF_MAX_LEN, E2E_TP_SIZE, DMX_DB_HOST, DMX_DB_PORT. Usage: python -m pytest tests/test_vllm_identical.py -q -s """ - -import json -import os -import shutil -import subprocess -import sys -import tempfile +from __future__ import annotations import pytest -import torch + +from tests._requirements import require_cuda, require_clickhouse, require_vllm +from tests.e2e_matrix import matrix_argv_from_env, run_single pytestmark = [ pytest.mark.gpu, @@ -47,215 +31,16 @@ pytest.mark.e2e, ] -_MODEL_REF_FILES = { - "gpt2": "gpt2_ref.py", - "qwen2_moe": "qwen2_moe_ref.py", - "qwen3": "qwen3_ref.py", - "llama": "llama_ref.py", -} - - -@pytest.mark.skipif( - not torch.backends.cuda.is_built(), reason="CUDA not built") -def test_vllm_identical(subtests): - """Bitwise comparison: ref model (disk) vs monitored model (ClickHouse).""" - - model_key = os.environ.get("E2E_MODEL", "gpt2") - hooks = os.environ.get("E2E_HOOK_SELECTION", "vllm-full") - max_len = int(os.environ.get("E2E_REF_MAX_LEN", "8192")) - enforce_eager = os.environ.get("E2E_ENFORCE_EAGER", "1") - - project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - models_dir = os.path.join( - project_root, "integration", "vllm", "vllm", - "model_executor", "models") - - ref_filename = _MODEL_REF_FILES.get(model_key) - if ref_filename is None: - pytest.skip(f"No ref model for {model_key}") - model_file = os.path.join(models_dir, ref_filename) - - keep_artifacts = os.environ.get("E2E_KEEP_ARTIFACTS", "0") == "1" - dump_compiled = os.environ.get("E2E_DUMP_COMPILED", "0") == "1" - artifact_dir = os.environ.get("E2E_ARTIFACT_DIR") - if artifact_dir: - run_dir = os.path.abspath(artifact_dir) - os.makedirs(run_dir, exist_ok=True) - else: - run_dir = tempfile.mkdtemp(prefix="vllm_identical_") - ref_dir = os.path.join(run_dir, "ref") - mon_dir = os.path.join(run_dir, "mon") - config_file = os.path.join(ref_dir, "ref_config.json") - result_file = os.path.join(run_dir, "result.json") - backup_file = os.path.join(run_dir, f"{ref_filename}.bak") - orig_logprobs_file = os.path.join(run_dir, "logprobs_orig.pt") - ref_logprobs_file = os.path.join(run_dir, "logprobs_ref.pt") - - print(f"\n{'=' * 60}") - print(f" vLLM identical check") - print(f" model={model_key} hooks={hooks} eager={enforce_eager}") - print(f" ref_max_len={max_len}") - print(f"{'=' * 60}") - - # Build env for subprocesses (inherit + add our vars) - sub_env = dict(os.environ) - sub_env["VLLM_DISABLE_COMPILE_CACHE"] = "1" - sub_env["E2E_ENFORCE_EAGER"] = enforce_eager - sub_env["DMX_HOOK_SELECTION"] = hooks - if dump_compiled: - sub_env["TORCH_LOGS"] = "+output_code" - - try: - # Backup ref model file - shutil.copy2(model_file, backup_file) - - # Enable hooks via preprocessor - print("\n [0/3] Enabling ref hooks...", flush=True) - os.makedirs(ref_dir, exist_ok=True) - sys.path.insert(0, models_dir) - from enable_ref_hooks import enable_ref_hooks - enable_ref_hooks( - model_file=model_file, - hooks=hooks, - max_len=max_len, - output_dir=ref_dir, - config_out=config_file, - ) - - # Step 0: Sanity check — original vs ref model logprobs - # Runs AFTER enabling hooks to verify D2D copies don't affect output. - print("\n [0/4] Sanity check: original model logprobs...", flush=True) - r0a = subprocess.run( - [sys.executable, "-m", "tests.vllm_logprob_runner", - "--output", orig_logprobs_file], - env=sub_env, capture_output=True, text=True, cwd=project_root, - ) - print(r0a.stdout[-1000:] if r0a.stdout else "", flush=True) - if dump_compiled and r0a.stderr: - with open(os.path.join(run_dir, "compile_orig.log"), "w") as f: - f.write(r0a.stderr) - if r0a.returncode != 0: - print(r0a.stderr[-2000:] if r0a.stderr else "", flush=True) - print(" WARNING: original logprob run failed, skipping sanity check") - orig_logprobs_file = None - - print(" [0/4] Sanity check: ref model logprobs...", flush=True) - ref_lp_env = dict(sub_env) - ref_lp_env["REF_CONFIG"] = config_file - r0b = subprocess.run( - [sys.executable, "-m", "tests.vllm_logprob_runner", - "--output", ref_logprobs_file, "--ref"], - env=ref_lp_env, capture_output=True, text=True, cwd=project_root, - ) - print(r0b.stdout[-1000:] if r0b.stdout else "", flush=True) - if dump_compiled and r0b.stderr: - with open(os.path.join(run_dir, "compile_ref_logprob.log"), "w") as f: - f.write(r0b.stderr) - if r0b.returncode != 0: - print(r0b.stderr[-2000:] if r0b.stderr else "", flush=True) - print(" WARNING: ref logprob run failed, skipping sanity check") - ref_logprobs_file = None - - # Step 0c: Monitored model logprobs (baseline vs monitored comparison) - mon_logprobs_file = os.path.join(run_dir, "logprobs_mon.pt") - print(" [0/4] Sanity check: monitored model logprobs...", flush=True) - r0c = subprocess.run( - [sys.executable, "-m", "tests.vllm_logprob_runner", - "--output", mon_logprobs_file, "--monitored"], - env=sub_env, capture_output=True, text=True, cwd=project_root, - ) - print(r0c.stdout[-1000:] if r0c.stdout else "", flush=True) - if dump_compiled and r0c.stderr: - with open(os.path.join(run_dir, "compile_mon_logprob.log"), "w") as f: - f.write(r0c.stderr) - if r0c.returncode != 0: - print(r0c.stderr[-2000:] if r0c.stderr else "", flush=True) - print(" WARNING: monitored logprob run failed, skipping") - mon_logprobs_file = None - - # Step 1: Reference run - print("\n [1/4] Reference run (RefDiskWorker)...", flush=True) - ref_env = dict(sub_env) - ref_env["REF_CONFIG"] = config_file - r1 = subprocess.run( - [sys.executable, "-m", "tests.vllm_ref_runner", - "--output-dir", ref_dir], - env=ref_env, capture_output=True, text=True, cwd=project_root, - ) - print(r1.stdout[-2000:] if r1.stdout else "", flush=True) - if keep_artifacts and r1.stdout: - with open(os.path.join(run_dir, "stdout_ref_runner.log"), "w") as f: - f.write(r1.stdout) - if dump_compiled and r1.stderr: - with open(os.path.join(run_dir, "compile_ref_runner.log"), "w") as f: - f.write(r1.stderr) - if r1.returncode != 0: - print(r1.stderr[-3000:] if r1.stderr else "", flush=True) - pytest.fail(f"Ref runner failed (rc={r1.returncode})") - - # Restore ref model from backup (before monitored run) - shutil.copy2(backup_file, model_file) - - # Step 2: Monitored run - print("\n [2/4] Monitored run (DMXGPUWorker)...", flush=True) - r2 = subprocess.run( - [sys.executable, "-m", "tests.vllm_monitored_runner", - "--output-dir", mon_dir], - env=sub_env, capture_output=True, text=True, cwd=project_root, - ) - print(r2.stdout[-2000:] if r2.stdout else "", flush=True) - if keep_artifacts and r2.stdout: - with open(os.path.join(run_dir, "stdout_mon_runner.log"), "w") as f: - f.write(r2.stdout) - if dump_compiled and r2.stderr: - with open(os.path.join(run_dir, "compile_mon_runner.log"), "w") as f: - f.write(r2.stderr) - if r2.returncode != 0: - print(r2.stderr[-3000:] if r2.stderr else "", flush=True) - pytest.fail(f"Monitored runner failed (rc={r2.returncode})") - - # Step 3: Comparator (includes logprob sanity check if available) - print("\n [3/4] Comparing (bitwise check)...", flush=True) - cmp_cmd = [ - sys.executable, "-m", "tests.vllm_identical_comparator", - "--ref-config", config_file, - "--mon-dir", mon_dir, - "--result-file", result_file, - ] - if orig_logprobs_file and os.path.exists(orig_logprobs_file): - cmp_cmd += ["--orig-logprobs", orig_logprobs_file] - if ref_logprobs_file and os.path.exists(ref_logprobs_file): - cmp_cmd += ["--ref-logprobs", ref_logprobs_file] - if mon_logprobs_file and os.path.exists(mon_logprobs_file): - cmp_cmd += ["--mon-logprobs", mon_logprobs_file] - r3 = subprocess.run( - cmp_cmd, - env=sub_env, capture_output=True, text=True, cwd=project_root, - ) - if r3.stdout: - # Always print LOGPROBS summary and PASS/FAIL lines first - for line in r3.stdout.splitlines(): - if "[LOGPROBS" in line or "ALL PASSED" in line or "FAILED (" in line: - print(line, flush=True) - # Then print tail for hidden state details - print(r3.stdout[-2000:], flush=True) - if r3.returncode != 0: - print(r3.stderr[-3000:] if r3.stderr else "", flush=True) - pytest.fail(f"Comparator failed (rc={r3.returncode})") - - # Report via subtests - with open(result_file) as f: - results = json.load(f) - - for test in results["tests"]: - with subtests.test(test["name"]): - assert test["passed"], test.get("detail", "") - finally: - # Always restore ref model file - if os.path.exists(backup_file): - shutil.copy2(backup_file, model_file) - if keep_artifacts: - print(f"\n [kept] run_dir = {run_dir}", flush=True) - else: - shutil.rmtree(run_dir, ignore_errors=True) +@require_cuda() +@require_vllm() +@require_clickhouse() +def test_vllm_identical(subtests) -> None: + """Bitwise: reference D2D buffers (disk) vs ring transport (ClickHouse).""" + cr = run_single(matrix_argv_from_env("vllm", "bitwise")) + if cr.error: + pytest.fail(f"matrix cell errored: {cr.error}") + assert cr.checks, "matrix produced no checks" + for chk in cr.checks: + with subtests.test(chk.name): + assert chk.passed, chk.detail diff --git a/tests/test_vllm_rowcnt.py b/tests/test_vllm_rowcnt.py index da0d466be..ee9092558 100644 --- a/tests/test_vllm_rowcnt.py +++ b/tests/test_vllm_rowcnt.py @@ -1,45 +1,27 @@ -"""vLLM E2E correctness test — three subprocesses, parent never touches CUDA. +"""vLLM row-count check — thin wrapper over the configurable matrix (plan §5). - 1. Reference: original model + FullHiddenStatesConnector -> disk - (skipped for models not supported by extract_hidden_states, e.g. GPT-2) - 2. Monitored: hooked model + DMXGPUWorker + ring transport -> ClickHouse - 3. Comparator: reads both, validates row counts + value comparison +Runs the monitored model (hooked + ring transport -> ClickHouse) and +validates schema + per-hook row counts (plus value comparison against the +reference when the model is supported by extract_hidden_states). The +orchestration (vllm_monitored_runner -> vllm_rowcnt_comparator) now lives in +:mod:`tests.e2e_matrix`; this wrapper drives the matrix's vLLM ``row_count`` +cell and asserts on its checks. -Environment variables: - E2E_MODEL "gpt2" (default) or "qwen3" - E2E_NUM_PROMPTS Number of prompts (default 8) - E2E_MAX_NEW_TOKENS Tokens to generate per prompt (default 20) - E2E_ENFORCE_EAGER "1" to disable torch.compile + CUDA graphs (default "0") - E2E_RING_PAYLOAD_MB Ring payload size in MB (default 4096) - E2E_RING_PINNED_MB Pinned staging size in MB (default 4096) - E2E_HOOK_SELECTION Public hook selection preset (default "vllm-full"); - translated to DMX_HOOK_SELECTION for subprocesses - E2E_COMPARE_LAYERS "all" or comma-separated layer IDs for value comparison. - Requires model supported by extract_hidden_states. - GPT-2 not supported -- value comparison skipped with warning. - E2E_TOLERANCE Max abs diff tolerance (default "0.01") - DMX_DB_HOST ClickHouse host (default "localhost") - DMX_DB_PORT ClickHouse port (default 9000) - -Requires: - - ClickHouse running on DMX_DB_HOST:DMX_DB_PORT - - VLLM_DISABLE_COMPILE_CACHE=1 (set automatically) - - LD_PRELOAD for libstdc++ if needed (caller's responsibility) +The test name is preserved because ``tests/tools/verify_vllm.sh`` invokes +this file and threads the model / ring-size / tolerance env vars the matrix +wrapper reads (E2E_MODEL, E2E_ENFORCE_EAGER, E2E_RING_PAYLOAD_MB, +E2E_RING_PINNED_MB, E2E_TOLERANCE, DMX_DB_HOST, DMX_DB_PORT). Usage: python -m pytest tests/test_vllm_rowcnt.py -q -s - E2E_MODEL=qwen3 E2E_COMPARE_LAYERS=all python -m pytest tests/test_vllm_rowcnt.py -q -s + E2E_MODEL=qwen3 python -m pytest tests/test_vllm_rowcnt.py -q -s """ - -import json -import os -import shutil -import subprocess -import sys -import tempfile +from __future__ import annotations import pytest -import torch + +from tests._requirements import require_cuda, require_clickhouse, require_vllm +from tests.e2e_matrix import matrix_argv_from_env, run_single pytestmark = [ pytest.mark.gpu, @@ -48,73 +30,16 @@ pytest.mark.e2e, ] -_MODEL_ALIASES = { - "gpt2": "gpt2", - "qwen2_moe": "Qwen/Qwen1.5-MoE-A2.7B", - "qwen3": "Qwen/Qwen3-0.6B", -} - -@pytest.mark.skipif( - not torch.backends.cuda.is_built(), reason="CUDA not built") -def test_vllm_rowcnt(subtests): - """vLLM row-count validation: monitored run + row-count check.""" - - model_key = os.environ.get("E2E_MODEL", "gpt2") - model_id = _MODEL_ALIASES.get(model_key, model_key) - - # Translate the public E2E_HOOK_SELECTION input into the internal - # DMX_HOOK_SELECTION runtime contract that the runner actually reads. - sub_env = dict(os.environ) - sub_env["DMX_HOOK_SELECTION"] = os.environ.get( - "E2E_HOOK_SELECTION", "vllm-full") - - run_dir = tempfile.mkdtemp(prefix="vllm_rowcnt_") - ref_dir = os.path.join(run_dir, "ref") - mon_dir = os.path.join(run_dir, "mon") - result_file = os.path.join(run_dir, "result.json") - project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - - # ref_dir still needed by comparator (with skipped marker) - os.makedirs(ref_dir, exist_ok=True) - with open(os.path.join(ref_dir, "meta.json"), "w") as f: - json.dump({"skipped": True}, f) - - print(f"\n{'=' * 60}") - print(f" vLLM row-count test") - print(f" model={model_id}") - print(f"{'=' * 60}") - - try: - # Step 1: Monitored run (hooked model + ring transport) - print("\n [1/2] Monitored run (hooked model + DMXGPUWorker)...", flush=True) - r2 = subprocess.run( - [sys.executable, "-m", "tests.vllm_monitored_runner", - "--output-dir", mon_dir], - env=sub_env, capture_output=True, text=True, cwd=project_root, - ) - if r2.returncode != 0: - pytest.fail(f"Monitored runner failed:\n{r2.stderr[-2000:]}") - - # Step 2: Comparator (CPU only, row-count validation) - print(" [2/2] Checking row counts...", flush=True) - r3 = subprocess.run( - [sys.executable, "-m", "tests.vllm_rowcnt_comparator", - "--ref-dir", ref_dir, - "--mon-dir", mon_dir, - "--result-file", result_file], - env=sub_env, capture_output=True, text=True, cwd=project_root, - ) - if r3.returncode != 0: - pytest.fail(f"Comparator failed:\n{r3.stderr[-2000:]}") - - # Read results - with open(result_file) as f: - results = json.load(f) - - # Report via subtests - for test in results["tests"]: - with subtests.test(test["name"]): - assert test["passed"], test.get("detail", "") - finally: - shutil.rmtree(run_dir, ignore_errors=True) +@require_cuda() +@require_vllm() +@require_clickhouse() +def test_vllm_rowcnt(subtests) -> None: + """vLLM row-count validation: monitored run + schema / row-count checks.""" + cr = run_single(matrix_argv_from_env("vllm", "row_count")) + if cr.error: + pytest.fail(f"matrix cell errored: {cr.error}") + assert cr.checks, "matrix produced no checks" + for chk in cr.checks: + with subtests.test(chk.name): + assert chk.passed, chk.detail diff --git a/tests/vllm_ref_runner.py b/tests/vllm_ref_runner.py index 1e244ffff..0298bcb08 100644 --- a/tests/vllm_ref_runner.py +++ b/tests/vllm_ref_runner.py @@ -51,6 +51,7 @@ def main(): os.environ.get("E2E_MAX_NUM_BATCHED_TOKENS", "512")), enforce_eager=enforce_eager, gpu_memory_utilization=float(os.environ.get("E2E_GPU_MEM_UTIL", "0.5")), + max_num_seqs=int(os.environ.get("E2E_MAX_NUM_SEQS", "32")), tensor_parallel_size=tp_size, ) if enable_ep: