diff --git a/thunder/benchmarks/__init__.py b/thunder/benchmarks/__init__.py index e5f7756c83..e1b33a6d0f 100644 --- a/thunder/benchmarks/__init__.py +++ b/thunder/benchmarks/__init__.py @@ -3,6 +3,7 @@ import tempfile import textwrap import time +from collections import UserDict from collections.abc import Callable from collections.abc import Sequence from dataclasses import dataclass @@ -2940,6 +2941,58 @@ def fn(self) -> Callable: return model +class TorchbenchBenchmark(Benchmark, metaclass=UserFacingBenchmarkMeta): + _args = ( + BenchmarkArg( + name="module_name", + description="The torchbenchmark module name (str).", + ), + BenchmarkArg( + name="device", + description="A device (str) to run on {'cpu' | 'cuda'}. Default is 'cuda'.", + ), + BenchmarkArg( + name="requires_grad", + description="Whether the model parameters require grad. Default is True.", + ), + ) + + @classmethod + @property + def name(cls) -> str: + return "torchbench" + + @classmethod + @property + def description(cls) -> str: + return "Torchbench fixture" + + @classmethod + @property + def args(cls) -> tuple[BenchmarkArg, ...]: + return cls._args + + def __init__(self, module_name, device: str = "cuda", requires_grad: bool = True): + import importlib + + module = importlib.import_module(f"torchbenchmark.models.{module_name}") + self.benchmark_cls = getattr(module, "Model", None) + + benchmark = self.benchmark_cls(test="train" if requires_grad else "eval", device=device) + + model, example = benchmark.get_module() + self.model = model + self.example_input = example + + def make_batch(self) -> tuple[list, dict]: + if isinstance(self.example_input, (dict, UserDict)): + return [], self.example_input + return self.example_input, {} + + def fn(self) -> Callable: + return self.model + + # TODO Add descriptions to the executors when listed, and list them alphabetically # TODO Allow querying benchmark for details # TODO Allow specifying benchmark arguments diff --git a/thunder/benchmarks/targets.py b/thunder/benchmarks/targets.py index a20a874ba5..9ea644b897 100644 --- a/thunder/benchmarks/targets.py +++ b/thunder/benchmarks/targets.py @@ -1,3 +1,5 @@ +import importlib +import warnings import os from collections.abc import Callable from enum import auto, Enum @@ -25,6 +27,7 @@ LitGPTGeluBenchmark, NanoGPTLayerNormBenchmark, ResNet50Benchmark, + TorchbenchBenchmark, thunder_apex_executor, thunder_apex_nvfuser_executor, thunder_cudnn_executor, @@ -812,3 +815,71 @@ def test_resnet50(benchmark, executor: Callable, compute_type: ComputeType): fn = executor(b.fn()) benchmark_for_compute_type(compute_type, benchmark, fn, args, kwargs) + + +# +# Torchbench benchmarks +# +# To setup torchbenchmark please follow the instructions to +# install it as a libraty here https://github.com/pytorch/benchmark . +# To install canary models make sure to add `--canary` to `python install.py`. + +torchbench_models = [] +torchbench_canary_models = [] +if importlib.util.find_spec("torchbenchmark"): + from torchbenchmark import _list_canary_model_paths, _list_model_paths + + torchbench_models = [os.path.basename(x) for x in _list_model_paths()] + torchbench_canary_models = [os.path.basename(x) for x in _list_canary_model_paths()] + + +@pytest.mark.skipif(not torchbench_models, reason="requires torchbenchmark to be installed") +@pytest.mark.parametrize( + "module_name,", + torchbench_models, + ids=torchbench_models, +) +@pytest.mark.parametrize( + "executor,", + executors, + ids=executors_ids, +) +@parametrize_compute_type +def test_torchbench(benchmark, module_name, executor, compute_type: ComputeType): + if not importlib.util.find_spec("torchbenchmark.models." + module_name): + pytest.skip(f"model {module_name} not installed") + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=FutureWarning) + b = TorchbenchBenchmark(module_name, device="cuda", requires_grad=is_requires_grad(compute_type)) + + args, kwargs = b.make_batch() + fn = executor(b.fn()) + + benchmark_for_compute_type(compute_type, benchmark, fn, args, kwargs) + + +@pytest.mark.skipif(not torchbench_canary_models, reason="requires torchbenchmark to be installed with flag --canary") +@pytest.mark.parametrize( + "module_name,", + torchbench_canary_models, + ids=torchbench_canary_models, +) +@pytest.mark.parametrize( + "executor,", + executors, + ids=executors_ids, +) +@parametrize_compute_type +def test_torchbench_canary(benchmark, module_name, executor, compute_type: ComputeType): + if not importlib.util.find_spec("torchbenchmark.models." + module_name): + pytest.skip(f"model {module_name} not installed") + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=FutureWarning) + b = TorchbenchBenchmark(module_name, device="cuda", requires_grad=is_requires_grad(compute_type)) + + args, kwargs = b.make_batch() + fn = executor(b.fn()) + + benchmark_for_compute_type(compute_type, benchmark, fn, args, kwargs)