Upgrade from CUDA 11.8 to CUDA 12.1

GETTING_STARTED.md

@@ -36,7 +36,7 @@ The specs on the benchmarking machines are:
 > **Prerequisites:**
 >
 > - Python minimum requirement >= 3.8
-> - CUDA 11.8
+> - CUDA 12.1
 > - NVIDIA Driver version 535.104.05
 
 To set up a virtual enviornment and install this repository
@@ -68,12 +68,19 @@ To set up a virtual enviornment and install this repository
 
     For **PyTorch**
 
+    Note: the below command assumes you have CUDA 12.1 installed locally.
+    This is the default in the provided Docker image.
+    We recommend you match this CUDA version but if you decide to run
+    with a different local CUDA version, please find the appropriate wheel
+    url to pass to the `pip install` command for `pytorch`.
+
     ```bash
     pip3 install -e '.[jax_cpu]'
-    pip3 install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/torch_stable.html'
+    pip3 install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/cu121'
     pip3 install -e '.[full]'
     ```
 
+
 <details>
 <summary>
 Per workload installations

README.md

@@ -63,7 +63,7 @@ pip3 install -e '.[full]'
 
 ```bash
 pip3 install -e '.[jax_cpu]'
-pip3 install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/torch_stable.html'
+pip3 install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/cu121'
 pip3 install -e '.[full]'
 ```
 

docker/Dockerfile

@@ -5,7 +5,7 @@
 # docker build -t <image_name> --build-arg framework=pytorch
 
 # To build Docker image
-FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04
+FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04
 
 # Installing machine packages
 RUN echo "Setting up machine"
@@ -36,8 +36,6 @@ ARG git_url=https://github.com/mlcommons/algorithmic-efficiency.git
 RUN git clone $git_url && cd /algorithmic-efficiency 
 RUN cd /algorithmic-efficiency && git checkout $branch
 
-RUN cd /algorithmic-efficiency && pip install -e '.[full]'
-
 RUN if [ "$framework" = "jax" ] ; then \
         echo "Installing Jax GPU" \
         && cd /algorithmic-efficiency \
@@ -47,17 +45,19 @@ RUN if [ "$framework" = "jax" ] ; then \
         echo "Installing Pytorch GPU" \
         && cd /algorithmic-efficiency \
         && pip install -e '.[jax_cpu]' \
-        && pip install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/torch_stable.html'; \
+        && pip install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/cu121'; \
     elif [ "$framework" = "both" ] ; then \
         echo "Installing Jax GPU and Pytorch GPU" \
         && cd /algorithmic-efficiency \
         && pip install -e '.[jax_gpu]' -f 'https://storage.googleapis.com/jax-releases/jax_cuda_releases.html' \
-        && pip install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/torch_stable.html'; \
+        && pip install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/cu121'; \
     else \
         echo "Invalid build-arg $framework: framework should be either jax, pytorch or both." >&2 \
         && exit 1 ; \
     fi
 
+RUN cd /algorithmic-efficiency && pip install -e '.[full]'
+
 RUN cd /algorithmic-efficiency && pip install -e '.[wandb]'
 
 RUN cd /algorithmic-efficiency && git fetch origin 

docker/Singularity.def

@@ -51,12 +51,12 @@ elif [ "$framework" = "pytorch" ] ; then \
 echo "Installing Pytorch GPU" \
 && cd /algorithmic-efficiency \
 && pip install -e '.[jax_cpu]' \
-&& pip install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/torch_stable.html'; \
+&& pip install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/cu121'; \
 elif [ "$framework" = "both" ] ; then \
 echo "Installing Jax GPU and Pytorch GPU" \
 && cd /algorithmic-efficiency \
 && pip install -e '.[jax_gpu]' -f 'https://storage.googleapis.com/jax-releases/jax_cuda_releases.html' \
-&& pip install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/torch_stable.html'; \
+&& pip install -e '.[pytorch_gpu]' -f 'https://download.pytorch.org/whl/cu121'; \
 else \
 echo "Invalid build-arg $framework: framework should be either jax, pytorch or both." >&2 \
 && exit 1 ; \

setup.cfg

@@ -123,25 +123,28 @@ jax_core_deps =
 
 # JAX CPU
 jax_cpu =
-  %(jax_core_deps)s
   jax==0.4.10
   jaxlib==0.4.10
+  %(jax_core_deps)s
 
 # JAX GPU
+# Note this installs both jax and jaxlib.
 jax_gpu =
+  jax==0.4.10
+  jaxlib==0.4.10+cuda12.cudnn88
   %(jax_core_deps)s
-  jax[cuda]==0.4.10
-  jaxlib==0.4.10+cuda11.cudnn86
 
 # PyTorch CPU
 pytorch_cpu =
   torch==2.1.0
   torchvision==0.16.0
 
 # PyTorch GPU
+# Note: omit the cuda suffix and installing from the appropriate
+# wheel will result in using locally installed CUDA.
 pytorch_gpu =
-  torch==2.1.0+cu118
-  torchvision==0.16.0+cu118
+  torch==2.1.0
+  torchvision==0.16.0
 
 # wandb
 wandb =