mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Disable some tests that fail on H100 in CI.

Browse Source 
PiperOrigin-RevId: 612637375
This commit is contained in:
Peter Hawkins 2024-03-04 16:59:05 -08:00 committed by jax authors
parent 3fe65e2005
commit 6207977fac
1 changed files with 18 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
tests/BUILD
View File

@ -416,6 +416,9 @@ jax_test(
backend_tags = {
"cpu": ["notsan"], # Test times out.
},
disable_configs = [
"gpu_h100", # TODO(b/328050517): ptxas compilation failures
],
shard_count = {
"cpu": 40,
"gpu": 40,
@ -731,6 +734,9 @@ jax_test(
name = "pytorch_interoperability_test",
srcs = ["pytorch_interoperability_test.py"],
disable_backends = ["tpu"],
disable_configs = [
"gpu_h100", # Pytorch H100 build times out in Google's CI.
],
tags = [
# PyTorch leaks dlpack metadata https://github.com/pytorch/pytorch/issues/117058, and
# compilation times out on CPU.
@ -858,6 +864,9 @@ jax_test(
"optonly",
],
},
disable_configs = [
"gpu_h100", # TODO(phawkins): numerical failure on h100
],
shard_count = {
"cpu": 40,
"gpu": 40,
@ -938,6 +947,9 @@ jax_test(
"cpu_x32": ["--jax_num_generated_cases=40"],
"gpu": ["--jax_num_generated_cases=40"],
},
disable_configs = [
"gpu_h100", # TODO(b/328050517): ptxas compilation failures
],
shard_count = {
"cpu": 50,
"gpu": 50,
@ -967,6 +979,9 @@ jax_test(
"noasan", # Times out under asan.
],
},
disable_configs = [
"gpu_h100", # TODO(b/328050517): ptxas compilation failures
],
shard_count = {
"cpu": 5,
"gpu": 20,
@ -1223,6 +1238,9 @@ jax_test(
backend_variant_args = {
"tpu_pjrt_c_api": ["--jax_num_generated_cases=1"],
},
disable_configs = [
"gpu_h100", # TODO(b/328050517): ptxas compilation failures
],
enable_configs = [
"gpu",
"cpu",