From 13e45c8953faa8e6fe13316ae61c06f16e94a322 Mon Sep 17 00:00:00 2001 From: Rahul Batra Date: Thu, 30 Mar 2023 18:34:47 +0000 Subject: [PATCH] [ROCm]: Run pmap test on specific number of GPUs --- build/rocm/run_multi_gpu.sh | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/build/rocm/run_multi_gpu.sh b/build/rocm/run_multi_gpu.sh index caad6f1a5..290e22f60 100755 --- a/build/rocm/run_multi_gpu.sh +++ b/build/rocm/run_multi_gpu.sh @@ -16,5 +16,17 @@ set -eux # run test module with multi-gpu requirements. We currently do not have a way to filter tests. # this issue is also tracked in https://github.com/google/jax/issues/7323 -python3 -m pytest --reruns 3 -x tests/pmap_test.py +cmd=$(lspci|grep 'controller'|grep 'AMD/ATI'|wc -l) +echo $cmd + +if [[ $cmd -gt 8 ]]; then + export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 && python3 -m pytest --reruns 3 -x tests/pmap_test.py +elif [[ $cmd -gt 4 ]]; then + export HIP_VISIBLE_DEVICES=0,1,2,3 && python3 -m pytest --reruns 3 -x tests/pmap_test.py +elif [[ $cmd -gt 2 ]]; then + export HIP_VISIBLE_DEVICES=0,1 && python3 -m pytest --reruns 3 -x tests/pmap_test.py +else + export HIP_VISIBLE_DEVICES=0 && python3 -m pytest --reruns 3 -x tests/pmap_test.py +fi + python3 -m pytest --reruns 3 -x tests/multi_device_test.py