rocm/jax-training:maxtext-v26.2

Manifest digest

sha256:a89643388487b1e2fc6b6ef7bd3c44378c05d217309c977a1c18c72d05ebcaeb

OS/ARCH

linux/amd64

Compressed size

22.12 GB

Last pushed

6 months by rocmdev

Type

Image

Manifest digest

sha256:a89643388487b1e2fc6b6ef7bd3c44378c05d217309c977a1c18c72d05ebcaeb

Image Layers

1ARG RELEASE0 B
2ARG LAUNCHPAD_BUILD_ARCH0 B
3LABEL org.opencontainers.image.ref.name=ubuntu0 B
4LABEL org.opencontainers.image.version=24.040 B
5ADD file ... in / 29.18 MB
6CMD ["/bin/bash"]0 B
7RUN /bin/sh -c apt-get update176.23 MB
8RUN /bin/sh -c apt-get update1.63 MB
9ARG GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a gfx9420 B
10ENV GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a gfx9420 B
11ARG ROCM_VERSION=7.1.10 B
12ARG ROCM_PATH=/opt/rocm-7.1.10 B
13ENV ROCM_PATH=/opt/rocm-7.1.10 B
14ARG ROCM_BUILD_JOB=0 B
15ARG ROCM_BUILD_NUM=0 B
16ARG PLUGIN_NAMESPACE=70 B
17ARG THEROCK_PATH0 B
18RUN |7 GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a6.06 GB
19ENV PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin:/opt/rocm/llvm/bin0 B
20ENV ln=-s /opt/rocm-/lib/librocblas.so.5 /opt/rocm-/lib/librocblas.so.40 B
21ENV LD_LIBRARY_PATH=/opt/rocm-7.1.1/lib:0 B
22ENV LLVM_PATH=/opt/rocm-7.1.1/llvm0 B
23ENV HCC_HOME=/opt/rocm-7.1.1/hcc0 B
24ENV HIP_PATH=/opt/rocm-7.1.1/0 B
25ENV OPENCL_ROOT=/opt/rocm-7.1.1/opencl0 B
26ENV PATH=/opt/rocm-7.1.1/hcc/bin:/opt/rocm-7.1.1//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin:/opt/rocm/llvm/bin0 B
27ENV PATH=/opt/rocm-7.1.1/bin:/opt/rocm-7.1.1/hcc/bin:/opt/rocm-7.1.1//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin:/opt/rocm/llvm/bin0 B
28ENV PATH=/opt/rocm-7.1.1/opencl/bin:/opt/rocm-7.1.1/bin:/opt/rocm-7.1.1/hcc/bin:/opt/rocm-7.1.1//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin:/opt/rocm/llvm/bin0 B
29ENV PATH=/root/bin:/root/.local/bin:/opt/rocm-7.1.1/opencl/bin:/opt/rocm-7.1.1/bin:/opt/rocm-7.1.1/hcc/bin:/opt/rocm-7.1.1//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin:/opt/rocm/llvm/bin0 B
30RUN |7 GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a98.46 MB
31RUN |7 GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a92.33 MB
32ARG JAX_VERSION0 B
33ARG XLA_COMMIT=4178517b7d7e5e3318ea4b701b01dc06847b73b00 B
34ARG JAX_COMMIT=5712de44e97c455faed1fd45532e821ca66d025a0 B
35ARG ROCM_JAX_COMMIT=be6011b1b9d3e61e103fcf8700f54f9742d05fb50 B
36LABEL com.amdgpu.rocm_version=7.1.1 com.amdgpu.python_version=3.12 com.amdgpu.jax_version= com.amdgpu.jax_commit=5712de44e97c455faed1fd45532e821ca66d025a0 B
37RUN |11 GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a173.23 MB
38ENV ROCM_PATH=/opt/rocm-7.1.10 B
39ENV PATH=/opt/rocm-7.1.1/opencl/bin:/opt/rocm-7.1.1/bin:/opt/rocm-7.1.1/hcc/bin:/opt/rocm-7.1.1//bin:/root/bin:/root/.local/bin:/opt/rocm-7.1.1/opencl/bin:/opt/rocm-7.1.1/bin:/opt/rocm-7.1.1/hcc/bin:/opt/rocm-7.1.1//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin:/opt/rocm/llvm/bin0 B
40ENV MAX_JOBS=1280 B
41RUN /bin/sh -c apt update1.63 GB
42RUN /bin/sh -c python -m4 MB
43ENV PATH=/opt/venv/bin:/opt/rocm-7.1.1/opencl/bin:/opt/rocm-7.1.1/bin:/opt/rocm-7.1.1/hcc/bin:/opt/rocm-7.1.1//bin:/root/bin:/root/.local/bin:/opt/rocm-7.1.1/opencl/bin:/opt/rocm-7.1.1/bin:/opt/rocm-7.1.1/hcc/bin:/opt/rocm-7.1.1//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin:/opt/rocm/llvm/bin0 B
44RUN /bin/sh -c pip install58.12 MB
45WORKDIR /workspace99 B
46ARG ROCR_REPO=https://github.com/ROCm/rocm-systems.git0 B
47ARG ROCR_BRANCH=release/rocm-rel-7.1.1.10 B
48ARG ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac669230 B
49RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923176 B
50RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923188 B
51RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac669231.44 MB
52RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923206 B
53RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923194 B
54RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923233 B
55RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac669234.85 KB
56RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac6692317.83 KB
57RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923217.23 MB
58RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac669232.38 MB
59WORKDIR /workspace32 B
60RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923368.87 KB
61ARG ROCM_REPO=https://github.com/ROCm/rocm-systems.git0 B
62ARG ROCM_BRANCH=release/rocm-rel-7.10 B
63ENV CLR_DIR=/workspace/rocm-systems/projects/clr0 B
64ENV HIP_DIR=/workspace/rocm-systems/projects/hip0 B
65RUN |5 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac6692313.9 MB
66ARG ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git0 B
67ARG ROCPROF_BRANCH=ed42157c3145a2f10f3c95c6059efd6f367064190 B
68RUN |7 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac6692314.06 MB
69ARG HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git0 B
70ENV HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git0 B
71ARG HIPBLASLT_BRANCH=fe818bf7619575499cb72225ced3e255aa08cbb50 B
72ENV HIPBLASLT_BRANCH=fe818bf7619575499cb72225ced3e255aa08cbb50 B
73ENV PYTORCH_ROCM_ARCH=gfx942;gfx9500 B
74ENV ROCM_AMDGPU_TARGETS=gfx942,gfx9500 B
75ENV Python_ROOT=/pyenv0 B
76RUN |9 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923176 B
77RUN |9 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac669231.54 GB
78RUN |9 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923187 B
79RUN |9 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923172 B
80RUN |9 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac6692394.81 KB
81RUN |9 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac6692394.81 KB
82RUN |9 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923162 B
83WORKDIR /workspace32 B
84ARG MAXTEXT_REPO=https://github.com/ROCm/maxtext.git0 B
85ARG MAXTEXT_BRANCH=yeandy/update-patches-scaling-patch-v2-checkpoint-restore0 B
86RUN |11 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac669231.65 GB
87ENV ROCM_JAX_REPO=https://github.com/ROCm/rocm-jax.git0 B
88ENV ROCM_JAX_BRANCH=6251484e78c34751b2098cc1026adb79d4d51fad0 B
89ARG JAX_VERSION=0.8.20 B
90ENV JAX_VERSION=0.8.20 B
91ENV JAX_COMMIT=cef02f3ae4abf69862294cca9370c721eb7eb2b40 B
92ENV XLA_COMMIT=e301b64cf5ca2f4d356fcb825b25feeee31897330 B
93ENV XLA_BRANCH=rocm-jaxlib-v0.8.20 B
94ENV JAX_BRANCH=rocm-jaxlib-v0.8.20 B
95WORKDIR /opt32 B
96RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923356.93 MB
97RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923170.24 MB
98RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac6692314.29 MB
99RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923165.49 MB
100RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923976.63 KB
101RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac669235.59 GB
102WORKDIR /opt/jax32 B
103RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac669232.12 GB
104ARG TE_REPO=https://github.com/ROCm/TransformerEngine.git0 B
105ARG TE_BRANCH=aec00a7f7ce1273ad1366dbf6d7074c9904999940 B
106WORKDIR /workspace32 B
107ENV USE_ROCM=10 B
108ENV HIP_PATH=/opt/rocm0 B
109ENV NVTE_FRAMEWORK=jax0 B
110ENV CMAKE_BUILD_PARALLEL_LEVEL=640 B
111ENV NVTE_ROCM_ARCH=gfx942;gfx9500 B
112ENV NVTE_USE_ROCM=10 B
113ENV PYTHONPATH=/home/TransformerEngine/3rdparty/hipify_torch0 B
114ENV NVTE_AITER_PREBUILT_BASE_URL=https://compute-artifactory.amd.com:5000/artifactory/rocm-generic-local/te-ci/aiter-prebuilts0 B
115RUN |14 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923571.94 KB
116RUN |14 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923872.13 MB
117ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B
118ENV NVTE_USE_HIPBLASLT=10 B
119ENV GPU_MAX_HW_QUEUES=20 B
120ENV HIP_FORCE_DEV_KERNARG=10 B
121ENV HSA_FORCE_FINE_GRAIN_PCIE=10 B
122ENV NVTE_FUSED_ATTN=10 B
123ENV NCCL_DEBUG=VERSION0 B
124ENV NVTE_CK_USES_BWD_V3=10 B
125ENV NVTE_CK_USES_FWD_V3=10 B
126ENV NVTE_CK_IS_V3_ATOMIC_FP32=00 B
127ENV NVTE_CK_HOW_V3_BF16_CVT=20 B
128ENV XLA_PYTHON_CLIENT_MEM_FRACTION=.970 B
129ENV 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
130RUN |14 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac669231.55 MB
131RUN |14 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923949 B
132WORKDIR /opt/jax32 B
133RUN |14 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac6692389.44 MB
134WORKDIR /workspace32 B
135RUN |14 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac6692368.61 MB
136ARG AINIC_BUNDLE_VERSION=1.117.5-a-380 B
137RUN |15 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac669239.92 MB
138RUN |15 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923227 B
139ARG UCX_VERSION=1.18.00 B
140RUN |16 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac6692386.44 MB
141ENV UCX_INSTALL_DIR=/workspace/ucx-1.18.0/install0 B
142ARG MPI_VERSION=4.1.60 B
143RUN |17 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac669233.84 MB
144ARG RCCL_REPO=https://github.com/ROCm/rccl.git0 B
145ARG RCCL_BRANCH=rocm-7.1.10 B
146RUN |19 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923302.74 MB
147ARG ANP_REPO=https://github.com/rocm/amd-anp.git0 B
148ARG ANP_BRANCH=tags/v1.3.00 B
149ENV RCCL_HOME=/workspace/rccl0 B
150ENV MPI_PATH=/ompi-4.1.6/install0 B
151RUN |21 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923107.31 MB
152RUN |21 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923516.17 KB
153RUN |21 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923131 B
154RUN |21 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923126 B
155RUN |21 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923126 B
156RUN |21 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923128 B
157ARG GIT_COMMIT_TAG=e5be0ef7ea8210e52a702430bc62fb90fd5a73880 B
158RUN |22 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923185 B
159bash351.24 KB
160bash31.55 MB
161bash444.98 MB
162bash0 B
163bash0 B
164bash0 B
165bash188 B
166bash260 B

Command

 ARG RELEASE