Single-GPU benchmarking in Thunder

Compile options

With the --compile option, you can test:

• torch.compile by specifying inductor

• torch eager mode by specifying eager, or

• Thunder by specifying thunder

To customize Thunder executors, in addition to nvFuser, you can append the any combination of the following to the string with an underscore:

• inductor_cat

• cudnn

• transformerengine

As an example, if you want to use cudnn as executor your terminal will looks something like:

python thunder/benchmarks/benchmark_litgpt.py --model_name Llama-2-7b-hf  --compile thunder_cudnn

and if you are testing torch.compile then it will look something like this:

python thunder/benchmarks/benchmark_litgpt.py --model_name Llama-2-7b-hf  --compile inductor

Comparing pytest runs

Another tool at your disposal is the comparison offered by pytest-benchmark:

pytest thunder/benchmarks/targets.py --benchmark-autosave -k "thunder]"
[... your changes ...]
pytest thunder/benchmarks/targets.py --benchmark-autosave -k "thunder]"
pytest-benchmark compare 0001 0002 --group-by='name'

By using --autosave pytest will save the results so that you can read or compare them later.

Creating a benchmarking class

As stated before, you need to create a class that inherits from thunder.benchmark.Benchmark as following:

from thunder.benchmarks import Benchmark, BenchmarkArg

class FooBenchmark(Benchmark):
    @classmethod
    @property
    def name(cls) -> str:
        return "foo_bench"

    @classmethod
    @property
    def description(cls) -> str:
        return "Benchmark for foo function"

The next step is to declare a list of accepted arguments from this benchmark as a property of the class and a class method that returns those arguments:

_args = (
    BenchmarkArg(name="device", description="A string representing the device. Default is 'cuda'."),
    BenchmarkArg(name="dtype", description="The dtype of the tensors. Default is thunder.float32."),
)

@classmethod
@property
def args(cls) -> tuple[BenchmarkArg, ...]:
    return cls._args

Now that the arguments are setup, the __init__() method must be implemented:

def __init__(self, device="cuda", dtype=thunder.float32):
    super().__init__(self)
    self.device: str = device
    self.dtype: dtypes.dtype = dtype

Next, you’ll want to create the data for your benchmark. To do so, you must implement a make_batch() method that prepares a valid input for the benchmark, possibly modified by the initialization arguments:

def make_batch(self) -> tuple[list, dict]:
    make = partial(make_tensor, device=self.device, dtype=self.dtype)
    return (make(10, 10),), {}

Now comes the best part, the fn() method, which should return the callable that will be benchmarked. The return callable should accept the output of make_batch()

def fn(self) -> Callable:
    def foo(a):
        return a + a

    return foo

If your benchmark doesn’t need any further steps you’d be done here however, consider the case where you want to benchmark a model, then you fn() method would look something like:

def fn(self) -> Callable:
    class FooNetwork(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.layer = torch.nn.Linear(10, 10)

        def forward(self, x):
            return self.layer(x)

    foo = FooNetwork().to(device=self.device, dtype=self.dtype).requires_grad_()
    return foo

Now this is just half of the test, what about the backward pass? In this case, you’ll need to implement a postprocess_for_backward() method to take care of that:

def postprocess_for_backward(self, out: torch.Tensor) -> torch.Tensor | None:
    # Check if backward it's needed at all
    if not self.requires_grad:
        return

    targets = make_tensor_like(out)  # fake targets
    loss = torch.nn.functional.mse_loss(out, targets)
    return loss

Declaring a test function and its parametrization

Now that your benchmarking class is ready you have nowhere to call it. To address this issue, let’s write a test_ prefixed function in thunder/benchmarks/targets.py that will use the newly created FooBenchmark class:

def test_foo(benchmark):
    bench: Benchmark = FooBenchmark(device="cuda", dtype=thunder.bfloat16)

    args, kwargs = bench.make_batch()
    benchmark(bench.fn(), *args, **kwargs)

Great! You are ready to benchmark foo()! But what if you want to test it with different Thunder executors? Here comes parametrization to help. To parametrize the function all it’s needed it’s the use of the @pytest.mark.parametrize decorator as following:

@pytest.mark.parametrize(
    "executor",
    (
        torch_executor,
        torch_compile_executor,
        thunder_executor,
    ),
    ids=("torch", "torch.compile", "thunder"),
)
def test_foo(benchmark, executor):
    bench: Benchmark = FooBenchmark(device="cuda", dtype=thunder.bfloat16)

    args, kwargs = bench.make_batch()
    fn = executor(bench.fn())

    benchmark(fn, *args, **kwargs)

Here you go, now you are ready to start benchmarking! For more information about the parametrization syntax you can get a look here.

Benchmarking forward and backward separately

As seen earlier, it’s possible to write benchmarks for models and not just standalone functions. What if you want to benchmark forward and backward pass separately? It’s possible by tweaking the test_ function you just declared in thunder/benchmarks/targets.py like so:

#[...previous parametrization omitted here...]
@parametrize_compute_type
def test_foo(benchmark, executor, compute_type: ComputeType):
    bench: Benchmark = FooBenchmark(device="cuda", dtype=thunder.bfloat16)

    args, kwargs = bench.make_batch()
    fn = executor(bench.fn())

    benchmark_for_compute_type(compute_type, benchmark, fn, *args, **kwargs)

And that’s as simple as that! Just add the decorator @parametrize_compute_type after your parametrization, add the compute_type argument, and use benchmark_for_compute_type to call the benchmark function.

Isolate benchmarks to avoid OutOfMemory errors

When running multiple benchmarks in sequence, pytest does not always do a good job cleaning up, and sometimes it happens that, while they work when called standalone, benchmarks fail anyway. The main problem we observed is that memory is not entirely freed before running the next benchmark, therefore the option --isolate-benchmarks comes in rescue. It will separate the benchmark runs, creating a sub-process for each benchmark configuration and run them one after the other. Logs of failures will be saved in the failed_benchmarks_logs folder and benchmark results will be saved in the form of json in the benchmarks_reports folder unless the THUNDER_BENCH_DIR environment variable is specified.