diff --git a/docs/examples/benchmarkable_tests.rst b/docs/examples/benchmarkable_tests.rst new file mode 100644 index 0000000000..825f49ce7b --- /dev/null +++ b/docs/examples/benchmarkable_tests.rst @@ -0,0 +1,112 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Benchmarkable Tests +=================== + +A benchmarkable test is an ordinary pytest test that returns a ``Case`` instead of asserting, so +one definition of setup, evaluation, reference and verification serves both correctness testing +and benchmarking. Running pytest normally checks correctness; adding ``--nvte-benchmark`` times +the same code instead. + +Writing the test +---------------- + +Build a ``Case`` from four callables and return it: + +.. code-block:: python + + from transformer_engine.common.testing import Case, benchmark + + def test_something(shape, dtype): + def setup(): + return make_inputs(shape, dtype) # deterministic + + def evaluate(state): + return te_implementation(state) # the Transformer Engine path + + def reference(state): + return naive_implementation(state) # what it should agree with + + def verify(actual, expected): + torch.testing.assert_close(actual, expected, **dtype_tols(dtype)) + + return Case(setup=setup, evaluate=evaluate, reference=reference, verify=verify) + +``setup`` must be deterministic, because benchmark mode calls it again for each timed variant. +``verify`` is required whenever ``reference`` is set; there is no default comparator, so build one +on ``tests/pytorch/utils.py::dtype_tols`` or ``tests/jax/utils.py::assert_allclose``. Raise +``CaseSkip`` from ``setup`` when a backend or architecture is unavailable and the test is skipped. + +Optional fields: ``reset(state)`` runs between timed samples for cases that mutate their state, +``time_reference=False`` records only the Transformer Engine path, and ``bytes_moved`` / ``flops`` +add ``bandwidth_GBps`` and ``tflops`` to the recorded numbers. + +Marking it for benchmarking +--------------------------- + +``@benchmark(argnames, values)`` gives an axis the values it should take when benchmarking. It +does not create an axis: the values replace those of an existing ``pytest.mark.parametrize`` with +the same argnames, so correctness parametrization is untouched. Coupled argnames are written +exactly as parametrized (``"m,n,k"``). + +Abridged from ``tests/pytorch/test_fused_rope.py``: + +.. code-block:: python + + @benchmark("dtype", [torch.bfloat16]) + @benchmark("seq_length", [8192]) + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) + @pytest.mark.parametrize("seq_length", [2048, 4096]) + def test_fused_rope(dtype, seq_length): + ... + return Case(setup=setup, evaluate=evaluate, reference=reference, verify=verify) + +An axis you do not declare keeps its full correctness values, so declare enough of them to keep +the benchmark matrix small -- a single benchmark shape usually means pinning most axes to one +value each. + +``@benchmark`` also applies to a class, where it covers every test method: + +.. code-block:: python + + @benchmark("b,s_q,s_kv,h", [(8, 2048, 2048, 16)]) + @pytest.mark.parametrize("b, s_q, s_kv, h", [...]) + class TestSoftmaxPrimitives: + @staticmethod + def test_forward(b, s_q, s_kv, h, dtype): + ... + return Case(setup=setup, evaluate=evaluate, reference=reference, verify=verify) + + @staticmethod + @benchmark.skip(reason="returns no Case") + def test_backward(b, s_q, s_kv, h, dtype): + ... + +Use ``@benchmark.skip`` or ``@benchmark.skipif(condition)`` for a test that returns a ``Case`` but +should not be benchmarked -- a correctness-only test written in this style, or one whose benchmark +you are temporarily disabling. + +Running benchmarks +------------------ + +.. code-block:: shell + + python3 -m pytest tests/pytorch/test_fused_rope.py --nvte-benchmark \ + --nvte-benchmark-report-dir /tmp/te-bench + +``--nvte-benchmark`` selects benchmark mode and deselects everything else. Each point is checked +for correctness once before it is timed, so a benchmark run also verifies the shapes it measures. + +Options, with defaults: + +* ``--nvte-benchmark-iterations`` (20) -- minimum timed samples per variant. +* ``--nvte-benchmark-warmup`` (5) -- untimed calls before sampling. +* ``--nvte-benchmark-inner-iterations`` (1) -- calls per timed sample. Raise it for kernels short + enough that host launch latency dominates. +* ``--nvte-benchmark-min-run-time`` (0.0) -- keep sampling until this many seconds have elapsed. +* ``--nvte-benchmark-no-reference`` (off) -- skip timing the reference variant. +* ``--nvte-benchmark-report-dir`` (unset) -- where to write the JSON, JSONL and CSV reports. + Without it, the collected numbers are discarded with a warning. diff --git a/docs/index.rst b/docs/index.rst index fcd15a7a11..aeac2e8282 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -59,6 +59,7 @@ Transformer Engine documentation examples/te_mixtral/tutorial_accelerate_hf_mixtral_with_te.ipynb examples/onnx/onnx_export.ipynb examples/te_jax_integration.rst + examples/benchmarkable_tests.rst examples/op_fuser/op_fuser.rst examples/gemm_profiling/gemm_profiling.rst diff --git a/pyproject.toml b/pyproject.toml index 2c9f224c14..1ecedfa2d1 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,3 +7,8 @@ requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "nin # Use legacy backend to import local packages in setup.py build-backend = "setuptools.build_meta:__legacy__" + +[tool.pytest.ini_options] +# Benchmarkable tests return a Case that only the nvte-benchmark plugin runs; without +# it pytest discards the Case and the test passes having verified nothing. +filterwarnings = ["error::pytest.PytestReturnNotNoneWarning"] diff --git a/setup.py b/setup.py index 2a8a9d7688..02af4736de 100644 --- a/setup.py +++ b/setup.py @@ -124,7 +124,7 @@ def setup_requirements() -> Tuple[List[str], List[str]]: "importlib-metadata>=1.0", "packaging", ] - test_reqs: List[str] = ["pytest>=8.2.1"] + test_reqs: List[str] = ["pytest>=8.2.1", "cuda-python>=12.0"] # Framework-specific requirements if not bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))): @@ -403,6 +403,11 @@ def git_check_submodules() -> None: ], ), extras_require=extras_require, + entry_points={ + "pytest11": [ + "nvte-benchmark = transformer_engine.common.testing.plugin", + ], + }, description="Transformer acceleration library", long_description=long_description, long_description_content_type="text/x-rst", diff --git a/tests/jax/pytest.ini b/tests/jax/pytest.ini index 490671a631..73cb71088e 100644 --- a/tests/jax/pytest.ini +++ b/tests/jax/pytest.ini @@ -27,3 +27,6 @@ filterwarnings= ignore:Scan loop is disabled for fused ring attention.*:UserWarning ignore:jax.extend.ffi.register_ffi_target is deprecated ignore:jax.extend.ffi.ffi_lowering is deprecated + # Benchmarkable tests return a Case that only the benchmarkable plugin runs; without + # it pytest discards the Case and the test passes having verified nothing. + error::pytest.PytestReturnNotNoneWarning diff --git a/tests/jax/test_softmax.py b/tests/jax/test_softmax.py index 7af9613538..d61b7b476d 100644 --- a/tests/jax/test_softmax.py +++ b/tests/jax/test_softmax.py @@ -16,6 +16,7 @@ from utils import assert_allclose +from transformer_engine.common.testing import Case, benchmark from transformer_engine.jax.cpp_extensions import is_softmax_kernel_available from transformer_engine.jax.cpp_extensions.attention import AttnSoftmaxType from transformer_engine.jax.softmax import SoftmaxFusionType, softmax @@ -98,15 +99,6 @@ def _setup_inputs(self): case _: raise ValueError(f"Unknown {self.softmax_fusion_type=}") - def test_forward(self): - """ - Test transformer_engine.jax.softmax.softmax fwd rule - """ - self._setup_inputs() - primitive_out = softmax(self.logits, self.mask, self.scale_factor, self.softmax_fusion_type) - reference_out = __class__.reference_softmax(self.logits, self.mask, self.scale_factor) - assert_allclose(primitive_out, reference_out, dtype=self.dtype) - def test_backward(self): """ Test transformer_engine.jax.softmax.softmax bwd rule @@ -148,10 +140,6 @@ class SoftmaxPrimitivesRunner(SoftmaxRunner): Jax Softmax Primitives runner """ - @catch_unsupported - def test_forward(self): - return super().test_forward() - @catch_unsupported def test_backward(self): return super().test_backward() @@ -187,6 +175,8 @@ def test_forward(self): # Run softmax primitives test +# The pinned shape must be one the fused kernel supports, or the benchmark times a raise. +@benchmark("b,s_q,s_kv,h", [(8, 2048, 2048, 16)]) @pytest.mark.parametrize( "b, s_q, s_kv, h", [ @@ -222,9 +212,36 @@ def test_forward(b, s_q, s_kv, h, scale_factor, softmax_fusion_type, dtype): Test forward with parameterized configs """ runner = SoftmaxPrimitivesRunner(b, s_q, s_kv, h, scale_factor, softmax_fusion_type, dtype) - runner.test_forward() + # Resolved here, outside the timed callables; depends only on construction-time fields. + supported = runner._is_support() + + def setup(): + runner._setup_inputs() + return runner + + def evaluate(state): + # Unsupported configs must raise from the primitive rather than compute a result. + if not supported: + with pytest.raises(AssertionError): + softmax(state.logits, state.mask, state.scale_factor, state.softmax_fusion_type) + return None + return softmax(state.logits, state.mask, state.scale_factor, state.softmax_fusion_type) + + def reference(state): + if not supported: + return None + return state.reference_softmax(state.logits, state.mask, state.scale_factor) + + def verify(actual, expected): + # Unsupported configs are checked by the expected raise in evaluate(). + if not supported: + return + assert_allclose(actual, expected, dtype=dtype) + + return Case(setup=setup, evaluate=evaluate, reference=reference, verify=verify) @staticmethod + @benchmark.skip(reason="returns no Case; keeps the class's axes out of collection") def test_backward(b, s_q, s_kv, h, scale_factor, softmax_fusion_type, dtype): """ Test forward with parameterized configs diff --git a/tests/pytorch/test_fused_rope.py b/tests/pytorch/test_fused_rope.py index 5cc4fc0f8a..951ee8e463 100644 --- a/tests/pytorch/test_fused_rope.py +++ b/tests/pytorch/test_fused_rope.py @@ -5,6 +5,7 @@ import math import torch import pytest +from transformer_engine.common.testing import Case, benchmark from transformer_engine.pytorch.attention.rope import ( RotaryPositionEmbedding, apply_rotary_pos_emb, @@ -29,6 +30,16 @@ def _non_overlapping_grad(output: Union[List[torch.Tensor], torch.Tensor]) -> to return torch.sum(output * t) +# start_positions must stay False here: True would hit the margin==0 skip below and +# drop half of the inherited margin values. +@benchmark("start_positions", [False]) +@benchmark("dtype", [torch.bfloat16]) +@benchmark("seq_length", [8192]) +@benchmark("hidden_size", [128]) +@benchmark("transpose", [None]) +@benchmark("tensor_format", ["sbhd"]) +@benchmark("cp_size", [1, 2]) +@benchmark("interleaved", [True]) @pytest.mark.parametrize("start_positions", [True, False]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) @pytest.mark.parametrize("seq_length", [2048, 4096]) @@ -52,7 +63,7 @@ def test_fused_rope( cp_size: int, interleaved: bool, start_positions: bool, -) -> None: +): if margin == 0 and start_positions == True: # This makes sure that the `start_positions` offsets being applied # are with the maximum length of the rope embeddings. @@ -60,67 +71,97 @@ def test_fused_rope( device = torch.device("cuda:0") batch_size, head_num = 2, 64 - t = torch.rand( - (seq_length - margin, batch_size, head_num, hidden_size), - dtype=dtype, - device=device, - ) - # Get arbitrary offsets to be used with RoPE for all the sequences - start_positions = ( - torch.randint(0, margin, (batch_size,), dtype=torch.int32, device=device) - if start_positions - else None - ) - - if tensor_format == "bshd": - t = t.transpose(0, 1).contiguous() - if transpose: - t = t.transpose(*transpose).contiguous().transpose(*transpose) - t.requires_grad = True + def setup(): + t = torch.rand( + (seq_length - margin, batch_size, head_num, hidden_size), + dtype=dtype, + device=device, + ) - rotary_pos_emb = RotaryPositionEmbedding(hidden_size, rotary_percent, interleaved=interleaved) - emb = rotary_pos_emb(seq_length * cp_size) - assert emb.is_contiguous() + # Get arbitrary offsets to be used with RoPE for all the sequences + sp = ( + torch.randint(0, margin, (batch_size,), dtype=torch.int32, device=device) + if start_positions + else None + ) - for cp_rank in range(cp_size): - # unfused - # The fused kernel computes in float32 internally, so we force the unfused func to use float32 - # for more accurate comparison - output_unfused = apply_rotary_pos_emb( - t.float(), - emb, - tensor_format=tensor_format, - start_positions=start_positions, - interleaved=interleaved, - fused=False, - cp_size=cp_size, - cp_rank=cp_rank, - ).to(dtype) - loss_unfused = loss_func(output_unfused) - loss_unfused.backward() - grad_unfused = t.grad.detach().clone() - t.grad = None + if tensor_format == "bshd": + t = t.transpose(0, 1).contiguous() + if transpose: + t = t.transpose(*transpose).contiguous().transpose(*transpose) + t.requires_grad = True - # fused - output_fused = apply_rotary_pos_emb( - t, - emb, - tensor_format=tensor_format, - start_positions=start_positions, - interleaved=interleaved, - fused=True, - cp_size=cp_size, - cp_rank=cp_rank, + rotary_pos_emb = RotaryPositionEmbedding( + hidden_size, rotary_percent, interleaved=interleaved ) - loss_fused = loss_func(output_fused) - loss_fused.backward() - grad_fused = t.grad.detach().clone() - t.grad = None + emb = rotary_pos_emb(seq_length * cp_size) + assert emb.is_contiguous() + return {"t": t, "start_positions": sp, "emb": emb} + + def reset(state): + # A no-op here (the cp_rank loop already clears `t.grad`); declared to + # exercise the runner's reset contract, whose presence pins this case to + # batchable=False / inner_iterations=1. + state["t"].grad = None + + def evaluate(state): + t, emb, sp = state["t"], state["emb"], state["start_positions"] + outputs = [] + for cp_rank in range(cp_size): + output_fused = apply_rotary_pos_emb( + t, + emb, + tensor_format=tensor_format, + start_positions=sp, + interleaved=interleaved, + fused=True, + cp_size=cp_size, + cp_rank=cp_rank, + ) + loss_fused = loss_func(output_fused) + loss_fused.backward() + grad_fused = t.grad.detach() + # detach() without clone(): the `t.grad = None` below is what keeps the + # capture alias-free -- the next backward() installs a fresh buffer + # instead of accumulating into what `grad_fused` aliases. Removing or + # reordering it silently corrupts every earlier capture. + t.grad = None + outputs.append((output_fused, grad_fused)) + return tuple(outputs) + + def reference(state): + t, emb, sp = state["t"], state["emb"], state["start_positions"] + outputs = [] + for cp_rank in range(cp_size): + # The fused kernel computes in float32 internally, so we force the unfused func to use float32 + # for more accurate comparison + output_unfused = apply_rotary_pos_emb( + t.float(), + emb, + tensor_format=tensor_format, + start_positions=sp, + interleaved=interleaved, + fused=False, + cp_size=cp_size, + cp_rank=cp_rank, + ).to(dtype) + loss_unfused = loss_func(output_unfused) + loss_unfused.backward() + grad_unfused = t.grad.detach() + # detach() without clone(); as in evaluate(), the `t.grad = None` below + # is what keeps the capture alias-free. + t.grad = None + outputs.append((output_unfused, grad_unfused)) + return tuple(outputs) + + def verify(actual, expected): + for (output_fused, grad_fused), (output_unfused, grad_unfused) in zip(actual, expected): + torch.testing.assert_close(output_fused, output_unfused) + torch.testing.assert_close(grad_fused, grad_unfused) + assert output_fused.is_contiguous() - torch.testing.assert_close(output_fused, output_unfused) - torch.testing.assert_close(grad_fused, grad_unfused) - assert output_fused.is_contiguous() + return Case(setup=setup, evaluate=evaluate, reference=reference, verify=verify, reset=reset) @pytest.mark.parametrize("margin", [10]) diff --git a/transformer_engine/common/testing/__init__.py b/transformer_engine/common/testing/__init__.py new file mode 100644 index 0000000000..985472819e --- /dev/null +++ b/transformer_engine/common/testing/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Unified correctness and benchmarking test utilities.""" + +from .case import Case, CaseSkip + +__all__ = ["Case", "CaseSkip", "benchmark"] # pylint: disable=undefined-all-variable + + +def __getattr__(name): + """Resolve ``benchmark`` lazily so importing this package does not require pytest.""" + if name == "benchmark": + from .decorator import benchmark as _benchmark + + return _benchmark + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/transformer_engine/common/testing/artifacts.py b/transformer_engine/common/testing/artifacts.py new file mode 100644 index 0000000000..5e0ae11048 --- /dev/null +++ b/transformer_engine/common/testing/artifacts.py @@ -0,0 +1,358 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Artifact helpers for benchmarkable runs.""" + +from __future__ import annotations + +import csv +from datetime import datetime, timezone +import json +import os +from pathlib import Path +import platform +import socket +import subprocess +import sys +from typing import Any + +from .case import axis_value +from .device import build_architectures, device_metadata + + +REPORT_SCHEMA_VERSION = "benchmark_report/v1" + +# .../transformer_engine/common/testing/artifacts.py -> testing -> common +# -> transformer_engine -> repo root +REPO_ROOT = Path(__file__).resolve().parents[3] + + +def write_run_artifacts( + output_dir: Path, + records: list[dict[str, Any]], + command: list[str], + selection: dict[str, Any], + sharding: dict[str, Any] | None = None, + report_name: str = "benchmark_report.json", + records_name: str = "benchmark_records.jsonl", + summary_name: str = "benchmark_summary.csv", +) -> dict[str, Path]: + """Write JSON, JSONL and CSV artifacts for one benchmark run.""" + output_dir.mkdir(parents=True, exist_ok=True) + report = build_report(records, command, selection, sharding=sharding) + + report_path = output_dir / report_name + records_path = output_dir / records_name + summary_path = output_dir / summary_name + + _write_json(report_path, report) + _write_jsonl(records_path, records) + _write_summary_csv(summary_path, records) + return {"report": report_path, "records": records_path, "summary": summary_path} + + +def build_report( + records: list[dict[str, Any]], + command: list[str], + selection: dict[str, Any], + sharding: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Build a self-contained machine-readable benchmark report.""" + status_counts: dict[str, int] = {} + for record in records: + status = str(record.get("status", "unknown")) + status_counts[status] = status_counts.get(status, 0) + 1 + + return { + "schema_version": REPORT_SCHEMA_VERSION, + "created_at_utc": datetime.now(timezone.utc).isoformat(), + "command": command, + "selection": selection, + "environment": collect_environment(command), + "sharding": sharding or {"enabled": False}, + "summary": { + "record_count": len(records), + "status_counts": status_counts, + }, + "records": records, + } + + +def collect_environment(command: list[str] | None = None) -> dict[str, Any]: + """Collect stable environment metadata without copying arbitrary environment variables. + + Framework versions are read from already-imported modules; never import one here. + ``build`` holds the architectures embedded in the loaded ``libtransformer_engine.so``. + """ + return { + "python": { + "version": sys.version, + "executable": sys.executable, + }, + "platform": { + "system": platform.system(), + "release": platform.release(), + "machine": platform.machine(), + "processor": platform.processor(), + "python_implementation": platform.python_implementation(), + }, + "host": { + "hostname": socket.gethostname(), + "cpu_count": os.cpu_count(), + }, + "git": _git_metadata(), + "frameworks": _framework_versions(), + "devices": {"cuda": device_metadata()}, + "build": build_architectures(), + "scheduler": _scheduler_metadata(), + "command": command or [], + } + + +def merge_worker_reports( + output_dir: Path, + worker_report_paths: list[Path], + command: list[str], + selection: dict[str, Any], + sharding: dict[str, Any], +) -> dict[str, Path]: + """Merge worker reports into the standard top-level artifact set. + + A worker report that cannot be read is recorded in the output rather than raising. + """ + merged_records: list[dict[str, Any]] = [] + worker_summaries = [] + unreadable_reports = [] + for report_path in sorted(worker_report_paths): + try: + with report_path.open("r", encoding="utf-8") as handle: + report = json.load(handle) + except (OSError, json.JSONDecodeError) as exc: + unreadable_reports.append({"path": str(report_path), "reason": str(exc)}) + continue + worker_summaries.append( + { + "path": str(report_path), + "summary": report.get("summary", {}), + "selection": report.get("selection", {}), + } + ) + merged_records.extend(report.get("records", [])) + + expected_records = _expected_records_from_worker_summaries(worker_summaries) + sharding = dict(sharding) + sharding["worker_reports"] = worker_summaries + sharding["unreadable_worker_reports"] = unreadable_reports + sharding["merge_validation"] = _validate_merged_records( + merged_records, + expected_records=expected_records, + unreadable_reports=unreadable_reports, + ) + return write_run_artifacts( + output_dir, + merged_records, + command, + selection, + sharding=sharding, + ) + + +def record_key(record: dict[str, Any]) -> str: + """Return a deterministic key for comparing benchmark records. + + ``params`` goes through ``axis_value``. Unlike the record writers, this ``json.dumps`` + has no ``default=str``, so a value it did not normalize raises here instead of becoming + a key that differs every run. + """ + key = { + "case_id": record.get("case_id"), + "framework": record.get("framework"), + "operation": record.get("operation"), + "variant": record.get("variant"), + "params": {name: axis_value(value) for name, value in record.get("params", {}).items()}, + } + return json.dumps(key, sort_keys=True, separators=(",", ":")) + + +def _validate_merged_records( + records: list[dict[str, Any]], + expected_records: list[dict[str, Any]] | None = None, + unreadable_reports: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + counts: dict[str, int] = {} + for record in records: + key = record_key(record) + counts[key] = counts.get(key, 0) + 1 + duplicates = sorted(key for key, count in counts.items() if count > 1) + missing = [] + if expected_records is not None: + expected_keys = sorted({record_key(record) for record in expected_records}) + missing = sorted(key for key in expected_keys if key not in counts) + unreadable = unreadable_reports or [] + return { + "record_count": len(records), + "unique_record_count": len(counts), + "duplicate_keys": duplicates, + "missing_keys": missing, + "unreadable_report_count": len(unreadable), + "valid": not duplicates and not missing and not unreadable, + } + + +def _expected_records_from_worker_summaries( + worker_summaries: list[dict[str, Any]], +) -> list[dict[str, Any]]: + expected_records = [] + for worker in worker_summaries: + selection = worker.get("selection", {}) + include_reference = selection.get("include_reference", True) + for case in selection.get("selected_cases", []): + variants = [] + if include_reference and case.get("has_reference"): + variants.append("reference") + variants.append("evaluation") + for variant in variants: + expected_records.append( + { + "case_id": case.get("case_id"), + "framework": case.get("framework"), + "operation": case.get("operation"), + "variant": variant, + "params": case.get("params", {}), + } + ) + return expected_records + + +def _write_json(path: Path, data: dict[str, Any]) -> None: + with path.open("w", encoding="utf-8") as handle: + # default=str: ``params`` hold pytest parametrize values, which are not always + # JSON-native (a JAX dtype is a bare type object, not a serializable instance). + json.dump(data, handle, indent=2, sort_keys=True, default=str) + handle.write("\n") + + +def _write_jsonl(path: Path, records: list[dict[str, Any]]) -> None: + with path.open("w", encoding="utf-8") as handle: + for record in records: + handle.write(json.dumps(record, sort_keys=True, default=str)) + handle.write("\n") + + +def _write_summary_csv(path: Path, records: list[dict[str, Any]]) -> None: + fields = [ + "status", + "framework", + "case_id", + "variant", + "component", + "operation", + "median_ms", + "mean_ms", + "p95_ms", + "reason", + ] + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + for record in records: + timing = record.get("timing", {}) + writer.writerow( + { + "status": record.get("status"), + "framework": record.get("framework"), + "case_id": record.get("case_id"), + "variant": record.get("variant"), + "component": record.get("component"), + "operation": record.get("operation"), + "median_ms": timing.get("median_ms"), + "mean_ms": timing.get("mean_ms"), + "p95_ms": timing.get("p95_ms"), + "reason": record.get("reason"), + } + ) + + +def _git_metadata() -> dict[str, Any]: + return { + "commit": _run_git(["rev-parse", "HEAD"]), + "branch": _run_git(["rev-parse", "--abbrev-ref", "HEAD"]), + "dirty": bool(_run_git(["status", "--porcelain"])), + } + + +def _run_git(args: list[str]) -> str | None: + try: + result = subprocess.run( + ["git", *args], + check=False, + capture_output=True, + encoding="utf-8", + cwd=REPO_ROOT, + ) + except OSError: + return None + if result.returncode != 0: + return None + return result.stdout.strip() + + +def _framework_versions() -> dict[str, str | None]: + versions: dict[str, str | None] = { + "transformer_engine": _module_version("transformer_engine"), + "torch": _module_version("torch"), + "jax": _module_version("jax"), + "jaxlib": _module_version("jaxlib"), + } + return versions + + +def _module_version(module_name: str) -> str | None: + module = sys.modules.get(module_name) + if module is None: + return None + return getattr(module, "__version__", None) + + +def _scheduler_metadata() -> dict[str, Any]: + names = [ + "CUDA_VISIBLE_DEVICES", + "SLURM_JOB_ID", + "SLURM_JOB_GPUS", + "SLURM_GPUS", + "SLURM_GPUS_ON_NODE", + "SLURM_STEP_GPUS", + ] + metadata = {name: os.environ.get(name) for name in names if os.environ.get(name) is not None} + visible_devices = visible_cuda_devices() + allocated_devices = scheduler_allocated_devices() + metadata.update( + { + "visible_cuda_devices": visible_devices, + "visible_gpu_count": len(visible_devices), + "scheduler_allocated_devices": allocated_devices, + "scheduler_allocated_gpu_count": len(allocated_devices), + } + ) + return metadata + + +def scheduler_allocated_devices() -> list[str]: + """Return the device list the job scheduler allocated to this process.""" + for name in ("SLURM_STEP_GPUS", "SLURM_JOB_GPUS", "CUDA_VISIBLE_DEVICES"): + devices = parse_device_list(os.environ.get(name)) + if devices: + return devices + return [] + + +def visible_cuda_devices() -> list[str]: + """Return the devices named by ``CUDA_VISIBLE_DEVICES``.""" + return parse_device_list(os.environ.get("CUDA_VISIBLE_DEVICES")) + + +def parse_device_list(raw: str | None) -> list[str]: + """Split a comma-separated device list, dropping empty entries.""" + if raw is None: + return [] + return [item.strip() for item in raw.split(",") if item.strip()] diff --git a/transformer_engine/common/testing/case.py b/transformer_engine/common/testing/case.py new file mode 100644 index 0000000000..d38ad9ab58 --- /dev/null +++ b/transformer_engine/common/testing/case.py @@ -0,0 +1,104 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""The Case contract returned by a benchmarkable test, and its axis_value renderer.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +import enum +import os +from typing import Any + +from .declaration import plugin_active + + +def _stable_name(value: Any) -> str | None: + """Return ``value``'s ``__qualname__``, else its ``__name__``, else ``None``.""" + for attribute in ("__qualname__", "__name__"): + name = getattr(value, attribute, None) + if isinstance(name, str): + return name + return None + + +def axis_value(value: Any) -> Any: + """Render one parametrize value canonically for artifact identity.""" + # The enum check must precede the int check: IntEnum is an int. + if isinstance(value, enum.Enum): + return value.name + if isinstance(value, type): + return value.__qualname__ + if value is None or isinstance(value, (bool, int, float, str)): + return value + name = _stable_name(value) + if name is not None: + return name + rendered = str(value) + # An address-bearing rendering is not stable across processes: reject it. + if " at 0x" in rendered: + raise ValueError( + f"Benchmark axis value {rendered!r} has no stable rendering: its repr embeds " + "a memory address, so no two runs would agree on its case_id. Parametrize " + "this axis with an enum or a named function/class instead of a bare object." + ) + return rendered + + +class CaseSkip(Exception): + """Raised by a Case callable when a feature is unavailable; the plugin skips.""" + + +def _require_plugin() -> None: + """Refuse to build a Case inside a pytest session that has no plugin to run it.""" + # pytest sets PYTEST_CURRENT_TEST only while a test is executing, so this fires + # exactly when a Case would be discarded and its test pass having verified nothing. + # It lives here, not in the decorators, because a Case-bearing test need not be + # decorated at all, and a class-level decorator never wraps its methods. + current = os.environ.get("PYTEST_CURRENT_TEST") + if not current or plugin_active(): + return + raise RuntimeError( + f"{current} built a Case, but the nvte-benchmark pytest plugin is not registered," + " so nothing would run it and this test asserts nothing. It autoloads from" + " transformer_engine's entry point, so reinstall transformer_engine (nvte-setup)" + " if pytest is not picking it up, and drop any '-p no:nvte-benchmark'." + ) + + +@dataclass +class Case: + """What a benchmarkable test returns: ``setup`` builds an opaque state object, + ``evaluate`` is the Transformer Engine path, and ``reference`` is the comparison + target in correctness mode and a timed baseline in benchmark mode. + """ + + setup: Callable[[], Any] + evaluate: Callable[[Any], Any] + reference: Callable[[Any], Any] | None = None + verify: Callable[[Any, Any], None] | None = None + reset: Callable[[Any], None] | None = None + batchable: bool = True + time_reference: bool = True + bytes_moved: int | None = None + flops: int | None = None + + # Setting `reset` implies not batchable: ``runner.py`` pins inner_iterations to 1. + + def __post_init__(self) -> None: + """Validate the plugin is present to run this Case, and that ``reference`` and + ``verify`` are set together.""" + _require_plugin() + if self.reference is None and self.verify is not None: + raise ValueError("Case defines verify but no reference to compare against.") + if self.reference is not None and self.verify is None: + raise ValueError( + "Case defines a reference but no verify, and there is no default " + "comparator. Pass a verify built on the framework's tolerance helper " + "(tests/pytorch/utils.py::dtype_tols, tests/jax/utils.py::assert_allclose)." + ) + + def run_verify(self, actual: Any, expected: Any) -> None: + """Compare one evaluate output against one reference output.""" + self.verify(actual, expected) diff --git a/transformer_engine/common/testing/compare.py b/transformer_engine/common/testing/compare.py new file mode 100644 index 0000000000..97b70c067d --- /dev/null +++ b/transformer_engine/common/testing/compare.py @@ -0,0 +1,267 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Compare benchmarkable reports against historical artifacts.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +import sys +from typing import Any + +from .artifacts import record_key + + +DEFAULT_RELATIVE_THRESHOLD = 0.05 +DEFAULT_ABSOLUTE_THRESHOLD_MS = 0.01 + + +def compare_reports( + baseline_report: dict[str, Any], + current_report: dict[str, Any], + relative_threshold: float = DEFAULT_RELATIVE_THRESHOLD, + absolute_threshold_ms: float = DEFAULT_ABSOLUTE_THRESHOLD_MS, +) -> dict[str, Any]: + """Compare completed records in two reports.""" + baseline_records = _completed_records_by_key(baseline_report) + current_records = _completed_records_by_key(current_report) + + regressions = [] + improvements = [] + unchanged = [] + incompatible = [] + hardware_compatible, incompatible_reason = _hardware_compatible(baseline_report, current_report) + + for key, current in sorted(current_records.items()): + baseline = baseline_records.get(key) + if baseline is None: + continue + if not hardware_compatible: + incompatible.append({"key": key, "reason": incompatible_reason}) + continue + + baseline_ms = _median_ms(baseline) + current_ms = _median_ms(current) + delta_ms = current_ms - baseline_ms + relative_delta = delta_ms / baseline_ms if baseline_ms > 0 else 0.0 + threshold = _threshold_for_record(current, relative_threshold, absolute_threshold_ms) + entry = { + "key": key, + "case_id": current.get("case_id"), + "variant": current.get("variant"), + "params": current.get("params", {}), + "baseline_median_ms": baseline_ms, + "current_median_ms": current_ms, + "delta_ms": delta_ms, + "relative_delta": relative_delta, + "relative_threshold": threshold["relative"], + "absolute_threshold_ms": threshold["absolute_ms"], + } + if delta_ms > threshold["absolute_ms"] and relative_delta > threshold["relative"]: + regressions.append(entry) + elif delta_ms < -threshold["absolute_ms"] and -relative_delta > threshold["relative"]: + improvements.append(entry) + else: + unchanged.append(entry) + + missing = sorted(key for key in baseline_records if key not in current_records) + new = sorted(key for key in current_records if key not in baseline_records) + return { + "schema_version": "benchmark_comparison/v1", + "summary": { + "baseline_records": len(baseline_records), + "current_records": len(current_records), + "regressions": len(regressions), + "improvements": len(improvements), + "unchanged": len(unchanged), + "missing": len(missing), + "new": len(new), + "incompatible": len(incompatible), + }, + "regressions": regressions, + "improvements": improvements, + "unchanged": unchanged, + "missing": missing, + "new": new, + "incompatible": incompatible, + } + + +def main() -> int: + """Compare two benchmark_report/v1 JSON files and write a comparison report.""" + parser = argparse.ArgumentParser(description="Compare benchmarkable JSON reports.") + parser.add_argument("--baseline", required=True, help="Historical benchmark_report.json") + parser.add_argument("--current", required=True, help="Current benchmark_report.json") + parser.add_argument("--output", required=True, help="Comparison JSON output path") + parser.add_argument( + "--relative-threshold", + type=float, + default=DEFAULT_RELATIVE_THRESHOLD, + help="Default relative regression threshold (default: %(default)s)", + ) + parser.add_argument( + "--absolute-threshold-ms", + type=float, + default=DEFAULT_ABSOLUTE_THRESHOLD_MS, + help="Default absolute regression threshold in ms (default: %(default)s)", + ) + parser.add_argument( + "--fail-on-regression", + action="store_true", + help="Exit nonzero when any regression is detected.", + ) + parser.add_argument( + "--fail-on-missing", + action="store_true", + help=( + "Exit nonzero when a baseline record has no completed counterpart in the current " + "report. A case that starts erroring disappears from the comparison entirely, so " + "without this flag a broken case looks like a clean run." + ), + ) + parser.add_argument( + "--fail-on-incompatible", + action="store_true", + help=( + "Exit nonzero when any record is hardware-incompatible with its baseline " + "counterpart. Every incompatible record short-circuits past the regression test " + "and 'missing' stays zero because the key is still present, so without this flag " + "a baseline from different hardware -- e.g. an H100 baseline against a B200 run -- " + "compares as clean by default." + ), + ) + args = parser.parse_args() + + baseline = _read_json(Path(args.baseline)) + current = _read_json(Path(args.current)) + comparison = compare_reports( + baseline, + current, + relative_threshold=args.relative_threshold, + absolute_threshold_ms=args.absolute_threshold_ms, + ) + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + with output.open("w", encoding="utf-8") as handle: + json.dump(comparison, handle, indent=2, sort_keys=True) + handle.write("\n") + + if comparison["incompatible"]: + reasons = sorted( + {entry["reason"] for entry in comparison["incompatible"] if entry.get("reason")} + ) + detail = f" ({'; '.join(reasons)})" if reasons else "" + print( + f"Warning: {len(comparison['incompatible'])} record(s) are hardware-incompatible " + f"with their baseline counterpart and were skipped for regression comparison{detail}.", + file=sys.stderr, + ) + + if args.fail_on_regression and comparison["regressions"]: + print( + f"Error: {len(comparison['regressions'])} regression(s) detected.", + file=sys.stderr, + ) + return 2 + if args.fail_on_missing and comparison["missing"]: + print( + f"Error: {len(comparison['missing'])} baseline record(s) missing from the current " + "report.", + file=sys.stderr, + ) + return 3 + if args.fail_on_incompatible and comparison["incompatible"]: + print( + f"Error: {len(comparison['incompatible'])} record(s) are hardware-incompatible " + "with their baseline counterpart.", + file=sys.stderr, + ) + return 4 + return 0 + + +def _read_json(path: Path) -> dict[str, Any]: + with path.open("r", encoding="utf-8") as handle: + return json.load(handle) + + +def _completed_records_by_key(report: dict[str, Any]) -> dict[str, dict[str, Any]]: + records = {} + for record in report.get("records", []): + if record.get("status") == "completed": + records[record_key(record)] = record + return records + + +def _median_ms(record: dict[str, Any]) -> float: + return float(record.get("timing", {}).get("median_ms", 0.0)) + + +def _threshold_for_record( + record: dict[str, Any], + default_relative: float, + default_absolute_ms: float, +) -> dict[str, float]: + override = record.get("regression_threshold") or {} + return { + "relative": float(override.get("relative", default_relative)), + "absolute_ms": float(override.get("absolute_ms", default_absolute_ms)), + } + + +def _hardware_compatible( + baseline_report: dict[str, Any], + current_report: dict[str, Any], +) -> tuple[bool, str | None]: + """Return whether the two reports were recorded on compatible hardware, plus a reason. + + Missing device metadata on either side means compatibility cannot be established, and is + treated as incompatible rather than as a match. + """ + baseline_device = _primary_device_identity(baseline_report) + current_device = _primary_device_identity(current_report) + if baseline_device is None and current_device is None: + return ( + False, + ( + "neither the baseline nor the current report carries device metadata -- " + "hardware compatibility cannot be established" + ), + ) + if baseline_device is None: + return ( + False, + ( + "the baseline report carries no device metadata -- hardware " + "compatibility cannot be established" + ), + ) + if current_device is None: + return ( + False, + ( + "the current report carries no device metadata -- hardware " + "compatibility cannot be established" + ), + ) + if baseline_device != current_device: + return False, "hardware metadata differs" + return True, None + + +def _primary_device_identity(report: dict[str, Any]) -> tuple[Any, ...] | None: + """Return the primary CUDA device's ``(name, compute_capability)``. + + Returns ``None`` when the report carries no CUDA device metadata. + """ + devices = report.get("environment", {}).get("devices", {}).get("cuda", {}).get("devices", []) + if not devices: + return None + device = devices[0] + return (device.get("name"), tuple(device.get("compute_capability", []))) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/transformer_engine/common/testing/declaration.py b/transformer_engine/common/testing/declaration.py new file mode 100644 index 0000000000..eb85006279 --- /dev/null +++ b/transformer_engine/common/testing/declaration.py @@ -0,0 +1,48 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Storage for benchmark axis declarations, kept free of pytest imports.""" + +from __future__ import annotations + +from typing import Any + +DECLARATION_ATTR = "_te_benchmark_axes" + +_PLUGIN_ACTIVE = False +"""Set true by ``plugin.py``'s ``pytest_configure`` once the plugin is registered.""" + + +def set_plugin_active(active: bool) -> None: + """Record whether the benchmarkable pytest plugin is registered in this session.""" + global _PLUGIN_ACTIVE # pylint: disable=global-statement + _PLUGIN_ACTIVE = bool(active) + + +def plugin_active() -> bool: + """Return whether the benchmarkable pytest plugin registered itself this session.""" + return _PLUGIN_ACTIVE + + +def normalize_argnames(argnames: Any) -> str: + """Return a canonical comma-joined key with no spaces, so a declaration and a + ``pytest.mark.parametrize`` mark written with different spacing still match.""" + if isinstance(argnames, str): + names = [name.strip() for name in argnames.split(",") if name.strip()] + else: + names = [str(name).strip() for name in argnames] + return ",".join(names) + + +def record_axis(holder: Any, argnames: Any, values: Any) -> None: + """Attach one benchmark axis declaration to a function, method, or class.""" + declarations = holder.__dict__.get(DECLARATION_ATTR) + if declarations is None: + declarations = {} + setattr(holder, DECLARATION_ATTR, declarations) + declarations[normalize_argnames(argnames)] = list(values) + + +def declared_axes(function: Any) -> dict[str, list]: + """Return the benchmark axis declarations attached to one test function.""" + return dict(getattr(function, DECLARATION_ATTR, None) or {}) diff --git a/transformer_engine/common/testing/decorator.py b/transformer_engine/common/testing/decorator.py new file mode 100644 index 0000000000..1608291b95 --- /dev/null +++ b/transformer_engine/common/testing/decorator.py @@ -0,0 +1,85 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Benchmarkable mode constants, plugin markers, and the ``benchmark`` decorator.""" + +import pytest + +from .declaration import record_axis + +MODE_CORRECTNESS = "correctness" +"""Run the correctness matrix, verified, with nothing timed.""" + +MODE_BENCHMARK = "benchmark" +"""Run only the benchmark matrix, gated by a one-time correctness check per point.""" + +CASE_MARKER = "nvte_case" +"""Case-bearing: the test returns a Case, so the plugin runs it instead of pytest.""" + +BENCHMARK_MARKER = "nvte_benchmark" +"""Benchmark-eligible: the test declared benchmark axes with ``@benchmark(...)``.""" + +SUPPRESS_MARKER = "nvte_no_benchmark" +"""Benchmarking suppressed by ``@benchmark.skip`` or a true ``@benchmark.skipif``.""" + +MARKERS = ( + (CASE_MARKER, "test returns a Case for the nvte-benchmark plugin to run."), + (BENCHMARK_MARKER, "Case-bearing test that declared benchmark axes."), + (SUPPRESS_MARKER, "Case-bearing test that is never a benchmark point."), +) + + +def _apply(holder, *names, **kwargs): + """Apply plugin markers by name to a function, a method, or a class.""" + for name in names: + holder = getattr(pytest.mark, name)(**kwargs)(holder) + return holder + + +class _Benchmark: + """Implements the ``benchmark`` decorator. + + An instance, not a module-level function carrying attributes, so ``skip`` and + ``skipif`` are real bound methods that ``help()``, ``inspect`` and pylint resolve. + """ + + def __call__(self, argnames, values): + """Mark a test Case-bearing and benchmark-eligible, and declare the values one axis + takes in benchmark mode. + + ``argnames`` accepts the same forms as ``pytest.mark.parametrize``. The values are + substituted into an existing ``parametrize`` mark on the test, its class, or its + module; undeclared axes keep their correctness values. Applies to a function, a + method, or a class. + """ + + def decorator(holder): + record_axis(holder, argnames, values) + return _apply(holder, CASE_MARKER, BENCHMARK_MARKER) + + return decorator + + def skip(self, holder=None, *, reason=""): + """Mark a test Case-bearing but never benchmarked. Usable bare or called with a + ``reason``, on a function, a method, or a class.""" + + def decorator(target): + target = _apply(target, CASE_MARKER) + return _apply(target, SUPPRESS_MARKER, reason=reason) + + return decorator if holder is None else decorator(holder) + + def skipif(self, condition, *, reason=""): + """Mark a test Case-bearing, and suppress benchmarking when ``condition`` is true. + + A false condition declares nothing: skip and skipif only ever subtract benchmarking. + """ + + def decorator(target): + target = _apply(target, CASE_MARKER) + return _apply(target, SUPPRESS_MARKER, reason=reason) if condition else target + + return decorator + + +benchmark = _Benchmark() diff --git a/transformer_engine/common/testing/device.py b/transformer_engine/common/testing/device.py new file mode 100644 index 0000000000..edff4fb296 --- /dev/null +++ b/transformer_engine/common/testing/device.py @@ -0,0 +1,190 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""CUDA device access for benchmarkable tests, via cuda-python rather than a framework.""" + +from __future__ import annotations + +import functools +import re +import shutil +import subprocess +from typing import Any + +# Cubin sections from ``cuobjdump --list-elf``, e.g. ``libtransformer_engine.4.sm_103a.cubin``. +# The ``[af]`` group must stay: dropping the Blackwell ``a``/``f`` suffix yields the wrong arch. +_ELF_ARCH = re.compile(r"\.sm_(\d+)([af]?)\.cubin") + + +def _runtime(): + """Return the cuda-python runtime module, or None when it is unavailable.""" + try: + from cuda.bindings import runtime + except ImportError: + try: # cuda-python < 12.8 + from cuda import cudart as runtime + except ImportError: + return None + return runtime + + +def cuda_available() -> bool: + """Return True when at least one CUDA device is visible.""" + runtime = _runtime() + if runtime is None: + return False + error, count = runtime.cudaGetDeviceCount() + return int(error) == 0 and count > 0 + + +def synchronize(output: Any = None) -> None: + """Block until all device work, including ``output``, has completed. + + ``output`` is duck-typed for host-async frameworks; the device synchronization + covers eager ones. Without cuda-python that synchronization is skipped, so timings + taken around this call are not trustworthy. + """ + _block_until_ready(output) + runtime = _runtime() + if runtime is None: + return + (error,) = runtime.cudaDeviceSynchronize() + if int(error) != 0: + # Must raise: a sticky error (illegal access in the kernel under test) + # surfaces here, and swallowing it records a fast, meaningless timing. + raise RuntimeError(f"cudaDeviceSynchronize failed: {error}") + + +def _block_until_ready(value: Any) -> None: + """Block on JAX-style arrays without importing jax.""" + if value is None: + return + if hasattr(value, "block_until_ready"): + value.block_until_ready() + elif isinstance(value, dict): + for item in value.values(): + _block_until_ready(item) + elif isinstance(value, (list, tuple)): + for item in value: + _block_until_ready(item) + + +def profiler_start() -> bool: + """Start CUDA profiler capture, reporting whether it actually started.""" + runtime = _runtime() + if runtime is None: + return False + error = runtime.cudaProfilerStart() + error = error[0] if isinstance(error, tuple) else error + return int(error) == 0 + + +def profiler_stop() -> bool: + """Stop CUDA profiler capture, reporting whether it stopped cleanly.""" + runtime = _runtime() + if runtime is None: + return False + error = runtime.cudaProfilerStop() + error = error[0] if isinstance(error, tuple) else error + return int(error) == 0 + + +def device_metadata() -> dict[str, Any]: + """Describe the visible CUDA devices without consulting any framework.""" + runtime = _runtime() + if runtime is None: + return {"available": False, "reason": "cuda-python is not installed"} + + error, count = runtime.cudaGetDeviceCount() + if int(error) != 0 or count == 0: + return {"available": False, "device_count": 0} + + error, current = runtime.cudaGetDevice() + if int(error) != 0: + return {"available": False, "reason": f"cudaGetDevice failed: {error}"} + + devices = [] + for index in range(count): + error, props = runtime.cudaGetDeviceProperties(index) + if int(error) != 0: + continue + devices.append( + { + "index": index, + "name": bytes(props.name).decode(errors="replace").rstrip("\x00"), + "total_memory": int(props.totalGlobalMem), + "compute_capability": [int(props.major), int(props.minor)], + "multi_processor_count": int(props.multiProcessorCount), + "uuid": bytes(props.uuid.bytes).hex(), + } + ) + return { + "available": True, + "device_count": count, + "current_device": int(current), + "devices": devices, + } + + +def device_architecture() -> str | None: + """Return the current device's arch as ``sm_``. + + Never raises: any failure yields None, which callers must treat as "unknown" rather + than as a proven architecture mismatch. + """ + runtime = _runtime() + if runtime is None: + return None + error, count = runtime.cudaGetDeviceCount() + if int(error) != 0 or count == 0: + return None + error, current = runtime.cudaGetDevice() + if int(error) != 0: + return None + error, props = runtime.cudaGetDeviceProperties(current) + if int(error) != 0: + return None + return f"sm_{int(props.major) * 10 + int(props.minor)}" + + +@functools.lru_cache(maxsize=None) +def build_architectures() -> dict[str, Any]: + """Return the SASS architectures embedded in ``libtransformer_engine.so``. + + Reads the loaded shared object with ``cuobjdump`` rather than trusting + ``NVTE_CUDA_ARCHS``, so it also holds for libraries built elsewhere. Cached because + ``cuobjdump`` takes seconds. Never raises: any failure yields + ``{"available": False, "reason": ...}``, which callers must treat as "unknown" + rather than as a proven architecture mismatch. + """ + from transformer_engine.common import _get_shared_object_file + + try: + so_path = _get_shared_object_file("core") + except (OSError, RuntimeError) as exc: + return {"available": False, "reason": f"could not locate libtransformer_engine.so: {exc}"} + + exe = shutil.which("cuobjdump") + if exe is None: + return {"available": False, "reason": "cuobjdump not found (CUDA toolkit absent)"} + + try: + result = subprocess.run( + [exe, "--list-elf", str(so_path)], + capture_output=True, + encoding="utf-8", + timeout=120, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + return {"available": False, "reason": f"cuobjdump failed: {exc}"} + + if result.returncode != 0: + stderr = (result.stderr or "").strip() + return {"available": False, "reason": f"cuobjdump exited {result.returncode}: {stderr}"} + + archs = sorted( + {f"sm_{m.group(1)}{m.group(2)}" for m in _ELF_ARCH.finditer(result.stdout)}, + key=lambda s: (int(re.sub(r"\D", "", s)), s), + ) + return {"available": True, "library": str(so_path), "cuda_architectures": archs} diff --git a/transformer_engine/common/testing/plugin.py b/transformer_engine/common/testing/plugin.py new file mode 100644 index 0000000000..207a9e8490 --- /dev/null +++ b/transformer_engine/common/testing/plugin.py @@ -0,0 +1,393 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Pytest plugin driving benchmarkable tests.""" + +from __future__ import annotations + +import inspect +import os +from pathlib import Path +import sys +from typing import Any +import warnings + +import pytest + +from .artifacts import write_run_artifacts +from .case import Case, CaseSkip +from .declaration import declared_axes, normalize_argnames, set_plugin_active +from .decorator import ( + BENCHMARK_MARKER, + CASE_MARKER, + MARKERS, + MODE_BENCHMARK, + MODE_CORRECTNESS, + SUPPRESS_MARKER, +) +from .device import build_architectures, cuda_available, device_architecture +from .runner import _run_benchmark_point, _run_correctness +from .timing import TIMING_METHOD + + +def pytest_addoption(parser): + """Register the benchmarkable command-line options, once per parser.""" + # Keyed to the parser, not to a module-level flag: pytest builds a fresh parser per + # pytest.main() call, and a sticky flag would hide --nvte-benchmark from a second run. + if getattr(parser, "_te_benchmarkable_registered", False): + return + parser._te_benchmarkable_registered = True # pylint: disable=protected-access + + group = parser.getgroup("benchmarkable") + group.addoption( + "--nvte-benchmark", + action="store_true", + default=False, + help=( + "Run the benchmark matrix only, gated by a one-time correctness check at " + "each benchmark point." + ), + ) + group.addoption( + "--nvte-benchmark-report-dir", + default=None, + help="Directory for benchmark_report/v1 artifacts.", + ) + group.addoption("--nvte-benchmark-warmup", type=int, default=5) + group.addoption("--nvte-benchmark-iterations", type=int, default=20) + group.addoption("--nvte-benchmark-inner-iterations", type=int, default=1) + group.addoption("--nvte-benchmark-min-run-time", type=float, default=0.0) + group.addoption( + "--nvte-benchmark-no-reference", + action="store_true", + default=False, + help="Skip timing the reference variant.", + ) + + +def _resolve_mode(config) -> str: + """Return ``MODE_BENCHMARK`` if ``--nvte-benchmark`` was passed, else ``MODE_CORRECTNESS``.""" + return MODE_BENCHMARK if config.getoption("--nvte-benchmark") else MODE_CORRECTNESS + + +def pytest_configure(config): + """Resolve the mode before test modules are imported, and register markers.""" + # Ahead of the double-fire guard, so a second registered copy of this hook still + # sets the flag; Case.__post_init__ reads it to detect a plugin-less session. + set_plugin_active(True) + + if getattr(config, "_benchmarkable_configured", False): + return + config._benchmarkable_configured = True # pylint: disable=protected-access + + for name, description in MARKERS: + config.addinivalue_line("markers", f"{name}: {description}") + + if _resolve_mode(config) == MODE_CORRECTNESS: + return + + if config.getoption("--nvte-benchmark-inner-iterations") < 1: + raise pytest.UsageError("--nvte-benchmark-inner-iterations must be at least 1.") + if config.getoption("--nvte-benchmark-iterations") < 1: + raise pytest.UsageError("--nvte-benchmark-iterations must be at least 1.") + + +def _is_case_bearing(node) -> bool: + """Whether the test returns a Case, so this plugin runs it instead of pytest.""" + # get_closest_marker, not node.keywords: keywords also carry every ancestor node's + # bare name, so an unmarked test under a path named nvte_case would match here. + return node.get_closest_marker(CASE_MARKER) is not None + + +def _is_benchmark_eligible(node) -> bool: + """Whether the test declared benchmark axes and nothing suppressed them.""" + # Suppression is a union over the test, its class and its module, and nothing + # re-enables: skip and skipif only ever subtract benchmarking. + if node.get_closest_marker(BENCHMARK_MARKER) is None: + return False + return next(node.iter_markers(name=SUPPRESS_MARKER), None) is None + + +def _substitutions(config) -> dict: + """Per-session record of which test definitions had a benchmark axis substituted.""" + registry = getattr(config, "_benchmarkable_substitutions", None) + if registry is None: + registry = {} + config._benchmarkable_substitutions = registry # pylint: disable=protected-access + return registry + + +def _definition_key(node) -> tuple: + """Key one test definition by the collector and name that ``pytest_generate_tests`` + and every item it generates share.""" + # Not the function object: a method inherited by two classes is one function but two + # definitions, and one class substituting an axis says nothing about the other. + return (node.parent, getattr(node, "originalname", None) or node.name) + + +# trylast is required: pluggy's LIFO ordering would otherwise run this hook before +# pytest's own -k/-m deselection, and the guard must see the final item list. +@pytest.hookimpl(trylast=True) +def pytest_collection_modifyitems(config, items): + """Narrow the session to benchmark points, then run the post-collection guard.""" + if _resolve_mode(config) == MODE_CORRECTNESS: + return + + substituted = _substitutions(config) + kept, dropped, suppressed, unmatched = [], [], 0, 0 + for item in items: + eligible = _is_benchmark_eligible(item) + if eligible and substituted.get(_definition_key(item), False): + kept.append(item) + continue + dropped.append(item) + if eligible: + unmatched += 1 + elif _is_case_bearing(item): + suppressed += 1 + if dropped: + config.hook.pytest_deselected(items=dropped) + items[:] = kept + + _guard_benchmark_collection(config, kept, suppressed, unmatched) + + +def _guard_benchmark_collection(config, kept, suppressed, unmatched) -> None: + """Fail loudly when a requested benchmark mode has nothing valid to time.""" + if not kept: + raise pytest.UsageError(_zero_points_message(suppressed, unmatched)) + + if config.option.collectonly: + return + + # Without cuda-python, device.synchronize() is a silent no-op: every measurement + # would be host launch time only and the report would still claim a device ran. + if not cuda_available(): + raise pytest.UsageError( + "--nvte-benchmark requires a CUDA device reachable through cuda-python; none is " + "available. Install cuda-python (the 'test' extra) or drop --nvte-benchmark." + ) + + _guard_build_architecture() + + +def _zero_points_message(suppressed, unmatched) -> str: + """Explain a zero-point benchmark session without blaming the wrong thing.""" + if suppressed: + detail = ( + f"all {suppressed} Case-bearing test(s) here carry @benchmark.skip or a true" + " @benchmark.skipif." + ) + elif unmatched: + detail = ( + f"{unmatched} benchmark-eligible test(s) here parametrize none of the axes" + " declared on their class." + ) + else: + detail = ( + "check the -k/-m expression, and that a test under the selected paths declares" + " axes with @benchmark(argnames, values)." + ) + return f"--nvte-benchmark collected zero benchmark points: {detail}" + + +def _guard_build_architecture() -> None: + """Fail loudly when the loaded library's embedded SASS cannot run on this device.""" + # Only a proven mismatch may raise: an unavailable build architecture (no + # cuobjdump, or the library was not located) or an unknown device architecture + # leaves the comparison undecidable and must not block the run. + build = build_architectures() + if not build.get("available"): + return + device_arch = device_architecture() + if device_arch is None: + return + + embedded = build["cuda_architectures"] + if device_arch.rstrip("af") in {arch.rstrip("af") for arch in embedded}: + return + + raise pytest.UsageError( + f"{build['library']} was built for " + f"{', '.join(embedded) or '(no architectures embedded)'}, but this device is " + f"{device_arch}, so no kernel can launch. Rebuild with an NVTE_CUDA_ARCHS " + "matching this device (use 100 for sm_100a/sm_103a)." + ) + + +def benchmark_settings(config) -> dict[str, Any]: + """Return the resolved benchmark settings for this session.""" + report_dir = config.getoption("--nvte-benchmark-report-dir") + return { + "mode": _resolve_mode(config), + "warmup": config.getoption("--nvte-benchmark-warmup"), + "iterations": config.getoption("--nvte-benchmark-iterations"), + "inner_iterations": config.getoption("--nvte-benchmark-inner-iterations"), + "min_run_time": config.getoption("--nvte-benchmark-min-run-time"), + "timing_method": TIMING_METHOD, + "report_dir": Path(report_dir) if report_dir else None, + "no_reference": config.getoption("--nvte-benchmark-no-reference"), + } + + +@pytest.hookimpl(hookwrapper=True) +def pytest_generate_tests(metafunc): + """Swap benchmark values into a test's existing parametrize marks. + + Substituting in place preserves the parametrization's positional structure, so ``pytest.param`` + ids and coupled argnames keep working. Class-level declarations are picked up via + ``metafunc.cls``; function-level entries win. Mutates ``own_markers``, not public pytest API + (pytest is pinned to 8.2.1). + """ + saved = [] + # Eligibility is checked before anything else: a suppressed or undeclared sibling + # method must not inherit its class's declarations, nor be failed by the check below. + eligible = _is_benchmark_eligible(metafunc.definition) + try: + if eligible and _resolve_mode(metafunc.config) == MODE_BENCHMARK: + own = declared_axes(metafunc.function) + declarations = dict(own) + if metafunc.cls is not None: + for key, values in declared_axes(metafunc.cls).items(): + declarations.setdefault(key, values) + node = metafunc.definition + while node is not None and declarations: + markers = getattr(node, "own_markers", None) + if markers: + for index, mark in enumerate(markers): + if mark.name != "parametrize": + continue + key = normalize_argnames(mark.args[0]) + if key in declarations: + saved.append((markers, index, mark)) + markers[index] = pytest.mark.parametrize( + mark.args[0], declarations[key] + ).mark + node = getattr(node, "parent", None) + _check_own_declarations_matched(metafunc, own, saved) + _substitutions(metafunc.config)[_definition_key(metafunc.definition)] = bool(saved) + yield + finally: + # Class- and module-level marks are shared, so a substituted mark left in place + # corrupts every sibling test for the rest of the session. The finally must + # therefore span the validation calls and the yield, not just the swap loop. + for markers, index, original in saved: + markers[index] = original + + +def _check_own_declarations_matched(metafunc, own, saved) -> None: + """Fail loudly when an axis declared on this test itself matches no parametrize mark: + a typo would otherwise silently benchmark the correctness values. An axis inherited from + a class is best-effort, since a sibling method need not parametrize it.""" + matched = {normalize_argnames(mark.args[0]) for _, _, mark in saved} + missing = sorted(set(own) - matched) + if not missing: + return + raise pytest.UsageError( + f"{metafunc.definition.nodeid}: @benchmark declared axes {missing}, which match no " + "pytest.mark.parametrize on this test, its class, or its module. Check the " + 'spelling, and declare a coupled group exactly as parametrized ("m,n,k").' + ) + + +@pytest.hookimpl(tryfirst=True) +def pytest_pyfunc_call(pyfuncitem): + """Run a Case-bearing test's Case according to the active mode.""" + if not _is_case_bearing(pyfuncitem): + return None + # An async test cannot return a Case, and leaving it to pytest keeps pytest's own + # warn-and-skip for it byte-identical. + if inspect.iscoroutinefunction(pyfuncitem.obj) or inspect.isasyncgenfunction(pyfuncitem.obj): + return None + + argnames = pyfuncitem._fixtureinfo.argnames # pylint: disable=protected-access + kwargs = {name: pyfuncitem.funcargs[name] for name in argnames} + case = pyfuncitem.obj(**kwargs) + if not isinstance(case, Case): + return _dispose_of_non_case(pyfuncitem, case) + + # CaseSkip from setup() is a coverage skip (unavailable backend or arch), not a + # failure. + if _resolve_mode(pyfuncitem.config) == MODE_CORRECTNESS: + try: + _run_correctness(case) + except CaseSkip as exc: + pytest.skip(str(exc)) + return True + + settings = benchmark_settings(pyfuncitem.config) + try: + records = _run_benchmark_point(case, settings, pyfuncitem) + except CaseSkip as exc: + pytest.skip(str(exc)) + store = getattr(pyfuncitem.config, "_benchmarkable_records", None) + if store is None: + store = [] + pyfuncitem.config._benchmarkable_records = store # pylint: disable=protected-access + store.extend(records) + return True + + +def _dispose_of_non_case(pyfuncitem, result): + """Handle a Case-bearing test that returned no Case, exactly as pytest would.""" + if declared_axes(pyfuncitem.function): + raise TypeError( + f"{pyfuncitem.nodeid}: @benchmark declares axes on this test, so it must return " + f"a Case; got {type(result).__name__}. A test that is Case-bearing but never " + "benchmarked is written @benchmark.skip." + ) + if _resolve_mode(pyfuncitem.config) == MODE_BENCHMARK: + pytest.skip("returns no Case, so it is not a benchmark point.") + # Mirrors _pytest/python.py::pytest_pyfunc_call, which owns this test in every other + # respect: an inherited declaration is a blanket statement, not a per-method claim. + if result is not None: + warnings.warn( + pytest.PytestReturnNotNoneWarning( + f"Expected None, but {pyfuncitem.nodeid} returned {result!r}, which will be " + "an error in a future version of pytest. Did you mean to use `assert` " + "instead of `return`?" + ) + ) + return True + + +def pytest_sessionfinish(session, exitstatus): # pylint: disable=unused-argument + """Write benchmark artifacts once the session completes. Guarded against a second + registered copy of the plugin, since this hook is not ``firstresult``.""" + config = session.config + if getattr(config, "_benchmarkable_report_written", False): + return + config._benchmarkable_report_written = True # pylint: disable=protected-access + records = getattr(config, "_benchmarkable_records", None) + if not records: + return + settings = benchmark_settings(config) + if settings["report_dir"] is None: + print( + f"\nWarning: {len(records)} benchmark record(s) were collected but " + "--nvte-benchmark-report-dir was not set, so they were discarded. Pass " + "--nvte-benchmark-report-dir to write a benchmark_report/v1 artifact.", + file=sys.stderr, + ) + return + + selection = { + "mode": settings["mode"], + "warmup": settings["warmup"], + "iterations": settings["iterations"], + "inner_iterations": settings["inner_iterations"], + "min_run_time": settings["min_run_time"], + "timing_method": settings["timing_method"], + "include_reference": not settings["no_reference"], + "args": list(config.invocation_params.args), + } + # Only argv[0]'s basename is kept: the full launcher path leaks the invoking user's + # home directory into every persisted report, and selection["args"] has the rest. + command = [os.path.basename(sys.argv[0])] + list(sys.argv[1:]) + paths = write_run_artifacts( + settings["report_dir"], + records, + command, + selection, + ) + print(f"\nWrote benchmark report: {paths['report']}") diff --git a/transformer_engine/common/testing/runner.py b/transformer_engine/common/testing/runner.py new file mode 100644 index 0000000000..7b96c1248b --- /dev/null +++ b/transformer_engine/common/testing/runner.py @@ -0,0 +1,166 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Case-execution engine: drives a Case's setup/evaluate/reference/verify cycle.""" + +from __future__ import annotations + +import os +import time + +from .case import Case, axis_value +from .device import synchronize +from .timing import TIMING_METHOD, WallClockSampler, timing_stats + + +def _framework_for(pyfuncitem) -> str: + """Infer the framework label for the report from the test's location under tests/. + + This is a reporting label only; nothing in the execution path branches on it. + """ + path = str(pyfuncitem.path) + if f"{os.sep}jax{os.sep}" in path or path.endswith(f"{os.sep}jax"): + return "jax" + return "pytorch" + + +def _run_correctness(case: Case) -> None: + """Run one setup/evaluate/reference/verify cycle with no timing.""" + state = case.setup() + actual = case.evaluate(state) + if case.reference is None: + return + # ``synchronize(actual)`` must precede ``reset(state)``: ``actual`` may alias + # ``state``, so an eager ``reset`` would race with in-flight ``evaluate`` work. + # ``reset`` must precede ``reference`` so it sees unmutated state. + synchronize(actual) + if case.reset is not None: + case.reset(state) + expected = case.reference(state) + synchronize(expected) + case.run_verify(actual, expected) + + +def _run_benchmark_point(case, settings, pyfuncitem): + """Gate once on correctness, then time evaluate and optionally reference.""" + state = case.setup() + + precondition_verified = False + if case.reference is not None: + actual = case.evaluate(state) + # Same ordering rule as ``_run_correctness``: synchronize before reset. + synchronize(actual) + if case.reset is not None: + case.reset(state) + expected = case.reference(state) + synchronize(expected) + case.run_verify(actual, expected) + precondition_verified = True + + variants = [("evaluation", case.evaluate)] + if case.reference is not None and case.time_reference and not settings["no_reference"]: + variants.insert(0, ("reference", case.reference)) + + # A case needing reset between calls cannot be batched: inner iterations are + # submitted back to back with no chance to reset, so timings would average over + # drifting state. + batchable = case.batchable and case.reset is None + inner = settings["inner_iterations"] if batchable else 1 + records = [] + for name, function in variants: + records.append( + _time_variant( + case, settings, pyfuncitem, name, function, inner, batchable, precondition_verified + ) + ) + return records + + +def _time_variant( + case, settings, pyfuncitem, variant, function, inner, batchable, precondition_verified +): + """Warm up, then time ``function``, with no verification inside the timed loop. + + A ``CaseSkip`` from ``setup()`` here means ``setup()`` is non-deterministic, which + the ``Case`` contract forbids, so it is left to propagate rather than caught. + """ + record = _base_record(pyfuncitem, variant, precondition_verified) + state = case.setup() + for _ in range(settings["warmup"]): + output = function(state) + synchronize(output) + if case.reset is not None: + case.reset(state) + synchronize() + + sampler = WallClockSampler(inner) + samples_ms = [] + start = time.perf_counter() + while ( + len(samples_ms) < settings["iterations"] + or time.perf_counter() - start < settings["min_run_time"] + ): + samples_ms.append(sampler(function, state)) + if case.reset is not None: + case.reset(state) + synchronize() + + stats = timing_stats(samples_ms) + record.update( + { + "status": "completed", + "warmup_iterations": settings["warmup"], + "iterations": len(samples_ms), + "inner_iterations": inner, + "batchable": batchable, + "timing_method": TIMING_METHOD, + "samples_ms": samples_ms, + "timing": stats, + "metrics": _metrics(case, stats["median_ms"]), + } + ) + return record + + +def _metrics(case, median_ms): + """Derive bandwidth and FLOPs metrics from the median timing, when available.""" + metrics = {} + if median_ms > 0: + if case.bytes_moved is not None: + metrics["bandwidth_GBps"] = case.bytes_moved / (median_ms / 1.0e3) / 1.0e9 + if case.flops is not None: + metrics["tflops"] = case.flops / (median_ms / 1.0e3) / 1.0e12 + return metrics + + +def _base_record(pyfuncitem, variant, precondition_verified): + """Build the record fields that are known before timing runs.""" + params = {name: axis_value(v) for name, v in pyfuncitem.callspec.params.items()} + return { + "schema_version": "benchmark_record/v1", + "status": "pending", + "variant": variant, + "framework": _framework_for(pyfuncitem), + "case_id": case_id_for(pyfuncitem, params), + "component": pyfuncitem.module.__name__.rsplit(".", maxsplit=1)[-1], + "operation": pyfuncitem.originalname, + "params": params, + "node_id": pyfuncitem.nodeid, + "precondition_verified": precondition_verified, + "tags": [], + "unit_test": pyfuncitem.nodeid, + "source": pyfuncitem.module.__name__, + "regression_threshold": None, + } + + +def case_id_for(pyfuncitem, params) -> str: + """Build the stable identity for a benchmark record. + + Keyed on module, test function and sorted named axis values rather than the + positional pytest node ID, which would rename cases when an axis is reordered. + ``params`` is already rendered by the caller; re-applying ``axis_value`` here is safe + only because it is idempotent on the primitives that rendering produces. + """ + axes = ".".join(f"{name}{axis_value(params[name])}" for name in sorted(params)) + return f"{pyfuncitem.module.__name__}.{pyfuncitem.originalname}.{axes}" diff --git a/transformer_engine/common/testing/timing.py b/transformer_engine/common/testing/timing.py new file mode 100644 index 0000000000..3fe8141e5c --- /dev/null +++ b/transformer_engine/common/testing/timing.py @@ -0,0 +1,56 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Timing backends for benchmarkable cases.""" + +from __future__ import annotations + +import math +import statistics +import time + +from .device import synchronize + + +TIMING_METHOD = "wall-clock-with-device-sync" + + +class WallClockSampler: + """Time ``inner_iterations`` calls with the host clock and one device sync.""" + + def __init__(self, inner_iterations: int) -> None: + self.inner_iterations = inner_iterations + + def __call__(self, function, state) -> float: + """Return the mean milliseconds per call across the inner iterations.""" + output = None + start = time.perf_counter() + for _ in range(self.inner_iterations): + output = function(state) + synchronize(output) + elapsed_ms = (time.perf_counter() - start) * 1000.0 + return elapsed_ms / self.inner_iterations + + +def timing_stats(samples_ms: list[float]) -> dict[str, float]: + """Compute timing statistics in milliseconds.""" + if not samples_ms: + return { + "median_ms": 0.0, + "mean_ms": 0.0, + "min_ms": 0.0, + "max_ms": 0.0, + "stddev_ms": 0.0, + "p95_ms": 0.0, + } + + ordered = sorted(samples_ms) + p95_index = max(0, math.ceil(0.95 * len(ordered)) - 1) + return { + "median_ms": statistics.median(ordered), + "mean_ms": statistics.fmean(ordered), + "min_ms": ordered[0], + "max_ms": ordered[-1], + "stddev_ms": statistics.pstdev(ordered) if len(ordered) > 1 else 0.0, + "p95_ms": ordered[p95_index], + }