Merge pull request #736 from rapidsai/branch-25.02

Forward-merge branch-25.02 into branch-25.04
rapidsai · Feb 18, 2025 · c833c43 · c833c43
2 parents a54fdfc + eb613ce
commit c833c43
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/context/entrypoint.sh b/context/entrypoint.sh
@@ -23,6 +23,14 @@ if [ "$EXTRA_PIP_PACKAGES" ]; then
     timeout ${PIP_TIMEOUT:-600} pip install $EXTRA_PIP_PACKAGES || exit $?
 fi
 
+if [ "$(uname -m)" = "aarch64" ]; then
+    # Check if the CUDA version is 12.8
+    if [[ "$CUDA_VERSION" = 12.8* ]]; then
+        export NCCL_CUMEM_HOST_ENABLE=0
+        echo "Set NCCL_CUMEM_HOST_ENABLE=0 for ARM with CUDA 12.8"
+    fi
+fi
+
 # Run whatever the user wants.
 if [ "${UNQUOTE}" = "true" ]; then
     exec $@