Use deterministic algo for test_fused_attn.py

Signed-off-by: Reese Wang <rewang@nvidia.com>
NVIDIA · Dec 10, 2024 · fea31d4 · fea31d4
1 parent f49f32f
commit fea31d4
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/tests/jax/conftest.py b/tests/jax/conftest.py
@@ -27,9 +27,6 @@ def enable_fused_attn():
     """
     if get_device_compute_capability(0) >= 90:
         os.environ["NVTE_FUSED_ATTN"] = "1"
-        os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0"
     yield
     if "NVTE_FUSED_ATTN" in os.environ:
         del os.environ["NVTE_FUSED_ATTN"]
-    if "NVTE_ALLOW_NONDETERMINISTIC_ALGO" in os.environ:
-        del os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"]
diff --git a/tests/jax/test_praxis_layers.py b/tests/jax/test_praxis_layers.py
@@ -42,6 +42,17 @@
 FP8_FORMATS = [Format.E4M3, Format.HYBRID]
 
 
+@pytest.fixture(autouse=True, scope="module")
+def set_deterministic_algo():
+    """
+    test_praxis_layers will run both flax and praxis interface and expect the results are the same.
+    We need to enable the deterministic algo to ensure the results are exactly the same.
+    """
+    os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0"
+    yield
+    del os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"]
+
+
 def compare_dict(ref_fd, test_fd, rtol=1e-05, atol=1e-08):
     for key in ref_fd:
         assert key in test_fd, f"{key} not found in test dict {test_fd}"