diff --git a/tests/helpers/testcases.py b/tests/helpers/testcases.py index dd7b320c2..60e9cb0d6 100644 --- a/tests/helpers/testcases.py +++ b/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) diff --git a/tests/test_layers_convolution.py b/tests/test_layers_convolution.py index 1036d3481..5f8f131b0 100644 --- a/tests/test_layers_convolution.py +++ b/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.""" diff --git a/tests/test_layers_rnn.py b/tests/test_layers_rnn.py index 43c6388b3..0bab659e9 100644 --- a/tests/test_layers_rnn.py +++ b/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)