Gpu image pipeline (#169)

* Add gpu image creation pipeline

Signed-off-by: Kunjan <kunjanp@google.com>

* Move jax[cuda] installation after requirements.txt to avoid clobbering of it by jax installation in requirements.txt

Signed-off-by: Kunjan <kunjanp@google.com>

---------

Signed-off-by: Kunjan <kunjanp@google.com>

Author: coolkp

.github/workflows/UploadDockerImages.yml

@@ -35,3 +35,16 @@ jobs:
     - name: build maxdiffusion jax nightly image
       run: |
         bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_nightly MODE=nightly PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_nightly
+
+  build-gpu-image:
+    runs-on: ["self-hosted", "e2", "cpu"]
+    steps:
+    - uses: actions/checkout@v3
+    - name: Cleanup old docker images
+      run: docker system prune --all --force
+    - name: build maxdiffusion jax stable stack gpu image
+      run: |
+        bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_stable_stack_gpu MODE=stable_stack PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_stable_stack_gpu BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:latest DEVICE=gpu
+    - name: build maxdiffusion jax nightly image
+      run: |
+         bash .github/workflows/build_and_upload_images.sh CLOUD_IMAGE_NAME=maxdiffusion_jax_nightly_gpu MODE=nightly PROJECT=tpu-prod-env-multipod LOCAL_IMAGE_NAME=maxdiffusion_jax_nightly DEVICE=gpu

.github/workflows/build_and_upload_images.sh

@@ -34,13 +34,15 @@ for ARGUMENT in "$@"; do
     echo "$KEY"="$VALUE"
 done
 
+export DEVICE="${DEVICE:-tpu}"
+
 if [[ ! -v CLOUD_IMAGE_NAME ]] || [[ ! -v PROJECT ]] || [[ ! -v MODE ]] ; then
   echo "You must set CLOUD_IMAGE_NAME, PROJECT and MODE"
   exit 1
 fi
 
 gcloud auth configure-docker us-docker.pkg.dev --quiet
-bash docker_build_dependency_image.sh LOCAL_IMAGE_NAME=$LOCAL_IMAGE_NAME MODE=$MODE
+bash docker_build_dependency_image.sh LOCAL_IMAGE_NAME=$LOCAL_IMAGE_NAME MODE=$MODE DEVICE=$DEVICE
 image_date=$(date +%Y-%m-%d)
 
 # Upload only dependencies image

maxdiffusion_gpu_dependencies.Dockerfile

@@ -22,8 +22,7 @@ RUN apt-get update && apt-get install -y google-cloud-sdk
 # Set environment variables for Google Cloud SDK
 ENV PATH="/usr/local/google-cloud-sdk/bin:${PATH}"
 
-# Upgrade libcusprase to work with Jax
-RUN apt-get update && apt-get install -y libcusparse-12-3
+
 
 ARG MODE
 ENV ENV_MODE=$MODE
@@ -46,5 +45,4 @@ RUN ls .
 RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION DEVICE=${ENV_DEVICE}"
 RUN --mount=type=cache,target=/root/.cache/pip bash setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} DEVICE=${ENV_DEVICE}
 
-
 WORKDIR /deps

setup.sh

@@ -55,6 +55,9 @@ if [[ -n $JAX_VERSION && ! ($MODE == "stable" || -z $MODE) ]]; then
   exit 1
 fi
 
+# Install dependencies from requirements.txt first
+pip3 install -U -r requirements.txt || echo "Failed to install dependencies in the requirements" >&2
+
 # Install JAX and JAXlib based on the specified mode
 if [[ "$MODE" == "stable" || ! -v MODE ]]; then
   # Stable mode
@@ -78,7 +81,7 @@ if [[ "$MODE" == "stable" || ! -v MODE ]]; then
         pip3 install "jax[cuda12]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
     fi
     export NVTE_FRAMEWORK=jax
-    pip3 install git+https://github.com/NVIDIA/TransformerEngine.git@stable
+    pip3 install transformer_engine[jax]==2.1.0
   fi
 
 elif [[ $MODE == "nightly" ]]; then
@@ -106,8 +109,5 @@ else
   exit 1
 fi
 
-# Install dependencies from requirements.txt
-pip3 install -U -r requirements.txt || echo "Failed to install dependencies in the requirements" >&2
-
 # Install maxdiffusion
 pip3 install -U . || echo "Failed to install maxdiffusion" >&2