Entrypoint /opt/nvidia/nvidia_entrypoint.sh; WorkingDir /opt/rosetta

NVTE_FRAMEWORK=jax pip install -e /opt/transformer-engine

GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA wget -O /opt/rosetta/rosetta/projects/pax/README.md https://raw.githubusercontent.com/NVIDIA/JAX-Toolbox/1b8b6e09fe4663ff900dbca96c054fa6cfe3da25/rosetta/rosetta/projects/pax/README.md

./nvidia_entrypoint.sh /opt/nvidia/nvidia_entrypoint.sh

GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA rmdir /opt/pax-mirror

RUN |2 GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA /bin/sh -c <<"EOF" bash -e -x git config --global user.email "${GIT_USER_EMAIL}" git config --global user.name "${GIT_USER_NAME}" cd /opt/paxml git remote add -f pax-mirror /opt/pax-mirror git fetch pax-mirror git cherry-pick $(git merge-base pax-mirror/main pax-mirror/patch/compute-capability-check)..pax-mirror/patch/compute-capability-check git cherry-pick pax-mirror/patch/final_te_configs git remote remove pax-mirror cd - EOF # buildkit

GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA pip install cuda-python

GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA pip install --no-cache-dir --upgrade certifi

RUN |2 GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA /bin/sh -c <<"EOF" bash -e -x pushd /opt/transformer-engine git remote add -f fork https://github.com/terrykong/TransformerEngine.git git checkout fix_wrong_sharding_only_fsdp popd EOF # buildkit

RUN |2 GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA /bin/sh -c <<"EOF" bash -e -x pushd /opt/xla-source git remote add -f fork https://github.com/terrykong/xla.git git checkout blas_init_fix_old_runtime export SRC_PATH_JAX=/opt/jax-source export SRC_PATH_XLA=/opt/xla-source export BAZEL_CACHE=/tmp bash build-jax.sh \
--bazel-cache ${BAZEL_CACHE} \
--src-path-jax ${SRC_PATH_JAX} \
--src-path-xla ${SRC_PATH_XLA} \
--sm all \
--clean rm -rf /tmp/* popd EOF # buildkit