Markers¶
pytest-torch registers these markers, so --strict-markers accepts them; there is nothing
to import besides pytest. Every one of them is used in
examples/test_markers.py.
| Purpose | Markers |
|---|---|
| Seeds and golden files | torch_seed, torch_golden |
| Devices and dtypes | torch_dtypes, torch_requires, torch_skip_on, torch_xfail_on |
| Modes for one test | torch_deterministic, torch_default_dtype, torch_matmul_precision, torch_no_host_sync, torch_detect_anomaly, torch_threads |
| Tolerances | torch_tolerance |
| Speed and memory | torch_max_latency, torch_max_memory, torch_memory_leak, torch_benchmark |
| Parallel runs | torch_gpu_exclusive |
The examples below assume import pytest, import torch and from torch import nn.
Where markers go¶
A marker applies to a test, to every test of a class, or to a whole module:
pytestmark = pytest.mark.torch_seed(1234) # every test in the module
@pytest.mark.torch_deterministic # every test in the class
class TestTraining:
@pytest.mark.torch_seed(0) # this test: overrides the module's seed
def test_step(self, torch_seed): ...
The closest marker wins: a test's own overrides its class's, which overrides the module's.
Two markers combine instead: every torch_requires that applies must hold, and
torch_tolerance merges all levels, the closest winning per key.
Seeds and golden files¶
torch_seed(n)¶
@pytest.mark.torch_seed(42)
def test_init(torch_seed, make_tensor):
assert torch_seed == 42
The seed the torch_seed fixture (and through it make_tensor and make_batch) uses, instead
of --torch-seed or the per-test seed of --torch-seed-mode=per-test.
torch_golden(name)¶
@pytest.mark.torch_golden("classifier_v1")
def test_matches_golden(model, make_tensor, golden):
golden.check(model(make_tensor(8, 16)), atol=1e-4)
The golden file the golden fixture uses: __golden__/classifier_v1.pt next to the test
file. Without the marker, the name is the test's; with it, several tests can share one file.
Devices and dtypes¶
torch_dtypes(*dtypes)¶
@pytest.mark.torch_dtypes(torch.float32, torch.bfloat16)
def test_batch_invariance(torch_device, torch_dtype, make_tensor):
...
Runs a test that requests torch_dtype once per dtype: [cpu-float32], [cpu-bfloat16], ...
Without it, the dtypes come from --torch-dtype, else float32 only. Combinations a device
can't run are skipped with the reason.
torch_requires(**requirements)¶
@pytest.mark.torch_requires(cuda=True, cuda_capability=(8, 0), torch=">=2.4")
def test_flash_attention(): ...
@pytest.mark.torch_requires(packages=["onnxruntime"], python=">=3.10")
def test_onnx_export(): ...
SKIPPED [1] tests/test_attention.py:12: needs CUDA
Skips unless every requirement holds:
| Requirement | Holds when |
|---|---|
cuda=True, xpu=True, mps=True |
That device is available |
gpu=True |
CUDA, XPU or MPS is available |
torch=">=2.4", python=">=3.10" |
The version matches the specifier (any PEP 440 specifier, e.g. ">=2.2,<2.6") |
compile=True |
torch.compile works on this platform and Python version |
gpu_memory_gb=16 |
The test's GPU (or the first one) has that much memory |
cuda_capability=(8, 0) |
The CUDA GPU's compute capability is at least that (Ampere for (8, 0)) |
packages=["onnxruntime"] |
Every package is importable (a string for one) |
Requirements are checked at collection, and the reason is reported at the test's location. An unknown keyword stops the run, so a typo can't silently skip a test.
torch_skip_on(*devices, dtype=None, reason=...)¶
@pytest.mark.torch_skip_on("mps", reason="float64 isn't supported on MPS")
def test_float64(torch_device):
...
Skips this test's runs on these devices only; the other devices still run it. "cuda"
matches every CUDA device (cuda:1 too), and dtype= (one or a list) narrows it to some
torch_dtype runs. The test must use the torch_device fixture.
torch_xfail_on(*devices, dtype=None, reason=..., raises=None, strict=True)¶
@pytest.mark.torch_xfail_on("mps", dtype=torch.bfloat16, reason="missing kernel", raises=NotImplementedError)
@pytest.mark.torch_dtypes(torch.float32, torch.bfloat16)
def test_layer(torch_device, torch_dtype):
...
XFAIL tests/test_layers.py::test_layer[mps-bfloat16] - missing kernel
An expected failure on these devices (and dtypes) only. It's strict by default: when a
torch release fixes the operation, the run reports XPASS(strict) as a failure, so the
marker doesn't outlive the bug. With raises=, any other error stays a real failure.
Modes for one test¶
These markers set a mode for the test, function-scoped fixtures included (so a model built in a fixture is affected), and restore the previous mode afterwards.
torch_deterministic¶
@pytest.mark.torch_deterministic
def test_reproducible_step(model, batch):
...
torch.use_deterministic_algorithms(True) during the test: a nondeterministic kernel raises
instead of varying silently. --torch-deterministic does the same for every test using
torch_seed.
torch_default_dtype(dtype)¶
@pytest.mark.torch_default_dtype(torch.float64)
def test_numerics():
assert nn.Linear(2, 2).weight.dtype == torch.float64
torch.set_default_dtype during the test, e.g. float64 for numerical and gradient checks.
Takes a dtype or its name.
torch_matmul_precision(precision)¶
@pytest.mark.torch_matmul_precision("highest")
def test_exact_float32(model, make_tensor):
...
torch.set_float32_matmul_precision during the test: "highest", "high" or "medium".
"highest" also turns off TF32 in cuDNN convolutions. On Ampere GPUs and newer, TF32 is
the usual reason exact float32 comparisons fail on CUDA but pass on CPU.
torch_no_host_sync¶
@pytest.mark.torch_no_host_sync
def test_forward_never_syncs(model, make_tensor, torch_device):
model(make_tensor(8, 16))
On CUDA, operations that synchronize the host with the GPU (.item(), .cpu(),
if tensor:) raise during the test, with the offending line in the
traceback. Such syncs stall the GPU pipeline in production. On other devices the test runs
normally and its report says it wasn't checked.
torch_detect_anomaly¶
@pytest.mark.torch_detect_anomaly
def test_gradients(model, batch):
model(batch).sum().backward()
Autograd anomaly detection during the test: a NaN produced in the backward pass raises,
naming the backward function (e.g. SqrtBackward0) and showing where its forward
operation ran. It slows autograd down, so use it to investigate rather than everywhere.
torch_threads(n)¶
@pytest.mark.torch_threads(1)
def test_single_threaded_timing(latency, model, make_tensor):
...
Torch uses n CPU threads during the test. --torch-num-threads sets the count for the
whole run; under pytest-xdist its default divides the cores between the workers.
Tolerances¶
torch_tolerance(**entries)¶
@pytest.mark.torch_tolerance(bfloat16=(5e-2, 1e-2), compile_parity={"atol": 1e-3})
def test_low_precision(model, make_tensor):
...
@pytest.mark.torch_tolerance(**{"cuda.float32": (1e-3, 1e-4)}) # device keys need **
class TestOnAmpere: ...
Default tolerances for comparisons that don't pass rtol/atol, with the keys of
[tool.pytest_torch.tolerance]: a dtype, a check, or either prefixed with a device. Values
are (rtol, atol) or a dict with either. See Tolerances.
Speed and memory¶
torch_max_latency(ms=...)¶
@pytest.mark.torch_max_latency(ms=50)
def test_latency(model, make_tensor, latency):
latency(model, make_tensor(32, 16))
The latency fixture fails when the median latency exceeds the budget.
torch_max_memory(mb=...)¶
@pytest.mark.torch_max_memory(mb=512)
def test_inference_fits(model, make_tensor):
with torch.inference_mode():
model(make_tensor(64, 3, 224, 224))
Failed: peak memory 812.4 MB exceeds the budget of 512 MB on cuda:0 (memory allocated during the test, above the 98.0 MB allocated before it)
Fails when the test body's peak memory, above what was allocated before it, exceeds mb.
Only CUDA and XPU keep the allocator statistics this needs; elsewhere the test runs
normally and its report says the budget wasn't checked.
torch_memory_leak(iterations=5)¶
HISTORY = []
@pytest.mark.torch_memory_leak
def test_training_step(model, batch):
loss = model(batch).sum()
HISTORY.append(loss) # should be loss.item()
Memory leak: every run of the test leaves more behind
live tensors after each run: 7, 10, 13, 16, 19, 22
new: () float32 on cpu with an autograd graph (SumBackward0), held by a list of 6 item(s)
Hint: tensors that still hold their autograd graph are being kept: store loss.item() or loss.detach() instead of loss.
Runs the test body once to warm up (lazy initialization and caches may allocate once),
then iterations more times with the same fixtures. It fails when the number of live
tensors, or the device memory allocated, grows after every run, and describes the new
tensors: shape, dtype, device, autograd graph, and the container holding them. Counting
tensors works on every device. Since the body runs several times, keep its side effects
repeatable. With --torch-repro-dir, a failure bundle holds the inputs of one run.
torch_benchmark¶
Marks a benchmark: deselected unless --torch-benchmark is given. Every
BenchmarkSuite carries it; mark your own slow measurements with
it to keep them out of normal runs.
Parallel runs¶
torch_gpu_exclusive(timeout=1800)¶
@pytest.mark.torch_gpu_exclusive
def test_full_size_model(model_loader, torch_device): ...
Under pytest-xdist, only one such test runs at a time on each GPU, across workers. The
others wait for it; after timeout seconds a waiting test errors with a TimeoutError. Use it for tests that need most of a GPU's memory. See
running in parallel.