rocm/jax-training:maxtext-v25.4

Manifest digest

sha256:fb3eb71cd74298a7b3044b7130cf84113f14d518ff05a2cd625c11ea5f6a7b01

OS/ARCH

linux/amd64

Compressed size

9.91 GB

Last pushed

over 1 year by rocmdev

Type

Image

Manifest digest

sha256:fb3eb71cd74298a7b3044b7130cf84113f14d518ff05a2cd625c11ea5f6a7b01

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.17 MB
6CMD ["/bin/bash"]0 B
7ARG ROCM_VERSION=6.3.10 B
8ARG AMDGPU_VERSION=6.3.603010 B
9ARG JAX_VERSION=0.4.310 B
10ARG PYTHON_VERSION=3.120 B
11RUN |4 ROCM_VERSION=6.3.3 AMDGPU_VERSION=6.3.60303 JAX_VERSION=0.4.31197.16 MB
12RUN |4 ROCM_VERSION=6.3.3 AMDGPU_VERSION=6.3.60303 JAX_VERSION=0.4.3116.63 KB
13RUN |4 ROCM_VERSION=6.3.3 AMDGPU_VERSION=6.3.60303 JAX_VERSION=0.4.311003.15 KB
14RUN |4 ROCM_VERSION=6.3.3 AMDGPU_VERSION=6.3.60303 JAX_VERSION=0.4.315.97 GB
15RUN |4 ROCM_VERSION=6.3.3 AMDGPU_VERSION=6.3.60303 JAX_VERSION=0.4.31577.91 MB
16ENV PATH=/opt/venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin/0 B
17SHELL [/bin/bash -c]0 B
18RUN /bin/bash -c apt update18.01 MB
19RUN /bin/bash -c apt install17.83 MB
20RUN /bin/bash -c apt install12.66 MB
21RUN /bin/bash -c apt install8.76 MB
22RUN /bin/bash -c pip install1.51 KB
23RUN /bin/bash -c pip install52.92 MB
24ENV NVTE_USE_HIPBLASLT=10 B
25ENV NVTE_FRAMEWORK=jax0 B
26ENV NVTE_ROCM_ARCH=gfx9420 B
27ENV Python_ROOT=/pyenv0 B
28ENV ROCM_PATH=/opt/rocm0 B
29WORKDIR /workspace99 B
30ENV HIPBLASLT_BRANCH=develop0 B
31ENV HIPBLASLT_COMMIT=78ec86229999aea66e86d70cf17ec1f6a11d0f870 B
32ENV PYTORCH_ROCM_ARCH=gfx9420 B
33RUN /bin/bash -c git clone5.29 KB
34RUN /bin/bash -c git clone523.69 MB
35RUN /bin/bash -c sed -i151.19 KB
36RUN /bin/bash -c sed -i151.19 KB
37ARG MAX_JOBS=1800 B
38RUN |1 MAX_JOBS=180 /bin/bash -c720.09 KB
39ENV AMDGPU_TARGETS=gfx9420 B
40ENV TE_COMMIT=f81a3eb4779e3d65e3a7ca0c4710554f68109d330 B
41RUN |1 MAX_JOBS=180 /bin/bash -c568.02 MB
42RUN |1 MAX_JOBS=180 /bin/bash -c300.56 MB
43RUN |1 MAX_JOBS=180 /bin/bash -c1.69 GB
44RUN |1 MAX_JOBS=180 /bin/bash -c2.33 KB
45RUN |1 MAX_JOBS=180 /bin/bash -c2.34 MB
46ENV NVTE_USE_HIPBLASLT=10 B
47ENV GPU_MAX_HW_QUEUES=20 B
48ENV HIP_FORCE_DEV_KERNARG=10 B
49ENV NVTE_FUSED_ATTN=10 B
50ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B
51ENV NCCL_CHECKS_DISABLE=10 B
52ENV NCCL_IB_GID_INDEX=30 B
53ENV NCCL_CROSS_NIC=00 B
54ENV NCCL_PROTO=Simple0 B
55ENV PROFILE_TB=10 B
56ENV NVTE_CK_V3_BF16_CVT=20 B
57ENV NVTE_CK_USES_BWD_V3=10 B
58ENV LD_LIBRARY_PATH=/opt/rocm/lib:0 B

Command

 ARG RELEASE