Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add GKE A3 Ultra support #940

Draft
wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

Add GKE A3 Ultra support #940

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
08d5a03
Add a3 ultra support
samos123 Jan 9, 2025
e842543
update Dockerfile
samos123 Jan 22, 2025
565a187
remove original gpu dockerfile contents
samos123 Jan 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 40 additions & 1 deletion Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,13 +97,52 @@ COPY . .
# GPU container spec. #
################################################################################

FROM base AS gpu
# Using `FROM base as GPU` causes INTERNAL: No valid engine configs for Matmul error.
# So we're using the NVIDIA provided cuda image instead which works.
FROM nvidia/cuda:12.6.3-cudnn-devel-ubuntu22.04 as gpu
Comment on lines +100 to +102
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's make sure this change doesn't break our GPU training on AWS.


# Copy from original base
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does it mean that we will need to keep the following commands consistent with those in BASE? Does Dockerfile support functions?

# Install curl and gpupg first so that we can use them to install google-cloud-cli.
# Any RUN apt-get install step needs to have apt-get update otherwise stale package
# list may occur when previous apt-get update step is cached. See here for more info:
# https://docs.docker.com/build/building/best-practices/#apt-get
RUN apt-get update && apt-get install -y curl gnupg

RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg && \
apt-get update -y && \
apt-get install -y apt-transport-https ca-certificates gcc g++ \
git screen ca-certificates google-perftools google-cloud-cli \
gcc g++ python3 python3-venv ibverbs-utils glibc-tools

RUN ln -s /usr/bin/python3 /usr/bin/python

# Setup.
RUN mkdir -p /root
WORKDIR /root
# Introduce the minimum set of files for install.
COPY README.md README.md
COPY pyproject.toml pyproject.toml
RUN mkdir axlearn && touch axlearn/__init__.py
# Setup venv to suppress pip warnings.
ENV VIRTUAL_ENV=/opt/venv
RUN python -m venv $VIRTUAL_ENV
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
# Install dependencies.
RUN pip install flit
RUN pip install --upgrade pip


# TODO(markblee): Support extras.
ENV PIP_FIND_LINKS=https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
RUN pip install .[core,gpu]
COPY . .

# TODO(samos123): remove this once axlearn upgrades to Jax 0.4.38.
RUN pip install -U "jax[gpu]==0.4.38" "jax==0.4.38" "jaxlib==0.4.38"
COPY . .


################################################################################
# Final target spec. #
################################################################################
Expand Down
Loading