Sign inSign up

rocm/jax-training:maxtext-v25.5

Manifest digest

sha256:4e0516358a227cae8f552fb866ec07e2edcf244756f02e7b40212abfbab5217b

OS/ARCH

linux/amd64

Compressed size

9.49 GB

Last pushed

over 1 year by rocmdev

Type

Image

Manifest digest

sha256:4e0516358a227cae8f552fb866ec07e2edcf244756f02e7b40212abfbab5217b

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.17 MB
6CMD ["/bin/bash"]0 B
7RUN /bin/sh -c apt-get update50.02 MB
8RUN /bin/sh -c apt-get update16.87 MB
9ARG GPU_DEVICE_TARGETS=gfx908 gfx90a gfx942 gfx10300 B
10ENV GPU_DEVICE_TARGETS=gfx908 gfx90a gfx942 gfx10300 B
11ARG ROCM_VERSION=6.3.40 B
12ARG ROCM_PATH=/opt/rocm-6.3.40 B
13ENV ROCM_PATH=/opt/rocm-6.3.40 B
14ARG ROCM_BUILD_JOB=0 B
15ARG ROCM_BUILD_NUM=0 B
16RUN |5 GPU_DEVICE_TARGETS=gfx908 gfx90a gfx9425.58 GB
17ENV PATH=/opt/rocm-6.3.4/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
18ENV PATH=/root/bin:/root/.local/bin:/opt/rocm-6.3.4/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
19RUN |5 GPU_DEVICE_TARGETS=gfx908 gfx90a gfx94214.6 MB
20ARG PYTHON_VERSION=3.100 B
21RUN |6 GPU_DEVICE_TARGETS=gfx908 gfx90a gfx9427.44 MB
22ENV PYENV_ROOT=/pyenv0 B
23ENV PATH=/pyenv/shims:/pyenv/bin:/root/bin:/root/.local/bin:/opt/rocm-6.3.4/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
24RUN |6 GPU_DEVICE_TARGETS=gfx908 gfx90a gfx94255.68 MB
25RUN |6 GPU_DEVICE_TARGETS=gfx908 gfx90a gfx94293.87 MB
26ARG JAX_VERSION=0.4.350 B
27ARG JAX_COMMIT=fd0366b169cbf7dab58b488c8357db7534a6cec50 B
28ARG XLA_COMMIT=5fbe8579b04c31fd0bec6f984c7520a46bccfb460 B
29ARG JAX_USE_CLANG0 B
30LABEL com.amdgpu.rocm_version=6.3.4 com.amdgpu.python_version=3.10 com.amdgpu.jax_version=0.4.35 com.amdgpu.jax_commit=fd0366b169cbf7dab58b488c8357db7534a6cec50 B
31RUN |10 GPU_DEVICE_TARGETS=gfx908 gfx90a gfx942103 B
32RUN |10 GPU_DEVICE_TARGETS=gfx908 gfx90a gfx942469.25 MB
33ARG TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e0 B
34RUN |1 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e /bin/sh -c32 B
35RUN |1 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e /bin/sh -c32 B
36RUN |1 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e /bin/sh -c4.96 MB
37RUN |1 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e /bin/sh -c12.68 MB
38RUN |1 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e /bin/sh -c12.04 MB
39ENV Python_ROOT=/pyenv0 B
40RUN |1 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e /bin/sh -c36.4 KB
41RUN |1 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e /bin/sh -c888.75 KB
42WORKDIR /workspace99 B
43RUN |1 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e /bin/sh -c462 B
44RUN |1 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e /bin/sh -c453.41 KB
45RUN |1 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e /bin/sh -c32 B
46ENV NVTE_USE_HIPBLASLT=10 B
47ENV NVTE_FRAMEWORK=jax0 B
48ENV NVTE_ROCM_ARCH=gfx9420 B
49RUN |1 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e /bin/sh -c453.56 MB
50RUN |1 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e /bin/sh -c300.59 MB
51RUN |1 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e /bin/sh -c1.71 GB
52RUN |1 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e /bin/sh -c2.53 KB
53ARG HIPBLASLT_DEV_NAME=hipblaslt-dev_0.13.0-ae9c477a_amd64.deb0 B
54ARG HIPBLASLT_NAME=hipblaslt_0.13.0-ae9c477a_amd64.deb0 B
55RUN |3 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e HIPBLASLT_DEV_NAME=hipblaslt-dev_0.13.0-ae9c477a_amd64.deb HIPBLASLT_NAME=hipblaslt_0.13.0-ae9c477a_amd64.deb32 B
56RUN |3 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e HIPBLASLT_DEV_NAME=hipblaslt-dev_0.13.0-ae9c477a_amd64.deb HIPBLASLT_NAME=hipblaslt_0.13.0-ae9c477a_amd64.deb32 B
57RUN |3 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e HIPBLASLT_DEV_NAME=hipblaslt-dev_0.13.0-ae9c477a_amd64.deb HIPBLASLT_NAME=hipblaslt_0.13.0-ae9c477a_amd64.deb177 B
58RUN |3 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e HIPBLASLT_DEV_NAME=hipblaslt-dev_0.13.0-ae9c477a_amd64.deb HIPBLASLT_NAME=hipblaslt_0.13.0-ae9c477a_amd64.deb32 B
59RUN |3 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e HIPBLASLT_DEV_NAME=hipblaslt-dev_0.13.0-ae9c477a_amd64.deb HIPBLASLT_NAME=hipblaslt_0.13.0-ae9c477a_amd64.deb694.43 MB
60RUN |3 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e HIPBLASLT_DEV_NAME=hipblaslt-dev_0.13.0-ae9c477a_amd64.deb HIPBLASLT_NAME=hipblaslt_0.13.0-ae9c477a_amd64.deb268.94 KB
61RUN |3 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e HIPBLASLT_DEV_NAME=hipblaslt-dev_0.13.0-ae9c477a_amd64.deb HIPBLASLT_NAME=hipblaslt_0.13.0-ae9c477a_amd64.deb32 B
62RUN |3 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e HIPBLASLT_DEV_NAME=hipblaslt-dev_0.13.0-ae9c477a_amd64.deb HIPBLASLT_NAME=hipblaslt_0.13.0-ae9c477a_amd64.deb32 B
63RUN |3 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e HIPBLASLT_DEV_NAME=hipblaslt-dev_0.13.0-ae9c477a_amd64.deb HIPBLASLT_NAME=hipblaslt_0.13.0-ae9c477a_amd64.deb170 B
64RUN |3 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e HIPBLASLT_DEV_NAME=hipblaslt-dev_0.13.0-ae9c477a_amd64.deb HIPBLASLT_NAME=hipblaslt_0.13.0-ae9c477a_amd64.deb32 B
65RUN |3 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e HIPBLASLT_DEV_NAME=hipblaslt-dev_0.13.0-ae9c477a_amd64.deb HIPBLASLT_NAME=hipblaslt_0.13.0-ae9c477a_amd64.deb97.4 KB
66RUN |3 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e HIPBLASLT_DEV_NAME=hipblaslt-dev_0.13.0-ae9c477a_amd64.deb HIPBLASLT_NAME=hipblaslt_0.13.0-ae9c477a_amd64.deb97.4 KB
67RUN |3 TE_COMMIT=b8b92dcc1e4c3a499a56759b21f677dcf7b67b7e HIPBLASLT_DEV_NAME=hipblaslt-dev_0.13.0-ae9c477a_amd64.deb HIPBLASLT_NAME=hipblaslt_0.13.0-ae9c477a_amd64.deb6.13 MB
68ENV NVTE_USE_HIPBLASLT=10 B
69ENV GPU_MAX_HW_QUEUES=20 B
70ENV HIP_FORCE_DEV_KERNARG=10 B
71ENV NVTE_FUSED_ATTN=10 B
72ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B
73ENV NCCL_CHECKS_DISABLE=10 B
74ENV NCCL_IB_GID_INDEX=30 B
75ENV NCCL_CROSS_NIC=00 B
76ENV NCCL_PROTO=Simple0 B
77ENV PROFILE_TB=10 B
78ENV NVTE_CK_V3_BF16_CVT=20 B
79ENV NVTE_CK_USES_BWD_V3=10 B
80ENV LD_LIBRARY_PATH=/opt/rocm/lib0 B
81RUN /bin/sh -c apt install31.26 MB

Command

 ARG RELEASE