[gpu] Provide an option to install CUDA 11 to support GCE A2 instances (

#803)

gpu/install_gpu_driver.sh

@@ -36,16 +36,19 @@ CUDA_VERSION=$(get_metadata_attribute 'cuda-version' '10.2')
 readonly CUDA_VERSION
 
 # Parameters for NVIDIA-provided Debian GPU driver
-readonly DEFAULT_NVIDIA_DEBIAN_GPU_DRIVER_URL='http://us.download.nvidia.com/XFree86/Linux-x86_64/450.51/NVIDIA-Linux-x86_64-450.51.run'
+readonly DEFAULT_NVIDIA_DEBIAN_GPU_DRIVER_URL='http://us.download.nvidia.com/tesla/450.51.06/NVIDIA-Linux-x86_64-450.51.06.run'
 NVIDIA_DEBIAN_GPU_DRIVER_URL=$(get_metadata_attribute 'gpu-driver-url' "${DEFAULT_NVIDIA_DEBIAN_GPU_DRIVER_URL}")
 readonly NVIDIA_DEBIAN_GPU_DRIVER_URL
 readonly DEFAULT_NVIDIA_DEBIAN_CUDA_URL_10_1='http://developer.download.nvidia.com/compute/cuda/10.1/Prod/local_installers/cuda_10.1.243_418.87.00_linux.run'
 readonly DEFAULT_NVIDIA_DEBIAN_CUDA_URL_10_2='http://developer.download.nvidia.com/compute/cuda/10.2/Prod/local_installers/cuda_10.2.89_440.33.01_linux.run'
+readonly DEFAULT_NVIDIA_DEBIAN_CUDA_URL_11_0='https://developer.download.nvidia.com/compute/cuda/11.0.3/local_installers/cuda_11.0.3_450.51.06_linux.run'
 
 if [[ "${CUDA_VERSION}" == "10.1" ]]; then
   readonly DEFAULT_NVIDIA_DEBIAN_CUDA_URL=${DEFAULT_NVIDIA_DEBIAN_CUDA_URL_10_1}
-else
+elif [[ "${CUDA_VERSION}" == "10.2" ]]; then 
   readonly DEFAULT_NVIDIA_DEBIAN_CUDA_URL=${DEFAULT_NVIDIA_DEBIAN_CUDA_URL_10_2}
+else
+  readonly DEFAULT_NVIDIA_DEBIAN_CUDA_URL=${DEFAULT_NVIDIA_DEBIAN_CUDA_URL_11_0}
 fi
 
 NVIDIA_DEBIAN_CUDA_URL=$(get_metadata_attribute 'cuda-url' "${DEFAULT_NVIDIA_DEBIAN_CUDA_URL}")
@@ -60,7 +63,7 @@ readonly NVIDIA_UBUNTU_REPOSITORY_CUDA_PIN="${NVIDIA_UBUNTU_REPOSITORY_URL}/cuda
 readonly DEFAULT_NCCL_REPO_URL='https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb'
 NCCL_REPO_URL=$(get_metadata_attribute 'nccl-repo-url' "${DEFAULT_NCCL_REPO_URL}")
 readonly NCCL_REPO_URL
-NCCL_VERSION=$(get_metadata_attribute 'nccl-version' '2.7.6')
+NCCL_VERSION=$(get_metadata_attribute 'nccl-version' '2.7.8')
 readonly NCCL_VERSION
 
 # Parameters for Ubuntu-provided NVIDIA GPU driver