Sign inSign up

rocm/jax-training:maxtext-v26.3.1

Manifest digest

sha256:75782ecbc38d8036f371abf50edd5f22f50b6b1637c6d5f68703423edb49f3bc

OS/ARCH

linux/amd64

Compressed size

16.66 GB

Last pushed

2 months by rocmdev

Type

Image

Manifest digest

sha256:75782ecbc38d8036f371abf50edd5f22f50b6b1637c6d5f68703423edb49f3bc

Image Layers

Image Layers

1ARG RELEASE0 B
2ARG LAUNCHPAD_BUILD_ARCH0 B
3LABEL org.opencontainers.image.version=24.040 B
4ADD file ... in / 28.36 MB
5CMD ["/bin/bash"]0 B
6ENV DEBIAN_FRONTEND=noninteractive0 B
7RUN /bin/sh -c apt update500.3 MB
8RUN /bin/sh -c wget https://repo.radeon.com/amdgpu-install/7.2.1/ubuntu/noble/amdgpu-install_7.2.1.70201-1_all.deb7.41 GB
9ENV ROCM_PATH=/opt/rocm-7.2.10 B
10ENV HIP_PATH=/opt/rocm-7.2.10 B
11ENV MAX_JOBS=1280 B
12ENV PATH=/opt/rocm-7.2.1/opencl/bin:/opt/rocm-7.2.1/bin:/opt/rocm-7.2.1/hcc/bin:/opt/rocm-7.2.1//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
13ENV LD_LIBRARY_PATH=/opt/rocm/lib0 B
14RUN /bin/sh -c apt update1.64 GB
15RUN /bin/sh -c python -m4 MB
16ENV PATH=/opt/venv/bin:/opt/rocm-7.2.1/opencl/bin:/opt/rocm-7.2.1/bin:/opt/rocm-7.2.1/hcc/bin:/opt/rocm-7.2.1//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
17RUN /bin/sh -c pip install29.85 MB
18WORKDIR /workspace99 B
19ARG ROCR_REPO=https://github.com/ROCm/rocm-systems.git0 B
20ARG ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f29810578040 B
21ARG ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b100 B
22RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10173 B
23RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10192 B
24RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b101.59 MB
25RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10217 B
26RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b101.27 GB
27RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b101.26 MB
28WORKDIR /workspace32 B
29ARG HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git0 B
30ENV HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git0 B
31ARG HIPBLASLT_BRANCH=bfcf25fa18b2a0197e5b18b4b8b5f81f13ba57810 B
32ENV HIPBLASLT_BRANCH=bfcf25fa18b2a0197e5b18b4b8b5f81f13ba57810 B
33ENV PYTORCH_ROCM_ARCH=gfx942;gfx9500 B
34ENV ROCM_AMDGPU_TARGETS=gfx942,gfx9500 B
35ENV Python_ROOT=/pyenv0 B
36RUN |5 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10176 B
37RUN |5 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b101.68 GB
38RUN |5 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10185 B
39RUN |5 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10203.99 KB
40WORKDIR /workspace32 B
41ARG MAXTEXT_REPO=https://github.com/ROCm/maxtext.git0 B
42ARG MAXTEXT_BRANCH=release/v26.30 B
43RUN |7 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b101.6 GB
44ENV ROCM_JAX_REPO=https://github.com/ROCm/rocm-jax.git0 B
45ENV ROCM_JAX_BRANCH=aae4d33f0c5aa711e8502ca4578ead71979d88180 B
46ARG JAX_VERSION=0.8.20 B
47ENV JAX_VERSION=0.8.20 B
48ENV JAX_COMMIT=4b67574ad93ee1c42dcf4498bff0830760a7777e0 B
49ENV XLA_COMMIT=107e89fa589603c5c4a0c60ff333130f0ee1b5c80 B
50ENV XLA_BRANCH=rocm-jaxlib-v0.8.20 B
51ENV JAX_BRANCH=rocm-jaxlib-v0.8.20 B
52WORKDIR /opt32 B
53RUN |8 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b1086.77 MB
54RUN |8 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10323.31 MB
55ARG TE_REPO=https://github.com/ROCm/TransformerEngine.git0 B
56ARG TE_BRANCH=9b312832f4a21bd5b73a6087de6bd1eeaf219b820 B
57WORKDIR /workspace32 B
58ENV USE_ROCM=10 B
59ENV HIP_PATH=/opt/rocm0 B
60ENV NVTE_FRAMEWORK=jax0 B
61ENV CMAKE_BUILD_PARALLEL_LEVEL=640 B
62ENV NVTE_ROCM_ARCH=gfx942;gfx9500 B
63ENV NVTE_USE_ROCM=10 B
64ENV PYTHONPATH=/home/TransformerEngine/3rdparty/hipify_torch0 B
65RUN |10 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10650.14 KB
66RUN |10 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10830.02 MB
67ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B
68ENV NVTE_USE_HIPBLASLT=10 B
69ENV GPU_MAX_HW_QUEUES=20 B
70ENV HIP_FORCE_DEV_KERNARG=10 B
71ENV HSA_FORCE_FINE_GRAIN_PCIE=10 B
72ENV NVTE_FUSED_ATTN=10 B
73ENV NCCL_DEBUG=VERSION0 B
74ENV NVTE_CK_USES_BWD_V3=10 B
75ENV NVTE_CK_USES_FWD_V3=10 B
76ENV NVTE_CK_IS_V3_ATOMIC_FP32=10 B
77ENV NVTE_CK_HOW_V3_BF16_CVT=20 B
78ENV XLA_PYTHON_CLIENT_MEM_FRACTION=.970 B
79ENV 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
80RUN |10 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10141.91 MB
81WORKDIR /workspace32 B
82ARG PRIMUS_REPO=https://github.com/AMD-AGI/Primus.git0 B
83ARG PRIMUS_BRANCH=main0 B
84RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10846.27 MB
85RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b101.55 MB
86ARG AINIC_BUNDLE_VERSION=1.117.5-a-770 B
87RUN |13 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10713.38 KB
88RUN |13 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10230 B
89ARG UCX_VERSION=1.18.00 B
90RUN |14 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b1087.65 MB
91ENV UCX_INSTALL_DIR=/workspace/ucx-1.18.0/install0 B
92ARG MPI_VERSION=4.1.60 B
93RUN |15 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b103.85 MB
94ARG RCCL_REPO=https://github.com/ROCm/rccl.git0 B
95ARG RCCL_BRANCH=release/rocm-rel-7.20 B
96RUN |17 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10241.66 MB
97ARG ANP_REPO=https://github.com/rocm/amd-anp.git0 B
98ARG ANP_BRANCH=tags/v1.3.00 B
99ENV RCCL_HOME=/workspace/rccl0 B
100ENV MPI_PATH=/ompi-4.1.6/install0 B
101RUN |19 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b1032 B
102RUN |19 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10564.43 KB
103RUN |19 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b10446 B
104RUN |19 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b1032 B
105ARG GIT_COMMIT_TAG=789c3cb481d5420e0bcd7ff2da5362ef83b61ca00 B
106ARG DOCKERFILE_PATH=jax_training/docker/Dockerfile_jax_0dot80 B
107RUN |21 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=e1a6bc5663304b9c586b3254b8920f2981057804 ROCR_CP_COMMIT=559d48b1f013a2a8e9decd2557508de7ac6c6b1032.03 KB
108COPY jax_training/docker/Dockerfile_jax_0dot8 /workspace/.manifest/Dockerfile # buildkit6.1 KB
109ENV NCCL_DMABUF_ENABLE=10 B

Command

 ARG RELEASE