Merge pull request #2463 from AI-Hypercomputer:kanglan/gpu_pyproj

PiperOrigin-RevId: 819598673

Author: Google-ML-Automation

.github/workflows/build_and_test_maxtext.yml

@@ -98,9 +98,49 @@ jobs:
       container_resource_option: "--privileged"
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
 
+  maxtext_gpu_unit_tests:
+    needs: build_and_upload_maxtext_package
+    uses: ./.github/workflows/run_tests_against_package.yml
+    strategy:
+        fail-fast: false
+        matrix:
+          image_type: ["py312"]
+          cuda: ["cuda12"]
+    with:
+      device_type: ${{ matrix.cuda }}
+      device_name: a100-40gb-4
+      image_type: ${{ matrix.image_type }}
+      cloud_runner: linux-x86-a2-48-a100-4gpu
+      pytest_marker: 'not cpu_only and not tpu_only and not integration_test'
+      pytest_addopts: '--ignore=tests/sft_hooks_test.py'
+      xla_python_client_mem_fraction: 0.65
+      tf_force_gpu_allow_growth: true
+      container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged"
+      is_scheduled_run: ${{ github.event_name == 'schedule' }}
+
+  maxtext_gpu_integration_tests:
+    needs: build_and_upload_maxtext_package
+    uses: ./.github/workflows/run_tests_against_package.yml
+    strategy:
+        fail-fast: false
+        matrix:
+          image_type: ["py312"]
+          cuda: ["cuda12"]
+    with:
+      device_type: ${{ matrix.cuda }}
+      device_name: a100-40gb-4
+      image_type: ${{ matrix.image_type }}
+      cloud_runner: linux-x86-a2-48-a100-4gpu
+      pytest_marker: 'not cpu_only and not tpu_only and integration_test'
+      pytest_addopts: '--ignore=tests/sft_hooks_test.py'
+      xla_python_client_mem_fraction: 0.65
+      tf_force_gpu_allow_growth: true
+      container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged"
+      is_scheduled_run: ${{ github.event_name == 'schedule' }}
+
   notify_failure:
     name: Notify failed build # creates an issue or modifies last open existing issue for failed build
-    needs: [maxtext_cpu_unit_tests, maxtext_tpu_unit_tests, maxtext_tpu_integration_tests]
+    needs: [maxtext_cpu_unit_tests, maxtext_tpu_unit_tests, maxtext_tpu_integration_tests, maxtext_gpu_unit_tests, maxtext_gpu_integration_tests]
     if: ${{ always() }}
     runs-on: ubuntu-latest
     permissions:

.github/workflows/run_tests_against_package.yml

@@ -31,6 +31,10 @@ on:
       pytest_marker:
         required: true
         type: string
+      pytest_addopts:
+        required: false
+        type: string
+        default: ''
       is_scheduled_run:
         required: true
         type: string
@@ -51,7 +55,7 @@ permissions:
   contents: read
 jobs:
   run:
-    runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }}
+    runs-on: ${{ inputs.cloud_runner }}
     container:
       image: gcr.io/tpu-prod-env-multipod/maxtext-unit-test-${{ inputs.device_type == 'cpu' && 'tpu' || inputs.device_type }}:${{ inputs.image_type != '' && inputs.image_type }}
       env:
@@ -94,4 +98,4 @@ jobs:
           export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/src/MaxText/test_assets
           export MAXTEXT_PKG_DIR=$(pwd)/src/MaxText
           # TODO: Fix the skipped tests and remove the deselect flags
-          .venv/bin/python3 -m pytest -v -m "${FINAL_PYTEST_MARKER}" --durations=0 --deselect "tests/aot_hlo_identical_test.py::AotHloIdenticalTest::test_default_hlo_match" --deselect "tests/tokenizer_test.py::TokenizerTest::test_detokenize"
+          .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" --durations=0 --deselect "tests/aot_hlo_identical_test.py::AotHloIdenticalTest::test_default_hlo_match" --deselect "tests/tokenizer_test.py::TokenizerTest::test_detokenize"

README.md

@@ -57,7 +57,10 @@ cd maxtext
 
 # 2. Install dependencies in editable mode
 pip install uv
+# install the tpu package
 uv pip install -e .[tpu] --resolution=lowest
+# or install the gpu package by running the following line
+# uv pip install -e .[cuda12] --resolution=lowest
 install_maxtext_github_deps
 ```
 

base_requirements/gpu-base-requirements.txt

@@ -0,0 +1,2 @@
+-r requirements.txt
+transformer-engine[jax]

clean_py_env_gpu.Dockerfile

@@ -0,0 +1,74 @@
+# This Dockerfile is used to build a Docker image containing all the
+# system-level dependencies required for maxtext. Since the maxtext
+# package itself includes all the necessary Python dependencies, this
+# image is designed to have a clean Python environment with just pip
+# and uv installed.
+#
+# Build a docker image by running:
+# `docker build --build-arg DEVICE=cuda12 -t <docker_image_name>:<tag> -f clean_py_env_gpu.Dockerfile .`
+# Or
+# `docker build --build-arg PYTHON_VERSION=3.11 --build-arg DEVICE=cuda12 -t <docker_image_name>:<tag> -f clean_py_env_gpu.Dockerfile .`
+#
+# How to upload the image to Google Container Registry (GCR):
+# e.g., DEVICE=cuda12 and PYTHON_VERSION=3.12
+#.  gcloud init
+#   gcloud auth configure-docker
+#   docker tag <docker_image_name>:<tag> gcr.io/tpu-prod-env-multipod/maxtext-unit-test-cuda12:py312
+#   docker push gcr.io/tpu-prod-env-multipod/maxtext-unit-test-cuda12:py312
+
+# Default to Python 3.12.
+ARG PYTHON_VERSION=3.12
+
+FROM nvcr.io/nvidia/cuda-dl-base:25.06-cuda12.9-devel-ubuntu24.04
+# TODO: enable cuda 13, e.g., nvcr.io/nvidia/cuda-dl-base:25.08-cuda13.0-devel-ubuntu24.04.
+
+# Set the working directory in the container
+WORKDIR /maxtext
+
+# Arguments
+ARG DEVICE
+
+# Environment variables
+ENV DEVICE=$DEVICE
+ENV PYTHON_VERSION=${PYTHON_VERSION}
+ENV PIP_NO_CACHE_DIR=1
+ENV PIP_ROOT_USER_ACTION=ignore
+ENV PIP_BREAK_SYSTEM_PACKAGES=1
+ENV CLOUD_SDK_VERSION=latest
+# Ensure apt package installations run without manual intervention
+ENV DEBIAN_FRONTEND=noninteractive
+# See all env variables
+RUN env
+
+# Update the package lists and install Python 3 and pip
+RUN apt-get update && apt-get install -y python3 python3-pip
+# Remove any existing pip at /usr/bin/pip
+RUN rm -f /usr/bin/pip
+# Set python3 and pip3 as the default python and pip commands
+RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 1 && \
+    update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 1
+
+
+# System level dependencies
+RUN apt-get update && apt-get install -y apt-utils git build-essential cmake curl pkg-config gnupg procps iproute2 ethtool lsb-release && rm -rf /var/lib/apt/lists/*
+# Add the Google Cloud SDK package repository
+RUN curl -fsSL https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg && \
+    echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee /etc/apt/sources.list.d/google-cloud-sdk.list
+# Install the Google Cloud SDK
+RUN apt-get update && apt-get install -y google-cloud-sdk && rm -rf /var/lib/apt/lists/*
+# Install gcsfuse
+RUN export GCSFUSE_REPO=gcsfuse-bullseye && \
+    echo "deb https://packages.cloud.google.com/apt $GCSFUSE_REPO main" | tee /etc/apt/sources.list.d/gcsfuse.list && \
+    curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - && \
+    apt update -y && apt -y install gcsfuse && \
+    rm -rf /var/lib/apt/lists/*
+# See all installed system level packages
+RUN apt list --installed
+
+RUN python3 --version
+RUN python3 -m pip list
+RUN python3 -m pip install uv
+RUN rm -rf /root/.cache/pip
+
+# Clean python env
+RUN python3 -m uv pip list