Sign inSign up

rocm/jax-training:maxtext-v25.7-jax060

Manifest digest

sha256:7352212ae033a76dca2b9dceffc23c1b5f1a61a7a560082cf747a9bf1acfc9ce

OS/ARCH

linux/amd64

Compressed size

15.58 GB

Last pushed

11 months by rocmdev

Type

Image

Manifest digest

sha256:7352212ae033a76dca2b9dceffc23c1b5f1a61a7a560082cf747a9bf1acfc9ce

Image Layers

Image Layers

1ARG RELEASE0 B
2ARG LAUNCHPAD_BUILD_ARCH0 B
3LABEL org.opencontainers.image.ref.name=ubuntu0 B
4LABEL org.opencontainers.image.version=22.040 B
5ADD file ... in / 28.16 MB
6CMD ["/bin/bash"]0 B
7LABEL [email protected]0 B
8ARG ROCM_VERSION=5.30 B
9ARG AMDGPU_VERSION=5.30 B
10ARG APT_PREF0 B
11RUN |3 ROCM_VERSION=6.4.1 AMDGPU_VERSION=6.4.1 APT_PREF=Package:240 B
12RUN |3 ROCM_VERSION=6.4.1 AMDGPU_VERSION=6.4.1 APT_PREF=Package:4 GB
13RUN |3 ROCM_VERSION=6.4.1 AMDGPU_VERSION=6.4.1 APT_PREF=Package:657 B
14WORKDIR /opt32 B
15ARG DEBIAN_FRONTEND=noninteractive0 B
16ARG ROCM_VERSION=6.4.10 B
17ENV GPU_ARCHS=gfx9420 B
18ENV ROCM_PATH=/opt/rocm MPI_HOME=/opt/ompi UCX_HOME=/opt/ucx PATH=/opt/ompi/bin:/opt/ucx/bin:/opt/cache/bin:/opt/rocm/llvm/bin:/opt/rocm/opencl/bin:/opt/rocm/hip/bin:/opt/rocm/hcc/bin:/opt/rocm/bin:/opt/conda/envs/py_3.10/bin:/opt/conda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
19RUN |2 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 /bin/sh160.28 MB
20RUN |2 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 /bin/sh41.7 MB
21WORKDIR /opt/rocm/32 B
22RUN |2 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 /bin/sh5.37 KB
23ARG HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git0 B
24ARG HIPBLASLT_COMMIT=499ece1c21a504b6c97197e82c2ff2d35b8667010 B
25RUN |4 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git2.21 GB
26RUN |4 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git188 B
27COPY /opt/MIOpen/build/*.deb build/ # buildkit211.71 MB
28RUN |4 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git220.55 KB
29RUN |4 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git296.67 MB
30RUN |4 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git243 B
31ENV RCCL_COMMIT=3c9c22bb5221377057c1943310b786ddab6e32840 B
32RUN |4 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git220.3 KB
33RUN |4 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git10.59 MB
34ENV JAX_BRANCH=rocm-jaxlib-v0.6.00 B
35ENV XLA_BRANCH=rocm-jaxlib-v0.6.00 B
36ENV JAX_COMMIT=3d6b521d8b32bf3de9300c1bd3629f858c081b150 B
37ENV XLA_COMMIT=510ea0628be3d3259def59f5a7b9a018119574740 B
38WORKDIR /opt32 B
39RUN |4 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git301.18 MB
40RUN |4 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git2.6 GB
41ENV NVTE_FRAMEWORK=jax0 B
42ENV NVTE_ROCM_ARCH=gfx9420 B
43ENV NVTE_USE_ROCM=10 B
44ARG TE_COMMIT=90d703dd4fca619be73947f39b522307f23d3bd00 B
45RUN |5 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git151 B
46RUN |5 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git3.34 GB
47WORKDIR /workspace99 B
48COPY google_forced_balance.patch /workspace/google_forced_balance.patch # buildkit871 B
49COPY device-info.patch /workspace/device-info.patch # buildkit1.68 KB
50COPY max_to_keep.patch /workspace/max_to_keep.patch # buildkit847 B
51COPY attention_mask.patch /workspace/attention_mask.patch # buildkit481 B
52RUN |5 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git2.06 GB
53RUN |5 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git1.38 KB
54RUN |5 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git154 B
55RUN |5 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git146 B
56RUN |5 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git146 B
57RUN |5 DEBIAN_FRONTEND=noninteractive ROCM_VERSION=6.4.1 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git148 B
58ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B
59ENV NVTE_USE_HIPBLASLT=10 B
60ENV GPU_MAX_HW_QUEUES=20 B
61ENV HIP_FORCE_DEV_KERNARG=10 B
62ENV HSA_FORCE_FINE_GRAIN_PCIE=10 B
63ENV NVTE_FUSED_ATTN=10 B
64ENV NCCL_DEBUG=VERSION0 B
65ENV NVTE_CK_USES_BWD_V3=10 B
66ENV NVTE_CK_USES_FWD_V3=00 B
67ENV NVTE_CK_IS_V3_ATOMIC_FP32=00 B
68ENV NVTE_CK_HOW_V3_BF16_CVT=20 B
69ENV NVTE_FUSED_ATTN_CK=10 B
70ENV NVTE_FUSED_ATTN_AOTRITON=00 B
71WORKDIR /workspace32 B
72COPY google_forced_balance.patch /workspace/google_forced_balance.patch # buildkit1.08 KB
73COPY device-info.patch /workspace/device-info.patch # buildkit1.67 KB
74COPY max_to_keep.patch /workspace/max_to_keep.patch # buildkit843 B
75COPY attention_mask.patch /workspace/attention_mask.patch # buildkit476 B
76COPY quantization.patch /workspace/quantization.patch # buildkit443 B
77RUN /bin/sh -c rm -rf342.27 MB
78RUN /bin/sh -c pip install1.38 KB
79RUN /bin/sh -c rm /workspace/google_forced_balance.patch153 B
80RUN /bin/sh -c rm /workspace/device-info.patch146 B
81RUN /bin/sh -c rm /workspace/max_to_keep.patch146 B
82RUN /bin/sh -c rm /workspace/attention_mask.patch147 B
83RUN /bin/sh -c rm /workspace/quantization.patch147 B

Command

 ARG RELEASE