Parameterize relative and absolute tolerances

Author: voltjia

tests/test_addmm.py

@@ -55,13 +55,17 @@ def addmm(input, mat1, mat2, beta=1, alpha=1):
 
 
 @pytest.mark.parametrize(
-    "device, dtype, atol",
-    matmul._device_dtype_config(get_available_devices(), fp16_atol=0.075),
+    "device, dtype, rtol, atol",
+    matmul._device_dtype_config(
+        get_available_devices(),
+        fp16_rtol=1e-5,
+        fp16_atol=0.075,
+    ),
 )
 @pytest.mark.parametrize("k", (512,))
 @pytest.mark.parametrize("n", (512,))
 @pytest.mark.parametrize("m", (512,))
-def test(m, n, k, dtype, device, atol):
+def test(m, n, k, dtype, device, rtol, atol):
     randn_dtype = dtype if dtype != torch.float8_e5m2 else torch.float16
 
     input = torch.randn((m, n), dtype=randn_dtype, device=device)
@@ -87,4 +91,4 @@ def test(m, n, k, dtype, device, atol):
         output = addmm(input, mat1, mat2, beta=beta, alpha=alpha)
         expected = torch.addmm(input, mat1, mat2, beta=beta, alpha=alpha)
 
-    assert torch.allclose(output, expected, atol=atol)
+    assert torch.allclose(output, expected, rtol=rtol, atol=atol)

tests/test_float8_capability.py

@@ -183,13 +183,13 @@ def _parameter_capabilities(parameters):
     capabilities = {}
 
     for parameter in parameters:
-        device, dtype, atol = parameter.values
+        device, dtype, rtol, atol = parameter.values
         references = tuple(
             mark.args[0]
             for mark in parameter.marks
             if mark.name == "requires_capability"
         )
-        capabilities[(device, dtype)] = (atol, references)
+        capabilities[(device, dtype)] = ((rtol, atol), references)
 
     return capabilities
 
@@ -198,17 +198,21 @@ def test_float8_parameters_use_device_specific_capabilities():
     import tests.test_matmul as matmul
 
     capabilities = _parameter_capabilities(
-        matmul._device_dtype_config(("cuda", "mlu"), fp16_atol=0.25)
+        matmul._device_dtype_config(
+            ("cuda", "mlu"),
+            fp16_rtol=0.5,
+            fp16_atol=0.25,
+        )
     )
 
-    assert capabilities[("cuda", torch.float16)] == (0.25, ())
-    assert capabilities[("mlu", torch.float16)] == (0.25, ())
+    assert capabilities[("cuda", torch.float16)] == ((0.5, 0.25), ())
+    assert capabilities[("mlu", torch.float16)] == ((0.5, 0.25), ())
     assert capabilities[("cuda", torch.float8_e5m2)] == (
-        0.125,
+        (1e-5, 0.125),
         ("tests.capabilities.float8:float8_e5m2_cuda",),
     )
     assert capabilities[("mlu", torch.float8_e5m2)] == (
-        0.125,
+        (1e-5, 0.125),
         ("tests.capabilities.float8:float8_e5m2_mlu",),
     )
 
@@ -222,6 +226,6 @@ def test_matmul_and_addmm_use_joint_device_dtype_parameters():
             mark.args[0] for mark in function.pytestmark if mark.name == "parametrize"
         )
 
-        assert "device, dtype, atol" in parameter_names
+        assert "device, dtype, rtol, atol" in parameter_names
         assert "device" not in parameter_names
-        assert "dtype, atol" not in parameter_names
+        assert "dtype, rtol, atol" not in parameter_names

tests/test_matmul.py

@@ -64,17 +64,18 @@ def matmul(lhs, rhs):
     return output
 
 
-def _device_dtype_config(devices, fp16_atol):
+def _device_dtype_config(devices, *, fp16_rtol, fp16_atol):
     config = []
 
     for device in devices:
-        config.append(pytest.param(device, torch.float16, fp16_atol))
+        config.append(pytest.param(device, torch.float16, fp16_rtol, fp16_atol))
 
         if hasattr(torch, "float8_e5m2"):
             config.append(
                 pytest.param(
                     device,
                     torch.float8_e5m2,
+                    1e-5,
                     0.125,
                     marks=pytest.mark.requires_capability(
                         f"tests.capabilities.float8:float8_e5m2_{device}"
@@ -86,13 +87,17 @@ def _device_dtype_config(devices, fp16_atol):
 
 
 @pytest.mark.parametrize(
-    "device, dtype, atol",
-    _device_dtype_config(get_available_devices(), fp16_atol=_FP16_ATOL),
+    "device, dtype, rtol, atol",
+    _device_dtype_config(
+        get_available_devices(),
+        fp16_rtol=_FP16_RTOL,
+        fp16_atol=_FP16_ATOL,
+    ),
 )
 @pytest.mark.parametrize("k", (512,))
 @pytest.mark.parametrize("n", (512,))
 @pytest.mark.parametrize("m", (512,))
-def test(m, n, k, dtype, device, atol):
+def test(m, n, k, dtype, device, rtol, atol):
     randn_dtype = dtype if dtype != torch.float8_e5m2 else torch.float16
 
     input = torch.randn((m, k), dtype=randn_dtype, device=device)
@@ -108,5 +113,4 @@ def test(m, n, k, dtype, device, atol):
         output = matmul(input, other)
         expected = torch.matmul(input, other)
 
-    rtol = _FP16_RTOL if dtype == torch.float16 else 1e-5
     assert torch.allclose(output, expected, rtol=rtol, atol=atol)