From 185f664c9cd823804bf360b23d1bf1f0a168edd9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20Bastien?= Date: Tue, 17 Dec 2024 10:13:11 -0500 Subject: [PATCH] Make test-maxtext.sh understand CUDA_VISIBLE_DEVICES (#1203) --- .github/container/test-maxtext.sh | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/container/test-maxtext.sh b/.github/container/test-maxtext.sh index 90e7c2488..96e5862f8 100755 --- a/.github/container/test-maxtext.sh +++ b/.github/container/test-maxtext.sh @@ -174,6 +174,9 @@ if [ $DTYPE == "fp8" ]; then fi GPUS_PER_NODE=$(nvidia-smi -L | grep -c '^GPU') +if [ "$CUDA_VISIBLE_DEVICES" != "" ]; then + GPUS_PER_NODE=`python -c 'import os; x=os.environ.get("CUDA_VISIBLE_DEVICES", ""); print(len(x.split(",")))'` +fi NGPUS=$((GPUS_PER_NODE * NODES)) # Heuristic to figure out ici and dcn of DP