Make test-maxtext.sh understand CUDA_VISIBLE_DEVICES (#1203)

NVIDIA · Dec 17, 2024 · 185f664 · 185f664
1 parent c7e1bb0
commit 185f664
Showing 1 changed file with 3 additions and 0 deletions.
diff --git 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