Sign inSign up

sanjeevrg/maxtext-gpu-pretraining:stable-493fba6

Manifest digest

sha256:8221d350c55e600b385db297c7fa9c3f154d14bbd6f3cdd3fa0e38c06290c6da

OS/ARCH

linux/amd64

Compressed size

9.99 GB

Last pushed

4 months by sanjeevrg

Type

Image

Manifest digest

sha256:8221d350c55e600b385db297c7fa9c3f154d14bbd6f3cdd3fa0e38c06290c6da

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.46.00 B
18ARG AWS_OFI_NCCL_VERSION=1.17.30 B
19ENV GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.26 MOFED_VERSION=5.4-rdmacore59.1 OPENUCX_VERSION=1.20.00 B
20ARG TARGETARCH=amd640 B
21RUN |10 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.26196.53 MB
22ENV 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
23ENV OMPI_MCA_coll_hcoll_enable=00 B
24ARG CUDA_VERSION=13.2.1.0090 B
25ARG CUDA_DRIVER_VERSION=595.58.030 B
26ARG NVVM_VERSION=13.2.780 B
27ARG DOCA_VERSION=3.2.00 B
28ENV CUDA_VERSION=13.2.1.009 CUDA_DRIVER_VERSION=595.58.03 NVVM_VERSION=13.2.78 DOCA_VERSION=3.2.00 B
29RUN |14 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.26154.59 MB
30RUN |14 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.2614.37 KB
31ENV _CUDA_COMPAT_PATH=/usr/local/cuda/compat ENV=/etc/shinit_v2 BASH_ENV=/etc/bash.bashrc SHELL=/bin/bash0 B
32LABEL com.nvidia.volumes.needed=nvidia_driver com.nvidia.cuda.version=9.00 B
33ARG NCCL_VERSION=2.29.7 CUBLAS_VERSION=13.4.0.1 CUFFT_VERSION=12.2.0.46 CURAND_VERSION=10.4.2.550 B
34ENV NCCL_VERSION=2.29.7 CUBLAS_VERSION=13.4.0.1 CUFFT_VERSION=12.2.0.46 CURAND_VERSION=10.4.2.550 B
35ADD docs.tgz / # buildkit11.26 KB
36RUN |47 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.26179 B
37ARG _LIBPATH_SUFFIX=0 B
38ENV 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
39RUN |48 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.2692.38 MB
40COPY /opt/amazon/aws-ofi-nccl /opt/amazon/aws-ofi-nccl # buildkit3.16 MB
41RUN |48 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.268.43 KB
42COPY entrypoint/ /opt/nvidia/ # buildkit5.98 KB
43ENV NVIDIA_PRODUCT_NAME=CUDA0 B
44ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"]0 B
45RUN |48 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.26408 B
46RUN |48 ENABLE_MITMPROXY=0 GDRCOPY_VERSION=2.5.1 HPCX_VERSION=2.2632 B
47ARG ENABLE_MITMPROXY=00 B
48RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c32 B
49RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c110.03 MB
50ARG CUDA_COMPONENT_LIST=cccl crt nvrtc driver-dev0 B
51ENV CUDA_COMPONENT_LIST=cccl crt nvrtc driver-dev0 B
52RUN |2 ENABLE_MITMPROXY=0 CUDA_COMPONENT_LIST=cccl crt214.88 MB
53ARG VERSION_LIST=CURAND_VERSION CUBLAS_VERSION CUFILE_VERSION0 B
54RUN |3 ENABLE_MITMPROXY=0 CUDA_COMPONENT_LIST=cccl crt1.65 GB
55ENV LIBRARY_PATH=/usr/local/cuda/lib64/stubs:0 B
56RUN |3 ENABLE_MITMPROXY=0 CUDA_COMPONENT_LIST=cccl crt32 B
57LABEL com.nvidia.nccl.version=2.29.7 com.nvidia.cublas.version=13.4.0.1 com.nvidia.cufft.version=12.2.0.46 com.nvidia.curand.version=10.4.2.550 B
58ARG ENABLE_MITMPROXY=00 B
59RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c32 B
60RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c32 B
61RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c64.13 MB
62RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c2.51 GB
63LABEL com.nvidia.nccl.version=2.29.7 com.nvidia.cublas.version=13.4.0.1 com.nvidia.cufft.version=12.2.0.46 com.nvidia.curand.version=10.4.2.550 B
64RUN |1 ENABLE_MITMPROXY=0 /bin/sh -c32 B
65ARG TARGETARCH=amd640 B
66ARG ENABLE_FIPS=00 B
67ARG ENABLE_MITMPROXY=00 B
68RUN |3 TARGETARCH=amd64 ENABLE_FIPS=0 ENABLE_MITMPROXY=032 B
69RUN |3 TARGETARCH=amd64 ENABLE_FIPS=0 ENABLE_MITMPROXY=021.25 MB
70COPY cuda-*.patch /tmp # buildkit503 B
71RUN |3 TARGETARCH=amd64 ENABLE_FIPS=0 ENABLE_MITMPROXY=0345.2 MB
72ENV LIBRARY_PATH=/usr/local/cuda/lib64/stubs:/usr/local/cuda/lib64/stubs:0 B
73RUN |3 TARGETARCH=amd64 ENABLE_FIPS=0 ENABLE_MITMPROXY=032 B
74ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.040 B
75ARG [email protected]0 B
76ARG GIT_USER_NAME=JAX Toolbox0 B
77ARG JAX_TOOLBOX_REF=9f281cabf200716ed0660b23cdf21b2b064ee11d0 B
78ARG NSIGHT_SYSTEMS_VERSION_OVERRIDE0 B
79ENV CUDA_BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.040 B
80ENV JAX_TOOLBOX_REF=9f281cabf200716ed0660b23cdf21b2b064ee11d0 B
81RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX52.63 MB
82RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX187 B
83RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX143 B
84ADD --chmod=777 git-clone.sh pip-finalize.sh /usr/local/bin/2.67 KB
85RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX5.28 MB
86RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX83.65 MB
87ADD pip-vcs-equivalency.patch /opt/pip/ # buildkit1.55 KB
88RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX71.63 MB
89ENV PIP_BREAK_SYSTEM_PACKAGES=10 B
90RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX14.59 KB
91RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX4.88 MB
92RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX172 B
93ADD symlnk-cudnn.sh /usr/local/bin # buildkit1 KB
94RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX890 B
95ADD symlnk-nccl.sh /usr/local/bin # buildkit717 B
96RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX1.49 KB
97ADD install-nccl-sanity-check.sh /usr/local/bin # buildkit600 B
98ADD nccl-sanity-check.cu /opt # buildkit2.13 KB
99RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX9.37 KB
100ADD parallel-launch /usr/local/bin/ # buildkit535 B
101ADD install-nsight-systems.sh /usr/local/bin # buildkit715 B
102RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX32 B
103RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX1.8 KB
104ADD symlnk-nvshmem.sh /usr/local/bin # buildkit686 B
105RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX235 B
106ENV MANIFEST_FILE=/opt/manifest.d/manifest.yaml0 B
107ADD manifest.yaml create-distribution.sh bump.sh /opt/manifest.d/6.98 KB
108COPY patches/ /opt/manifest.d/patches/ # buildkit4.76 KB
109RUN /bin/sh -c if [559 B
110RUN /bin/sh -c apt-get update43.93 MB
111RUN /bin/sh -c apt-get install2.09 MB
112RUN /bin/sh -c echo "deb263 B
113RUN /bin/sh -c curl https://packages.cloud.google.com/apt/doc/apt-key.gpg943 B
114RUN /bin/sh -c apt-get update155.25 MB
115ENV 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
116RUN /bin/sh -c apt-get update32 B
117ARG MODE0 B
118ENV ENV_MODE=stable0 B
119ARG JAX_VERSION0 B
120ENV ENV_JAX_VERSION=0 B
121ARG DEVICE0 B
122ENV ENV_DEVICE=gpu0 B
123ARG PACKAGE_DIR0 B
124ENV PACKAGE_DIR=src0 B
125ENV MAXTEXT_ASSETS_ROOT=/deps/src/maxtext/assets0 B
126ENV MAXTEXT_TEST_ASSETS_ROOT=/deps/tests/assets0 B
127ENV MAXTEXT_PKG_DIR=/deps/src/maxtext0 B
128ENV MAXTEXT_REPO_ROOT=/deps0 B
129WORKDIR /deps94 B
130COPY src/dependencies/extra_deps/ src/dependencies/extra_deps/ # buildkit1.05 KB
131COPY src/dependencies/requirements/ src/dependencies/requirements/ # buildkit5.87 KB
132COPY src/dependencies/scripts/ src/dependencies/scripts/ # buildkit10.85 KB
133COPY src/maxtext/integration/vllm/ src/maxtext/integration/vllm/ # buildkit18.09 KB
134RUN |4 MODE=stable JAX_VERSION= DEVICE=gpu32 B
135RUN |4 MODE=stable JAX_VERSION= DEVICE=gpu4.19 GB
136COPY src/maxtext/ src/maxtext/ # buildkit9.6 MB
137COPY tests*/ tests/ # buildkit21.54 MB
138COPY pytest.ini* pytest.ini # buildkit869 B
139COPY benchmarks*/ benchmarks/ # buildkit279.6 KB
140ARG INCLUDE_TEST_ASSETS=false0 B
141RUN |5 MODE=stable JAX_VERSION= DEVICE=gpu32 B
142ENV PYTHONPATH=/deps/src0 B
143ARG PACKAGE_DIR0 B
144ENV PACKAGE_DIR=src0 B
145ENV MAXTEXT_ASSETS_ROOT=/deps/src/maxtext/assets0 B
146ENV MAXTEXT_TEST_ASSETS_ROOT=/deps/tests/assets0 B
147ENV MAXTEXT_PKG_DIR=/deps/src/maxtext0 B
148ENV MAXTEXT_REPO_ROOT=/deps0 B
149WORKDIR /deps32 B
150COPY src/maxtext/assets/ /deps/src/maxtext/assets # buildkit32 B
151COPY src/maxtext/ src/maxtext/ # buildkit32 B

Command

 ARG RELEASE