fix(tests): dynamic tolerance for cuDNN TF32 precision on Ampere+ GPUs

tests/helpers/testcases.py

@@ -6,7 +6,9 @@
 
 """TestCases for aihwkit tests."""
 
+import math
 import os
+from functools import lru_cache
 from typing import Type
 from unittest import SkipTest, TestCase
 
@@ -16,6 +18,77 @@
 SKIP_CUDA_TESTS = os.getenv("SKIP_CUDA_TESTS") or not cuda.is_compiled()
 
 
+def _atol_to_decimal(atol):
+    """Convert an absolute tolerance to the ``decimal`` parameter used by
+    ``numpy.testing.assert_array_almost_equal``.
+
+    The numpy function checks ``abs(a - b) < 1.5 * 10**(-decimal)``,
+    so we solve for the largest ``decimal`` that still admits ``atol``.
+    """
+    if atol <= 0:
+        return 6  # effectively exact
+    return max(0, -math.ceil(math.log10(atol / 1.5)))
+
+
+@lru_cache(maxsize=None)
+def _probe_cuda_conv3d_tolerance(in_channels, kernel_size, n_trials=10):
+    """Measure cuDNN-vs-CUBLAS numerical divergence for Conv3d.
+
+    Uses Xavier-normalized weights so that output magnitude is ~O(1),
+    making the measured absolute tolerance directly comparable to test
+    scenarios with properly initialised weights.
+
+    Returns the max absolute difference observed across *n_trials*.
+    """
+    import torch
+    import torch.nn.functional as F
+
+    if not torch.cuda.is_available():
+        return 0.0
+
+    k = kernel_size
+    dot_dim = in_channels * k ** 3
+    max_diff = 0.0
+    with torch.no_grad():
+        for _ in range(n_trials):
+            x = torch.randn(3, in_channels, k + 1, k + 2, k + 3, device="cuda")
+            w = torch.randn(3, in_channels, k, k, k, device="cuda") / math.sqrt(dot_dim)
+            y_cudnn = F.conv3d(x, w, padding=k // 2)  # pylint: disable=not-callable
+            with torch.backends.cudnn.flags(enabled=False):
+                y_ref = F.conv3d(x, w, padding=k // 2)  # pylint: disable=not-callable
+            max_diff = max(max_diff, (y_cudnn - y_ref).abs().max().item())
+    return max_diff
+
+
+@lru_cache(maxsize=None)
+def _probe_cuda_rnn_tolerance(input_size, hidden_size, num_layers,
+                              bidirectional=False, n_trials=10):
+    """Measure cuDNN-vs-non-cuDNN divergence for ``torch.nn.RNN``.
+
+    Returns the max absolute difference observed across *n_trials* on
+    a forward pass (no training).
+    """
+    import torch
+
+    if not torch.cuda.is_available():
+        return 0.0
+
+    seq_len = 10 if bidirectional else 3
+    max_diff = 0.0
+    with torch.no_grad():
+        for _ in range(n_trials):
+            rnn = torch.nn.RNN(
+                input_size, hidden_size, num_layers,
+                bidirectional=bidirectional,
+            ).cuda()
+            x = torch.randn(seq_len, 3, input_size, device="cuda")
+            y_cudnn = rnn(x)[0]
+            with torch.backends.cudnn.flags(enabled=False):
+                y_ref = rnn(x)[0]
+            max_diff = max(max_diff, (y_cudnn - y_ref).abs().max().item())
+    return max_diff
+
+
 class AihwkitTestCase(TestCase):
     """Test case that contains common asserts and functions for aihwkit."""
 
@@ -68,3 +141,33 @@ def setUp(self) -> None:
             raise SkipTest("not compiled with CUDA support")
 
         super().setUp()
+
+    def get_cuda_decimal(self, base_atol, training_steps=0):
+        """Derive a ``decimal`` value for ``assert_array_almost_equal``
+        from a measured *base_atol* (the forward-pass tolerance probed at
+        test-session start).
+
+        On CPU (``self.use_cuda == False``) the default tight precision
+        (``decimal=6``) is returned regardless of the measured tolerance.
+
+        Args:
+            base_atol: absolute tolerance measured by one of the ``_probe_*``
+                helpers for a single forward pass on the current GPU.
+            training_steps: number of forward+backward+update iterations the
+                comparison spans.  Use 0 for a pure forward-pass comparison.
+                Each step roughly doubles the accumulated error.
+
+        Returns:
+            ``decimal`` value suitable for
+            ``numpy.testing.assert_array_almost_equal``.
+        """
+        if not self.use_cuda:
+            return 6
+
+        if base_atol <= 0:
+            return 6
+
+        # Safety margin: 3× for forward-only.  Each training step
+        # compounds the per-layer error (empirically ~2× per step).
+        margin = 3.0 * (2.0 ** training_steps)
+        return _atol_to_decimal(base_atol * margin)

tests/test_layers_convolution.py

@@ -32,7 +32,7 @@
 
 from .helpers.decorators import parametrize_over_layers
 from .helpers.layers import Conv1d, Conv1dCuda, Conv2d, Conv2dCuda, Conv3d, Conv3dCuda
-from .helpers.testcases import ParametrizedTestCase
+from .helpers.testcases import ParametrizedTestCase, _probe_cuda_conv3d_tolerance
 from .helpers.tiles import FloatingPoint, Inference, TorchInference, Custom, QuantizedTorchInference
 
 
@@ -650,7 +650,9 @@ def test_torch_original_layer(self):
         self.set_weights_from_digital_model(analog_model, model)
 
         y_analog = analog_model(x)
-        self.assertTensorAlmostEqual(y_analog, y)
+        base_atol = _probe_cuda_conv3d_tolerance(in_channels=2, kernel_size=4)
+        decimal = self.get_cuda_decimal(base_atol)
+        self.assertTensorAlmostEqual(y_analog, y, decimal=decimal)
 
     def test_torch_train_original_layer(self):
         """Test the forward and update pass, having the digital layer as reference."""

tests/test_layers_rnn.py

@@ -23,7 +23,7 @@
     LSTMCombinedWeight,
     LSTMCombinedWeightCuda,
 )
-from .helpers.testcases import ParametrizedTestCase
+from .helpers.testcases import ParametrizedTestCase, _probe_cuda_rnn_tolerance
 from .helpers.tiles import (
     FloatingPoint,
     Inference,
@@ -193,8 +193,10 @@ def get_parameters(model, analog_if) -> dict:
             rnn_analog.cuda()
             rnn.cuda()
 
+        base_atol = _probe_cuda_rnn_tolerance(input_size, hidden_size, num_layers)
+        fwd_decimal = self.get_cuda_decimal(base_atol)
         with no_grad():
-            self.assertTensorAlmostEqual(rnn(y_in)[0], rnn_analog(y_in)[0])
+            self.assertTensorAlmostEqual(rnn(y_in)[0], rnn_analog(y_in)[0], decimal=fwd_decimal)
 
         # First train analog and make sure weights differ.
         pred_analog = self.train_once(rnn_analog, y_in, y_out, True)
@@ -205,9 +207,10 @@ def get_parameters(model, analog_if) -> dict:
             for weight, weight_org in zip(analog_weights.values(), weights_org.values()):
                 assert_raises(AssertionError, assert_array_almost_equal, weight[0], weight_org[0])
 
-        # Compare with RNN.
+        # Compare after training (2 iterations of fwd+bwd+update).
+        train_decimal = self.get_cuda_decimal(base_atol, training_steps=2)
         pred = self.train_once(rnn, y_in, y_out, False)
-        assert_array_almost_equal(pred, pred_analog)
+        assert_array_almost_equal(pred, pred_analog, decimal=train_decimal)
 
         rnn_pars = get_parameters(rnn, False)
         rnn_analog_pars = get_parameters(rnn_analog, True)
@@ -220,7 +223,9 @@ def get_parameters(model, analog_if) -> dict:
 
         for par_name, par_item in rnn_pars.items():
             assert_array_almost_equal(
-                par_item.detach().cpu().numpy(), rnn_analog_pars[par_name].detach().cpu().numpy()
+                par_item.detach().cpu().numpy(),
+                rnn_analog_pars[par_name].detach().cpu().numpy(),
+                decimal=train_decimal,
             )
 
     def test_bidir_layer_training(self):
@@ -297,8 +302,12 @@ def get_parameters(model, analog_if) -> dict:
             rnn_analog.cuda()
             rnn.cuda()
 
+        base_atol = _probe_cuda_rnn_tolerance(
+            input_size, hidden_size, num_layers, bidirectional=True
+        )
+        fwd_decimal = self.get_cuda_decimal(base_atol)
         with no_grad():
-            self.assertTensorAlmostEqual(rnn(y_in)[0], rnn_analog(y_in)[0])
+            self.assertTensorAlmostEqual(rnn(y_in)[0], rnn_analog(y_in)[0], decimal=fwd_decimal)
 
         # First train analog and make sure weights differ.
         pred_analog = self.train_once_bidir(rnn_analog, y_in, y_out, True)
@@ -309,9 +318,10 @@ def get_parameters(model, analog_if) -> dict:
             for weight, weight_org in zip(analog_weights.values(), weights_org.values()):
                 self.assertNotAlmostEqualTensor(weight[0], weight_org[0])
 
-        # Compare with RNN.
+        # Compare after training (2 iterations of fwd+bwd+update).
+        train_decimal = self.get_cuda_decimal(base_atol, training_steps=2)
         pred = self.train_once_bidir(rnn, y_in, y_out, False)
-        assert_array_almost_equal(pred, pred_analog)
+        assert_array_almost_equal(pred, pred_analog, decimal=train_decimal)
 
         rnn_pars = get_parameters(rnn, False)
         rnn_analog_pars = get_parameters(rnn_analog, True)
@@ -324,7 +334,9 @@ def get_parameters(model, analog_if) -> dict:
 
         for par_name, par_item in rnn_pars.items():
             assert_array_almost_equal(
-                par_item.detach().cpu().numpy(), rnn_analog_pars[par_name].detach().cpu().numpy()
+                par_item.detach().cpu().numpy(),
+                rnn_analog_pars[par_name].detach().cpu().numpy(),
+                decimal=train_decimal,
             )
 
 
@@ -434,9 +446,14 @@ def get_parameters(model, analog_if) -> dict:
             rnn_analog.cuda()
             rnn.cuda()
 
+        # LSTMCombinedWeight uses a single weight tile with dim = input_size + hidden_size
+        base_atol = _probe_cuda_rnn_tolerance(input_size, hidden_size, num_layers)
+        fwd_decimal = self.get_cuda_decimal(base_atol)
         with no_grad():
             assert_array_almost_equal(
-                rnn(y_in)[0].detach().clone().cpu(), rnn_analog(y_in)[0].detach().clone().cpu()
+                rnn(y_in)[0].detach().clone().cpu(),
+                rnn_analog(y_in)[0].detach().clone().cpu(),
+                decimal=fwd_decimal,
             )
 
         # First train analog and make sure weights differ.
@@ -449,7 +466,8 @@ def get_parameters(model, analog_if) -> dict:
             for weight, weight_org in zip(analog_weights.values(), weights_org.values()):
                 self.assertNotAlmostEqualTensor(weight[0], weight_org[0])
 
-        # Compare with RNN.
+        # Compare after training (2 iterations of fwd+bwd+update).
+        train_decimal = self.get_cuda_decimal(base_atol, training_steps=2)
         pred = self.train_once(rnn, y_in, y_out, False)
 
-        assert_array_almost_equal(pred, pred_analog)
+        assert_array_almost_equal(pred, pred_analog, decimal=train_decimal)