honor deterministic path

Signed-off-by: MahmoudAshraf97 <hassouna97.ma@gmail.com>

Author: MahmoudAshraf97

nemo/collections/asr/parts/submodules/subsampling.py

@@ -716,12 +716,14 @@ def forward(self, x, lengths):
         x = x.unsqueeze(1)  # (batch, 1, time, features)
         current_lengths = lengths
 
-        # Tracing and export cannot capture a Triton launch, and the fused kernel returns no
-        # input gradient.
+        # Tracing and export cannot capture a Triton launch, the fused kernel returns no input
+        # gradient, and its weight gradients accumulate through atomics, so their summation order
+        # varies between runs.
         if (
             self.fuse_triton
             and x.is_cuda
             and not x.requires_grad
+            and not torch.are_deterministic_algorithms_enabled()
             and not (torch.jit.is_tracing() or torch.compiler.is_exporting())
         ):
             x, current_lengths, mask = self._forward_fused(x, current_lengths)

tests/collections/asr/test_fast_subsampling.py

@@ -324,6 +324,27 @@ def test_export_takes_the_pytorch_path(tmp_path, dynamo):
     assert len(convolutions) == 5
 
 
+@pytest.mark.unit
+@pytest.mark.skipif(not CUDA_TRITON_AVAILABLE, reason="CUDA and Triton are required")
+def test_deterministic_algorithms_take_the_pytorch_path(monkeypatch):
+    """The weight gradients accumulate through atomics, whose summation order varies per run."""
+    module = _build()
+    assert module.conv.fuse_triton
+    x, lengths = _inputs(batch=2, time=520)
+
+    def unreachable(*args, **kwargs):
+        raise AssertionError("the fused path ran under deterministic algorithms")
+
+    monkeypatch.setattr("nemo.collections.asr.parts.submodules.subsampling.fused_conv_relu_dw", unreachable)
+    torch.use_deterministic_algorithms(True)
+    try:
+        out, _ = module.conv(x, lengths)
+    finally:
+        torch.use_deterministic_algorithms(False)
+
+    assert torch.isfinite(out).all()
+
+
 @pytest.mark.unit
 @pytest.mark.skipif(not CUDA_TRITON_AVAILABLE, reason="CUDA and Triton are required")
 def test_cpu_input_uses_the_pytorch_path():