Skip to content

Commit 8bf5cb1

Browse files
kanglantkanglant
committed
Create maxtext gpu package and unit/integration tests
Use exclusion rules instead of direct pinning to disable unsupported JAX versions. This will allow newer JAX versions in future releases without requiring manual updates.
1 parent be918e9 commit 8bf5cb1

8 files changed

Lines changed: 377 additions & 5 deletions

File tree

‎.github/workflows/build_and_test_maxtext.yml‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,9 +98,49 @@ jobs:
9898
container_resource_option: "--privileged"
9999
is_scheduled_run: ${{ github.event_name == 'schedule' }}
100100

101+
maxtext_gpu_unit_tests:
102+
needs: build_and_upload_maxtext_package
103+
uses: ./.github/workflows/run_tests_against_package.yml
104+
strategy:
105+
fail-fast: false
106+
matrix:
107+
image_type: ["py312"]
108+
cuda: ["cuda12"]
109+
with:
110+
device_type: ${{ matrix.cuda }}
111+
device_name: a100-40gb-4
112+
image_type: ${{ matrix.image_type }}
113+
cloud_runner: linux-x86-a2-48-a100-4gpu
114+
pytest_marker: 'not cpu_only and not tpu_only and not integration_test'
115+
pytest_addopts: '--ignore=tests/sft_hooks_test.py'
116+
xla_python_client_mem_fraction: 0.65
117+
tf_force_gpu_allow_growth: true
118+
container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged"
119+
is_scheduled_run: ${{ github.event_name == 'schedule' }}
120+
121+
maxtext_gpu_integration_tests:
122+
needs: build_and_upload_maxtext_package
123+
uses: ./.github/workflows/run_tests_against_package.yml
124+
strategy:
125+
fail-fast: false
126+
matrix:
127+
image_type: ["py312"]
128+
cuda: ["cuda12"]
129+
with:
130+
device_type: ${{ matrix.cuda }}
131+
device_name: a100-40gb-4
132+
image_type: ${{ matrix.image_type }}
133+
cloud_runner: linux-x86-a2-48-a100-4gpu
134+
pytest_marker: 'not cpu_only and not tpu_only and integration_test'
135+
pytest_addopts: '--ignore=tests/sft_hooks_test.py'
136+
xla_python_client_mem_fraction: 0.65
137+
tf_force_gpu_allow_growth: true
138+
container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged"
139+
is_scheduled_run: ${{ github.event_name == 'schedule' }}
140+
101141
notify_failure:
102142
name: Notify failed build # creates an issue or modifies last open existing issue for failed build
103-
needs: [maxtext_cpu_unit_tests, maxtext_tpu_unit_tests, maxtext_tpu_integration_tests]
143+
needs: [maxtext_cpu_unit_tests, maxtext_tpu_unit_tests, maxtext_tpu_integration_tests, maxtext_gpu_unit_tests, maxtext_gpu_integration_tests]
104144
if: ${{ always() }}
105145
runs-on: ubuntu-latest
106146
permissions:

‎.github/workflows/run_tests_against_package.yml‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,10 @@ on:
3131
pytest_marker:
3232
required: true
3333
type: string
34+
pytest_addopts:
35+
required: false
36+
type: string
37+
default: ''
3438
is_scheduled_run:
3539
required: true
3640
type: string
@@ -51,7 +55,7 @@ permissions:
5155
contents: read
5256
jobs:
5357
run:
54-
runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }}
58+
runs-on: ${{ inputs.cloud_runner }}
5559
container:
5660
image: gcr.io/tpu-prod-env-multipod/maxtext-unit-test-${{ inputs.device_type == 'cpu' && 'tpu' || inputs.device_type }}:${{ inputs.image_type != '' && inputs.image_type }}
5761
env:
@@ -94,4 +98,4 @@ jobs:
9498
export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/src/MaxText/test_assets
9599
export MAXTEXT_PKG_DIR=$(pwd)/src/MaxText
96100
# TODO: Fix the skipped tests and remove the deselect flags
97-
.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"
101+
.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‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,10 @@ cd maxtext
5757

5858
# 2. Install dependencies in editable mode
5959
pip install uv
60+
# install the tpu package
6061
uv pip install -e .[tpu] --resolution=lowest
62+
# or install the gpu package by running the following line
63+
# uv pip install -e .[cuda12] --resolution=lowest
6164
install_maxtext_github_deps
6265
```
6366

‎base_requirements/gpu-base-requirements.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
-r requirements.txt
2+
transformer-engine[jax]

‎clean_py_env_gpu.Dockerfile‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
# This Dockerfile is used to build a Docker image containing all the
2+
# system-level dependencies required for maxtext. Since the maxtext
3+
# package itself includes all the necessary Python dependencies, this
4+
# image is designed to have a clean Python environment with just pip
5+
# and uv installed.
6+
#
7+
# Build a docker image by running:
8+
# `docker build --build-arg DEVICE=cuda12 -t <docker_image_name>:<tag> -f clean_py_env_gpu.Dockerfile .`
9+
# Or
10+
# `docker build --build-arg PYTHON_VERSION=3.11 --build-arg DEVICE=cuda12 -t <docker_image_name>:<tag> -f clean_py_env_gpu.Dockerfile .`
11+
#
12+
# How to upload the image to Google Container Registry (GCR):
13+
# e.g., DEVICE=cuda12 and PYTHON_VERSION=3.12
14+
#. gcloud init
15+
# gcloud auth configure-docker
16+
# docker tag <docker_image_name>:<tag> gcr.io/tpu-prod-env-multipod/maxtext-unit-test-cuda12:py312
17+
# docker push gcr.io/tpu-prod-env-multipod/maxtext-unit-test-cuda12:py312
18+
19+
# Default to Python 3.12.
20+
ARG PYTHON_VERSION=3.12
21+
22+
FROM nvcr.io/nvidia/cuda-dl-base:25.06-cuda12.9-devel-ubuntu24.04
23+
# TODO: enable cuda 13, e.g., nvcr.io/nvidia/cuda-dl-base:25.08-cuda13.0-devel-ubuntu24.04.
24+
25+
# Set the working directory in the container
26+
WORKDIR /maxtext
27+
28+
# Arguments
29+
ARG DEVICE
30+
31+
# Environment variables
32+
ENV DEVICE=$DEVICE
33+
ENV PYTHON_VERSION=${PYTHON_VERSION}
34+
ENV PIP_NO_CACHE_DIR=1
35+
ENV PIP_ROOT_USER_ACTION=ignore
36+
ENV PIP_BREAK_SYSTEM_PACKAGES=1
37+
ENV CLOUD_SDK_VERSION=latest
38+
# Ensure apt package installations run without manual intervention
39+
ENV DEBIAN_FRONTEND=noninteractive
40+
# See all env variables
41+
RUN env
42+
43+
# Update the package lists and install Python 3 and pip
44+
RUN apt-get update && apt-get install -y python3 python3-pip
45+
# Remove any existing pip at /usr/bin/pip
46+
RUN rm -f /usr/bin/pip
47+
# Set python3 and pip3 as the default python and pip commands
48+
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 1 && \
49+
update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 1
50+
51+
52+
# System level dependencies
53+
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/*
54+
# Add the Google Cloud SDK package repository
55+
RUN curl -fsSL https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg && \
56+
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
57+
# Install the Google Cloud SDK
58+
RUN apt-get update && apt-get install -y google-cloud-sdk && rm -rf /var/lib/apt/lists/*
59+
# Install gcsfuse
60+
RUN export GCSFUSE_REPO=gcsfuse-bullseye && \
61+
echo "deb https://packages.cloud.google.com/apt $GCSFUSE_REPO main" | tee /etc/apt/sources.list.d/gcsfuse.list && \
62+
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - && \
63+
apt update -y && apt -y install gcsfuse && \
64+
rm -rf /var/lib/apt/lists/*
65+
# See all installed system level packages
66+
RUN apt list --installed
67+
68+
RUN python3 --version
69+
RUN python3 -m pip list
70+
RUN python3 -m pip install uv
71+
RUN rm -rf /root/.cache/pip
72+
73+
# Clean python env
74+
RUN python3 -m uv pip list

0 commit comments

Comments
 (0)