Sign inSign up

sanjeevrg/maxtext-gpu-pretraining:stable-2026-06-29

Manifest digest

sha256:31dc68b2308afb4f4d04df4f64ab0010f813f7e27a49d091eeca4e54d144887e

OS/ARCH

linux/amd64

Compressed size

9.99 GB

Last pushed

3 months by sanjeevrg

Type

Image

Manifest digest

sha256:31dc68b2308afb4f4d04df4f64ab0010f813f7e27a49d091eeca4e54d144887e

Image Layers

Image Layers

1ARG RELEASE0 B
2ARG LAUNCHPAD_BUILD_ARCH0 B
3LABEL org.opencontainers.image.version=24.040 B
4ADD file ... in / 29.18 MB
5CMD ["/bin/bash"]0 B
6ARG ENABLE_MITMPROXY=00 B
7RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c32 B
8RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c1 MB
9RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c32 B
10RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c3.43 MB
11ARG GDRCOPY_VERSION=2.5.10 B
12ARG HPCX_VERSION=2.260 B
13ARG RDMACORE_VERSION=59.10 B
14ARG MOFED_VERSION=5.4-rdmacore59.10 B
15ARG OPENUCX_VERSION=1.20.00 B
16ARG OPENMPI_VERSION=4.1.70 B
17ARG EFA_VERSION=1.48.00 B
18ENV GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.26 MOFED_VERSION=5.4-rdmacore59.1 OPENUCX_VERSION=1.20.00 B
19ARG TARGETARCH=amd640 B
20RUN |9 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.26195.6 MB
21ENV OPAL_PREFIX=/usr/local/mpi PATH=/usr/local/mpi/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/ucx/bin:/opt/amazon/efa/bin0 B
22ENV OMPI_MCA_coll_hcoll_enable=00 B
23ARG CUDA_VERSION=13.2.1.0090 B
24ARG CUDA_DRIVER_VERSION=595.58.030 B
25ARG NVVM_VERSION=13.2.780 B
26ARG DOCA_VERSION=3.2.00 B
27ENV CUDA_VERSION=13.2.1.009 CUDA_DRIVER_VERSION=595.58.03 NVVM_VERSION=13.2.78 DOCA_VERSION=3.2.00 B
28RUN |13 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.26154.58 MB
29RUN |13 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.2614.4 KB
30ENV _CUDA_COMPAT_PATH=/usr/local/cuda/compat ENV=/etc/shinit_v2 BASH_ENV=/etc/bash.bashrc SHELL=/bin/bash0 B
31LABEL com.nvidia.volumes.needed=nvidia_driver com.nvidia.cuda.version=9.00 B
32ARG NCCL_VERSION=2.30.4 CUBLAS_VERSION=13.4.1.1 CUFFT_VERSION=12.2.0.46 CURAND_VERSION=10.4.2.550 B
33ENV NCCL_VERSION=2.30.4 CUBLAS_VERSION=13.4.1.1 CUFFT_VERSION=12.2.0.46 CURAND_VERSION=10.4.2.550 B
34ADD docs.tgz / # buildkit11.26 KB
35RUN |48 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.26185 B
36ARG _LIBPATH_SUFFIX=0 B
37ENV PATH=/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/mpi/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/ucx/bin:/opt/amazon/efa/bin LD_LIBRARY_PATH=/usr/local/cuda/compat/lib:/usr/local/nvidia/lib:/usr/local/nvidia/lib64 NVIDIA_VISIBLE_DEVICES=all NVIDIA_DRIVER_CAPABILITIES=compute,utility,video0 B
38RUN |49 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.2692.53 MB
39RUN |49 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.261.44 MB
40COPY entrypoint/ /opt/nvidia/ # buildkit5.98 KB
41ENV NVIDIA_PRODUCT_NAME=CUDA0 B
42ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"]0 B
43RUN |49 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.26408 B
44RUN |49 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.2632 B
45ARG ENABLE_MITMPROXY=00 B
46ARG ENABLE_FIPS=00 B
47RUN |2 ENABLE_MITMPROXY=0 ENABLE_FIPS=0 /bin/sh32 B
48RUN |2 ENABLE_MITMPROXY=0 ENABLE_FIPS=0 /bin/sh111.01 MB
49ARG CUDA_COMPONENT_LIST=cccl crt nvrtc driver-dev0 B
50ENV CUDA_COMPONENT_LIST=cccl crt nvrtc driver-dev0 B
51RUN |3 ENABLE_MITMPROXY=0 ENABLE_FIPS=0 CUDA_COMPONENT_LIST=cccl214.88 MB
52ARG VERSION_LIST=CURAND_VERSION CUBLAS_VERSION CUFILE_VERSION0 B
53RUN |4 ENABLE_MITMPROXY=0 ENABLE_FIPS=0 CUDA_COMPONENT_LIST=cccl1.66 GB
54ENV LIBRARY_PATH=/usr/local/cuda/lib64/stubs:0 B
55RUN |4 ENABLE_MITMPROXY=0 ENABLE_FIPS=0 CUDA_COMPONENT_LIST=cccl32 B
56LABEL com.nvidia.nccl.version=2.30.4 com.nvidia.cublas.version=13.4.1.1 com.nvidia.cufft.version=12.2.0.46 com.nvidia.curand.version=10.4.2.550 B
57ARG ENABLE_MITMPROXY=00 B
58RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c32 B
59RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c32 B
60RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c64.13 MB
61RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c2.51 GB
62LABEL com.nvidia.nccl.version=2.30.4 com.nvidia.cublas.version=13.4.1.1 com.nvidia.cufft.version=12.2.0.46 com.nvidia.curand.version=10.4.2.550 B
63RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c32 B
64ARG TARGETARCH=amd640 B
65ARG ENABLE_FIPS=00 B
66ARG ENABLE_MITMPROXY=00 B
67RUN |3 TARGETARCH=amd64 ENABLE_FIPS=0 ENABLE_MITMPROXY=032 B
68RUN |3 TARGETARCH=amd64 ENABLE_FIPS=0 ENABLE_MITMPROXY=021.14 MB
69COPY cuda-*.patch /tmp # buildkit503 B
70RUN |3 TARGETARCH=amd64 ENABLE_FIPS=0 ENABLE_MITMPROXY=0345.52 MB
71ENV LIBRARY_PATH=/usr/local/cuda/lib64/stubs:/usr/local/cuda/lib64/stubs:0 B
72RUN |3 TARGETARCH=amd64 ENABLE_FIPS=0 ENABLE_MITMPROXY=032 B
73ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.040 B
74ARG [email protected]0 B
75ARG GIT_USER_NAME=JAX Toolbox0 B
76ARG JAX_TOOLBOX_REF=4847ab0b2f900768ec3e082a1a9340f55bca0c630 B
77ARG NSIGHT_SYSTEMS_VERSION_OVERRIDE0 B
78ENV CUDA_BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.040 B
79ENV JAX_TOOLBOX_REF=4847ab0b2f900768ec3e082a1a9340f55bca0c630 B
80RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX46.19 MB
81RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX188 B
82RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX145 B
83ADD --chmod=777 git-clone.sh pip-finalize.sh /usr/local/bin/2.66 KB
84RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX5.29 MB
85RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX84.35 MB
86ADD pip-vcs-equivalency.patch /opt/pip/ # buildkit1.55 KB
87RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX72.31 MB
88ENV PIP_BREAK_SYSTEM_PACKAGES=10 B
89RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX14.59 KB
90RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX4.89 MB
91RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX172 B
92ADD symlnk-cudnn.sh /usr/local/bin # buildkit1 KB
93RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX939 B
94ADD symlnk-nccl.sh /usr/local/bin # buildkit715 B
95RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX1.85 KB
96ADD install-nccl-sanity-check.sh /usr/local/bin # buildkit597 B
97ADD nccl-sanity-check.cu /opt # buildkit2.13 KB
98RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX9.37 KB
99ADD parallel-launch /usr/local/bin/ # buildkit535 B
100ADD install-nsight-systems.sh /usr/local/bin # buildkit714 B
101RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX32 B
102RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX1.8 KB
103ADD symlnk-nvshmem.sh /usr/local/bin # buildkit685 B
104RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX234 B
105ENV MANIFEST_FILE=/opt/manifest.d/manifest.yaml0 B
106ADD manifest.yaml create-distribution.sh bump.sh /opt/manifest.d/6.97 KB
107COPY patches/ /opt/manifest.d/patches/ # buildkit4.76 KB
108RUN /bin/sh -c if [561 B
109RUN /bin/sh -c apt-get update36.36 MB
110RUN /bin/sh -c apt-get install2.09 MB
111RUN /bin/sh -c echo "deb265 B
112RUN /bin/sh -c curl https://packages.cloud.google.com/apt/doc/apt-key.gpg946 B
113RUN /bin/sh -c apt-get update156.56 MB
114ENV PATH=/usr/local/google-cloud-sdk/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/mpi/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/ucx/bin:/opt/amazon/efa/bin0 B
115RUN /bin/sh -c apt-get update32 B
116ARG MODE0 B
117ENV ENV_MODE=stable0 B
118ARG JAX_VERSION0 B
119ENV ENV_JAX_VERSION=0 B
120ARG DEVICE0 B
121ENV ENV_DEVICE=gpu0 B
122ARG PACKAGE_DIR0 B
123ENV PACKAGE_DIR=src0 B
124ENV MAXTEXT_ASSETS_ROOT=/deps/src/maxtext/assets0 B
125ENV MAXTEXT_TEST_ASSETS_ROOT=/deps/tests/assets0 B
126ENV MAXTEXT_PKG_DIR=/deps/src/maxtext0 B
127ENV MAXTEXT_REPO_ROOT=/deps0 B
128WORKDIR /deps94 B
129COPY src/dependencies/extra_deps/ src/dependencies/extra_deps/ # buildkit1.09 KB
130COPY src/dependencies/requirements/ src/dependencies/requirements/ # buildkit6.38 KB
131COPY src/dependencies/scripts/ src/dependencies/scripts/ # buildkit13.58 KB
132COPY src/maxtext/integration/vllm/ src/maxtext/integration/vllm/ # buildkit22.38 KB
133RUN |4 MODE=stable JAX_VERSION= DEVICE=gpu32 B
134RUN |4 MODE=stable JAX_VERSION= DEVICE=gpu4.19 GB
135COPY src/maxtext/ src/maxtext/ # buildkit9.69 MB
136COPY tests*/ tests/ # buildkit21.62 MB
137COPY pytest.ini* pytest.ini # buildkit885 B
138COPY benchmarks*/ benchmarks/ # buildkit279.61 KB
139ARG INCLUDE_TEST_ASSETS=false0 B
140RUN |5 MODE=stable JAX_VERSION= DEVICE=gpu32 B
141ENV PYTHONPATH=/deps/src0 B
142ARG PACKAGE_DIR0 B
143ENV PACKAGE_DIR=src0 B
144ENV MAXTEXT_ASSETS_ROOT=/deps/src/maxtext/assets0 B
145ENV MAXTEXT_TEST_ASSETS_ROOT=/deps/tests/assets0 B
146ENV MAXTEXT_PKG_DIR=/deps/src/maxtext0 B
147ENV MAXTEXT_REPO_ROOT=/deps0 B
148WORKDIR /deps32 B
149COPY src/maxtext/assets/ /deps/src/maxtext/assets # buildkit32 B
150COPY src/maxtext/ src/maxtext/ # buildkit32 B

Command

 ARG RELEASE