Determinism fixes

Author: ericspod

monai/utils/misc.py

@@ -374,23 +374,18 @@ def set_determinism(
         for func in additional_settings:
             func(seed)
 
-    if torch.backends.flags_frozen():
-        warnings.warn("PyTorch global flag support of backends is disabled, enable it to set global `cudnn` flags.")
-        torch.backends.__allow_nonbracketed_mutation_flag = True
-
-    if seed is not None:
-        torch.backends.cudnn.deterministic = True
-        torch.backends.cudnn.benchmark = False
-    else:  # restore the original flags
-        torch.backends.cudnn.deterministic = _flag_deterministic
-        torch.backends.cudnn.benchmark = _flag_cudnn_benchmark
+    with torch.backends.__allow_nonbracketed_mutation():  # FIXME: better method without accessing private member
+        if seed is not None:
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+        else:  # restore the original flags
+            torch.backends.cudnn.deterministic = _flag_deterministic
+            torch.backends.cudnn.benchmark = _flag_cudnn_benchmark
+
     if use_deterministic_algorithms is not None:
-        if hasattr(torch, "use_deterministic_algorithms"):  # `use_deterministic_algorithms` is new in torch 1.8.0
-            torch.use_deterministic_algorithms(use_deterministic_algorithms)
-        elif hasattr(torch, "set_deterministic"):  # `set_deterministic` is new in torch 1.7.0
-            torch.set_deterministic(use_deterministic_algorithms)
-        else:
-            warnings.warn("use_deterministic_algorithms=True, but PyTorch version is too old to set the mode.")
+        # environment variable must be set to enable determinism for algorithms, alternative value is ":16:8"
+        os.environ["CUBLAS_WORKSPACE_CONFIG"] = os.environ.get("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+        torch.use_deterministic_algorithms(use_deterministic_algorithms)
 
 
 def list_to_dict(items):

tests/utils/test_set_determinism.py

@@ -42,7 +42,8 @@ def test_values(self):
         self.assertEqual(seed, get_seed())
         a = np.random.randint(seed)
         b = torch.randint(seed, (1,))
-        # tset when global flag support is disabled
+ 
+        # test when global flag support is disabled
         torch.backends.disable_global_flags()
         set_determinism(seed=seed)
         c = np.random.randint(seed)
@@ -60,12 +61,23 @@ def setUp(self):
 
     @SkipIfBeforePyTorchVersion((1, 8))  # beta feature
     @skip_if_no_cuda
-    def test_algo(self):
+    def test_algo_not_deterministic(self):
+        """
+        Test `avg_pool3d_backward_cuda` correctly raises an exception since it lacks a deterministic implementation.
+        """
         with self.assertRaises(RuntimeError):
             x = torch.randn(20, 16, 50, 44, 31, requires_grad=True, device="cuda:0")
             y = torch.nn.AvgPool3d((3, 2, 2), stride=(2, 1, 2))(x)
             y.sum().backward()
 
+    @skip_if_no_cuda
+    def test_algo_cublas_env(self):
+        """
+        Test `torch.mm` does not raise an exception with the CUBLAS_WORKSPACE_CONFIG environment variable correctly set.
+        """
+        x = torch.rand(5, 5, device="cuda:0")
+        _ = torch.mm(x, x)
+
     def tearDown(self):
         set_determinism(None, use_deterministic_algorithms=False)