Skip 128k test in T4

microsoft · Jul 3, 2024 · b68b6e9 · b68b6e9
1 parent 3b8ff12
commit b68b6e9
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/onnxruntime/test/python/transformers/test_sparse_attention.py b/onnxruntime/test/python/transformers/test_sparse_attention.py
@@ -1129,7 +1129,8 @@ def run_relevance_test(self, sm: int):
         device = torch.device("cuda", device_id)
         with torch.no_grad():
             # Test long sequence when GPU memory is enough (need about 12 GB for 128K sequence length)
-            if torch.cuda.get_device_properties(device_id).total_memory > 13 * 1024 * 1024 * 1024:
+            # The 128k tests fails randomly in T4 GPU, increase memory threshold for now.
+            if torch.cuda.get_device_properties(device_id).total_memory > 20 * 1024 * 1024 * 1024:
                 self.run_relevance_no_past_128k(sm, device)
                 self.run_relevance_past_128k(sm, device)
             self.run_relevance_no_past(sm, device)