Sign inSign up

rocm/jax-training:maxtext-v26.9

Manifest digest

sha256:afbcbaf0414b689c8d52efe05c2ed386fe5b9191633b4101c1d39791905def03

OS/ARCH

linux/amd64

Compressed size

13.5 GB

Last pushed

10 days by rocm

Type

Image

Manifest digest

sha256:afbcbaf0414b689c8d52efe05c2ed386fe5b9191633b4101c1d39791905def03

Image Layers

Image Layers

1ARG RELEASE0 B
2ARG LAUNCHPAD_BUILD_ARCH0 B
3LABEL org.opencontainers.image.version=24.040 B
4ADD file ... in / 29.21 MB
5CMD ["/bin/bash"]0 B
6WORKDIR /workspace/99 B
7ENV MAX_JOBS=1280 B
8ENV PYTORCH_ROCM_ARCH=gfx942;gfx9500 B
9ENV ROCM_AMDGPU_TARGETS=gfx942,gfx9500 B
10ENV ROCPROFILER_QUEUE_INTERPOSITION=00 B
11ENV DEBUG_HIP_DYNAMIC_QUEUES=00 B
12ENV DEBIAN_FRONTEND=noninteractive0 B
13ARG LIBC_SECURITY_REFRESH=2026-10-010 B
14RUN |1 LIBC_SECURITY_REFRESH=2026-10-01 /bin/sh -c234.65 MB
15COPY /uv /usr/local/bin/uv # buildkit20.96 MB
16RUN |1 LIBC_SECURITY_REFRESH=2026-10-01 /bin/sh -c9.65 KB
17ENV PATH=/opt/venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
18ENV VIRTUAL_ENV=/opt/venv0 B
19COPY /dist/pip-*.whl /tmp/ # buildkit1.71 MB
20RUN |1 LIBC_SECURITY_REFRESH=2026-10-01 /bin/sh -c29.32 MB
21ENV THE_ROCK_VERSION=10.3.0a202610260 B
22RUN |1 LIBC_SECURITY_REFRESH=2026-10-01 /bin/sh -c4.09 GB
23RUN |1 LIBC_SECURITY_REFRESH=2026-10-01 /bin/sh -c4.09 GB
24COPY /bazelisk /usr/local/bin/bazel # buildkit5.68 MB
25ENV ROCM_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel0 B
26ENV ROCM_HOME=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel0 B
27ENV HIP_PLATFORM=amd0 B
28ENV HIP_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel0 B
29ENV HIP_CLANG_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/llvm/bin0 B
30ENV HIP_INCLUDE_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/include0 B
31ENV HIP_LIB_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib0 B
32ENV HIP_DEVICE_LIB_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib/llvm/amdgcn/bitcode0 B
33ENV PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/bin:/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/llvm/bin:/opt/venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
34ENV LIBRARY_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib:/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib:/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib640 B
35ENV CPATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/include0 B
36ENV PKG_CONFIG_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib/pkgconfig0 B
37ENV LD_LIBRARY_PATH=0 B
38ARG JAX_VERSION=0.12.00 B
39ENV JAX_VERSION=0.12.00 B
40ENV JAX_PJRT_VERSION=0.12.0+rocm10.3.0a202610260 B
41ENV JAX_PLUGIN_VERSION=0.12.0+rocm10.3.0a202610260 B
42RUN |2 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 /bin/sh273 MB
43ARG TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb0 B
44RUN |3 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb401.23 MB
45COPY /mount.gcsfuse /usr/sbin/mount.gcsfuse # buildkit18.99 MB
46ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B
47ENV NVTE_USE_HIPBLASLT=10 B
48ENV GPU_MAX_HW_QUEUES=20 B
49ENV HIP_FORCE_DEV_KERNARG=10 B
50ENV HSA_FORCE_FINE_GRAIN_PCIE=10 B
51ENV NVTE_FUSED_ATTN=10 B
52ENV NCCL_DEBUG=VERSION0 B
53ENV NVTE_CK_USES_BWD_V3=10 B
54ENV NVTE_CK_USES_FWD_V3=10 B
55ENV NVTE_CK_IS_V3_ATOMIC_FP32=10 B
56ENV NVTE_CK_HOW_V3_BF16_CVT=20 B
57ENV XLA_PYTHON_CLIENT_MEM_FRACTION=.90 B
58ENV XLA_FLAGS=--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_all_gather_combine_threshold_bytes=85899345920 B
59WORKDIR /workspace32 B
60ARG MAXTEXT_REPO=https://github.com/ROCm/maxtext.git0 B
61ARG MAXTEXT_BRANCH=release/v26.90 B
62RUN |5 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb1.52 GB
63ENV MAXTEXT_PATH=/workspace/maxtext0 B
64ARG TF_REPO=https://github.com/ROCm/tensorflow-upstream.git0 B
65ARG TF_BRANCH=upstream-v2.21.00 B
66RUN |7 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb503.05 KB
67COPY /909041e4802c4b9a2223ca04099f35bf1dbbd460.tar.gz /opt/bazel_distdir/ # buildkit243.86 MB
68RUN |7 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb254.41 MB
69ARG PRIMUS_REPO=https://github.com/AMD-AGI/Primus.git0 B
70ARG PRIMUS_BRANCH=db1a41601833294231daadb3cdbb2933eb805c730 B
71RUN |9 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb1.6 GB
72RUN |9 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb32.14 MB
73COPY /mount.gcsfuse /usr/sbin/mount.gcsfuse # buildkit18.99 MB
74COPY /gcsfuse /usr/bin/gcsfuse # buildkit24.97 MB
75RUN |9 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb25.73 MB
76ARG AINIC_BUNDLE_VERSION=1.117.5-a-1960 B
77RUN |10 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb342.1 KB
78RUN |10 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb231 B
79ARG UCX_VERSION=1.18.00 B
80RUN |11 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb86.44 MB
81ENV UCX_INSTALL_DIR=/workspace/ucx-1.18.0/install0 B
82ARG MPI_VERSION=4.1.60 B
83RUN |12 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb3.84 MB
84RUN |12 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb587 B
85RUN |12 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb32 B
86ARG GIT_COMMIT_TAG=aa44dbf5a9a7edc0fc76694adc624023df6c50e70 B
87ARG DOCKERFILE_PATH=pipeline/jax_training/docker/Dockerfile0 B
88RUN |14 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb15.99 KB
89COPY pipeline/jax_training/docker/Dockerfile /workspace/.manifest/docker-build-recipe.txt # buildkit11.47 KB
90ENV NVTE_FRAMEWORK=jax0 B
91RUN |14 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb216.83 MB
92RUN |14 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb42.35 MB
93RUN |14 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb13.05 MB
94RUN |14 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb12.61 KB
95ARG MAXDIFFUSION_REPO=https://github.com/AI-Hypercomputer/maxdiffusion.git0 B
96ARG MAXDIFFUSION_COMMIT=68e069659f0af80694559e29939a17f879fe7f6a0 B
97RUN |16 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb267.06 MB
98RUN |16 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb3.18 KB
99RUN |16 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb13.36 KB
100RUN |16 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb6.26 KB
101RUN |16 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.12.0 TE_VERSION=2.19.0.dev0+rocm10.2.0a20261002.bfe94fb898 B

Command

 ARG RELEASE