sanjeevrg/maxtext-gpu-pretraining:stable-85690da

Manifest digest

sha256:bc00488f0843ede471d084207458da9ba447cbf3f3bbbeb58b8338bf166e7fa2

OS/ARCH

linux/amd64

Compressed size

9.98 GB

Last pushed

3 months by sanjeevrg

Type

Image

Manifest digest

sha256:bc00488f0843ede471d084207458da9ba447cbf3f3bbbeb58b8338bf166e7fa2

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

22

ENV 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/bin

0 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

38

ENV 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,video

0 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

57

LABEL 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.55

0 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

63

LABEL 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.55

0 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=0fde20f86429a863dabf208d9169abbc046388330 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=0fde20f86429a863dabf208d9169abbc046388330 B

81RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX52.73 MB

82RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX189 B

83RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX144 B

84ADD --chmod=777 git-clone.sh pip-finalize.sh /usr/local/bin/2.66 KB

85RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX5.29 MB

86RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX84.36 MB

87ADD pip-vcs-equivalency.patch /opt/pip/ # buildkit1.56 KB

88RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX72.31 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=JAX174 B

93ADD symlnk-cudnn.sh /usr/local/bin # buildkit1023 B

94RUN |5 BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04 [email protected] GIT_USER_NAME=JAX891 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 # buildkit685 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 [560 B

110RUN /bin/sh -c apt-get update36.26 MB

111RUN /bin/sh -c apt-get install2.09 MB

112RUN /bin/sh -c echo "deb264 B

113RUN /bin/sh -c curl https://packages.cloud.google.com/apt/doc/apt-key.gpg949 B

114RUN /bin/sh -c apt-get update156.04 MB

115

ENV 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/bin

0 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/ # buildkit6.21 KB

132COPY src/dependencies/scripts/ src/dependencies/scripts/ # buildkit11.93 KB

133COPY src/maxtext/integration/vllm/ src/maxtext/integration/vllm/ # buildkit22.38 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.68 MB

137COPY tests*/ tests/ # buildkit21.59 MB

138COPY pytest.ini* pytest.ini # buildkit876 B

139COPY benchmarks*/ benchmarks/ # buildkit279.61 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