Skip to content

Performance and determinism

Latency budgets

@pytest.mark.torch_max_latency(ms=50)
def test_latency(model, make_tensor, latency):
    latency(model, make_tensor(32, 16))

Performance regressions

def test_speed(model, make_tensor, perf_baseline):
    perf_baseline(model, make_tensor(32, 64), latency_tolerance=0.3)

pytest --torch-update-perf     # record baselines once per machine
pytest                         # later runs compare against them

Baselines live in .torch-perf-baselines.json, keyed by test and hardware fingerprint, so numbers from different machines never mix. A regression is new > baseline * (1 + tolerance) + slack; faster is never a failure. A missing baseline skips the test (--torch-perf-strict fails it instead). Peak memory is tracked on CUDA and Intel XPU.

To benchmark a model over batch sizes, dtypes and torch.compile, with an end-of-run report, see Benchmark suites.

Memory budgets and leaks

@pytest.mark.torch_max_memory(mb=512)       # peak memory of the test body (CUDA, XPU)
def test_inference_fits(model, make_tensor):
    model(make_tensor(64, 3, 224, 224))


@pytest.mark.torch_memory_leak(iterations=5)  # fails when every run leaves more behind
def test_training_step_does_not_leak(model, batch):
    train_step(model, batch)

The leak check names the tensors that accumulate and, when they hold an autograd graph, says to store loss.item() instead of loss. See torch_max_memory and torch_memory_leak.

Host-device syncs

from pytest_torch import assert_no_host_sync

assert_no_host_sync(lambda: model(x_cuda))   # .item(), .cpu(), `if tensor:` raise (CUDA only)

Nondeterminism

from pytest_torch import assert_deterministic, diagnose_nondeterminism

assert_deterministic(model, x, runs=5)
print(diagnose_nondeterminism(model, x).verdict)
# Randomness-driven: the variance disappears when the RNG is reseeded before each run.

The diagnosis separates RNG usage from kernel-level variance and checks whether torch.use_deterministic_algorithms(True) fixes it (pytest --torch-deterministic applies it to torch_seed).

Running in parallel (pytest-xdist)

pytest -n 4 --torch-xdist-gpus=0,1 --torch-gpu-memory-fraction=auto
  • Workers are spread round-robin over GPUs (CUDA or Intel XPU); torch_device maps "cuda" (or "xpu") to the worker's cuda:N (xpu:N) while test ids stay identical.
  • --torch-gpu-memory-fraction=auto splits each GPU's memory between the workers sharing it.
  • @pytest.mark.torch_gpu_exclusive serializes heavy tests per GPU across workers.
  • Each worker uses its share of the CPU cores (--torch-num-threads=auto, the default): 4 workers on 16 cores get 4 threads each. Without this, every worker starts one thread per core, and the oversubscribed machine can slow tests down by orders of magnitude. Your own OMP_NUM_THREADS or MKL_NUM_THREADS is left alone; --torch-num-threads=N sets a count, none keeps torch's default, and @pytest.mark.torch_threads(1) pins one test.
  • Record performance baselines without parallel workers: timings taken while other workers share the machine are not a fair reference.

Better failures in CI

pytest --torch-github-annotations --torch-failure-artifacts=failure-artifacts

A heatmap PNG of every tensor mismatch is saved and its path added to the error, and GitHub Actions annotations mark failing tests on the pull request.