Sign inSign up

rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0

Manifest digest

sha256:9b263b38c15974cb1d1f1e174341cbc5fcb9522aada6f92eb03f7b9ac1f76e25

OS/ARCH

linux/amd64

Compressed size

15.4 GB

Last pushed

2 months by rocmdev

Type

Image

Manifest digest

sha256:9b263b38c15974cb1d1f1e174341cbc5fcb9522aada6f92eb03f7b9ac1f76e25

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
6WORKDIR /workspace/99 B
7ENV MAX_JOBS=1280 B
8ENV PYTORCH_ROCM_ARCH=gfx942;gfx9500 B
9ENV ROCM_AMDGPU_TARGETS=gfx942,gfx9500 B
10ENV DEBIAN_FRONTEND=noninteractive0 B
11RUN /bin/sh -c apt update274.92 MB
12RUN /bin/sh -c python -m4 MB
13ENV PATH=/opt/venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
14RUN /bin/sh -c pip install32.16 MB
15ENV THE_ROCK_VERSION=7.14.0a202605260 B
16ENV THEROCK_TARBALL=https://rocm.nightlies.amd.com/tarball-multi-arch/therock-dist-linux-multiarch-7.14.0a20260526.tar.gz0 B
17RUN /bin/sh -c mkdir -p10.96 GB
18ENV PATH=/opt/rocm/lib:/opt/venv/bin:/opt/rocm/bin:/opt/venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
19ENV LD_LIBRARY_PATH=/opt/rocm/lib:/opt/rocm/lib/rocm_sysdeps/lib0 B
20ENV ROCM_PATH=/opt/rocm0 B
21ENV HIP_PLATFORM=amd0 B
22ENV HIP_DEVICE_LIB_PATH=/opt/rocm/lib/llvm/amdgcn/bitcode0 B
23RUN /bin/sh -c echo 'deb235 B
24RUN /bin/sh -c wget -O2.63 KB
25RUN /bin/sh -c sudo apt-get41.39 MB
26RUN /bin/sh -c sudo apt-get240.89 MB
27RUN /bin/sh -c pip install174.92 KB
28WORKDIR /workspace32 B
29ARG MAXTEXT_REPO=https://github.com/ROCm/maxtext.git0 B
30ARG MAXTEXT_BRANCH=release/v26.40 B
31RUN |2 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 /bin/sh1.61 GB
32RUN |2 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 /bin/sh253.32 MB
33ARG JAX_VERSION=0.9.10 B
34ENV JAX_VERSION=0.9.10 B
35ENV GFX_FAMILY=gfx94X-dcgpu0 B
36ENV JAX_PJRT_VERSION=0.9.1+rocm7.14.0a202605260 B
37ENV JAX_PLUGIN_VERSION=0.9.1+rocm7.14.0a202605260 B
38RUN |3 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 JAX_VERSION=0.9.1278.12 MB
39ARG TE_REPO=https://github.com/ROCm/TransformerEngine.git0 B
40ARG TE_BRANCH=635d7c085c39a6d9bfe4881c7d3efab7a46d71290 B
41WORKDIR /workspace32 B
42ENV USE_ROCM=10 B
43ENV HIP_PATH=/opt/rocm0 B
44ENV NVTE_FRAMEWORK=jax0 B
45ENV CMAKE_BUILD_PARALLEL_LEVEL=640 B
46ENV NVTE_ROCM_ARCH=gfx942;gfx9500 B
47ENV NVTE_USE_ROCM=10 B
48ENV PYTHONPATH=/home/TransformerEngine/3rdparty/hipify_torch0 B
49RUN |5 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 JAX_VERSION=0.9.1650.14 KB
50RUN |5 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 JAX_VERSION=0.9.1708.27 MB
51ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B
52ENV NVTE_USE_HIPBLASLT=10 B
53ENV GPU_MAX_HW_QUEUES=20 B
54ENV HIP_FORCE_DEV_KERNARG=10 B
55ENV HSA_FORCE_FINE_GRAIN_PCIE=10 B
56ENV NVTE_FUSED_ATTN=10 B
57ENV NCCL_DEBUG=VERSION0 B
58ENV NVTE_CK_USES_BWD_V3=10 B
59ENV NVTE_CK_USES_FWD_V3=10 B
60ENV NVTE_CK_IS_V3_ATOMIC_FP32=10 B
61ENV NVTE_CK_HOW_V3_BF16_CVT=20 B
62ENV XLA_PYTHON_CLIENT_MEM_FRACTION=.970 B
63ENV 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
64RUN |5 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 JAX_VERSION=0.9.11.16 MB
65WORKDIR /workspace32 B
66ARG PRIMUS_REPO=https://github.com/AMD-AGI/Primus.git0 B
67ARG PRIMUS_BRANCH=main0 B
68RUN |7 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 JAX_VERSION=0.9.1910.03 MB
69RUN |7 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 JAX_VERSION=0.9.126.25 MB
70ARG AINIC_BUNDLE_VERSION=1.117.5-a-770 B
71RUN |8 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 JAX_VERSION=0.9.12.65 MB
72RUN |8 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 JAX_VERSION=0.9.1228 B
73ARG UCX_VERSION=1.18.00 B
74RUN |9 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 JAX_VERSION=0.9.186.44 MB
75ENV UCX_INSTALL_DIR=/workspace/ucx-1.18.0/install0 B
76ARG MPI_VERSION=4.1.60 B
77RUN |10 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 JAX_VERSION=0.9.13.84 MB
78RUN |10 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 JAX_VERSION=0.9.1446 B
79RUN |10 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 JAX_VERSION=0.9.132 B
80ARG GIT_COMMIT_TAG=b0b705ce0533ead5490070daa402ca625f1489b50 B
81ARG DOCKERFILE_PATH=jax_training_the_rock/docker/Dockerfile0 B
82RUN |12 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.4 JAX_VERSION=0.9.117.97 KB
83COPY jax_training_the_rock/docker/Dockerfile /workspace/.manifest/Dockerfile # buildkit3.43 KB
84ENV NCCL_DMABUF_ENABLE=10 B

Command

 ARG RELEASE