fix distributed jax test

Signed-off-by: Charlene Yang <[email protected]>
NVIDIA · Dec 4, 2023 · 7b8ff67 · 7b8ff67
1 parent 225b560
commit 7b8ff67
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py
@@ -73,6 +73,9 @@ def generate_inputs(self, shape, mesh_resource, with_bias, attn_mask_type, dtype
     @pytest.mark.parametrize('dtype', DTYPES)
     def test_self_attn(self, device_count, mesh_shape, mesh_axes, mesh_resource, data_shape,
                        attn_bias_type, attn_mask_type, dtype):
+        # TODO (cyang): remove this when cudnn fe v1 has included it for dbias cases
+        os.environ["CUDNN_FRONTEND_ATTN_DP_WORKSPACE_LIMIT"] = "-1"
+
         dropout_prob = 0.0
         is_training = True
         scaling_factor = 1.0