Sign inSign up

rocm/jax-training:latest

Index digest

sha256:b034a6769b58a72f6584e104d3821fe11d51f79225e8ff9d7b88cac81fc6867a

OS/ARCH

linux/amd64

Compressed size

12.09 GB

Last pushed

20 days by rocm

Type

Image

Manifest digest

sha256:21abbf2a142582ec62c53834b8679feb081eefaece75852d62a3575ab2636a18

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 ROCPROFILER_QUEUE_INTERPOSITION=00 B
11ENV DEBUG_HIP_DYNAMIC_QUEUES=00 B
12ENV DEBIAN_FRONTEND=noninteractive0 B
13RUN /bin/sh -c apt update276.59 MB
14RUN /bin/sh -c python -m4 MB
15ENV PATH=/opt/venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
16RUN /bin/sh -c pip install32.16 MB
17ENV THEROCK_TARBALL=https://repo.amd.com/rocm/tarball-multi-arch/therock-dist-linux-multiarch-7.14.0.tar.gz0 B
18RUN /bin/sh -c mkdir -p7.97 GB
19ENV 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
20ENV LD_LIBRARY_PATH=/opt/rocm/lib:/opt/rocm/lib/rocm_sysdeps/lib0 B
21ENV ROCM_PATH=/opt/rocm0 B
22ENV HIP_PLATFORM=amd0 B
23ENV HIP_DEVICE_LIB_PATH=/opt/rocm/lib/llvm/amdgcn/bitcode0 B
24RUN /bin/sh -c echo 'deb236 B
25RUN /bin/sh -c wget -O2.62 KB
26RUN /bin/sh -c sudo apt-get34.46 MB
27RUN /bin/sh -c sudo apt-get241.64 MB
28RUN /bin/sh -c pip install174.91 KB
29WORKDIR /workspace32 B
30ARG MAXTEXT_REPO=https://github.com/ROCm/maxtext.git0 B
31ARG MAXTEXT_BRANCH=release/v26.50 B
32RUN |2 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 /bin/sh1.61 GB
33ARG TF_REPO=https://github.com/ROCm/tensorflow-upstream.git0 B
34ARG TF_BRANCH=upstream-v2.21.00 B
35ARG BAZELISK_VERSION=v1.25.00 B
36RUN |5 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 TF_REPO=https://github.com/ROCm/tensorflow-upstream.git3.03 MB
37RUN |5 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 TF_REPO=https://github.com/ROCm/tensorflow-upstream.git264.61 MB
38ARG JAX_VERSION=0.10.00 B
39ENV JAX_VERSION=0.10.00 B
40ENV JAX_PJRT_VERSION=0.10.0+rocm7.14.00 B
41ENV JAX_PLUGIN_VERSION=0.10.0+rocm7.14.00 B
42RUN |6 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 TF_REPO=https://github.com/ROCm/tensorflow-upstream.git183.16 MB
43ARG TE_VERSION=2.15.0.dev0+rocm7.15.0a20260707.72d01a00 B
44RUN |7 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 TF_REPO=https://github.com/ROCm/tensorflow-upstream.git330.19 MB
45ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B
46ENV NVTE_USE_HIPBLASLT=10 B
47ENV GPU_MAX_HW_QUEUES=20 B
48ENV HIP_FORCE_DEV_KERNARG=10 B
49ENV HSA_FORCE_FINE_GRAIN_PCIE=10 B
50ENV NVTE_FUSED_ATTN=10 B
51ENV NCCL_DEBUG=VERSION0 B
52ENV NVTE_CK_USES_BWD_V3=10 B
53ENV NVTE_CK_USES_FWD_V3=10 B
54ENV NVTE_CK_IS_V3_ATOMIC_FP32=10 B
55ENV NVTE_CK_HOW_V3_BF16_CVT=20 B
56ENV XLA_PYTHON_CLIENT_MEM_FRACTION=.90 B
57ENV 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
58WORKDIR /workspace32 B
59ARG PRIMUS_REPO=https://github.com/AMD-AGI/Primus.git0 B
60ARG PRIMUS_BRANCH=main0 B
61RUN |9 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 TF_REPO=https://github.com/ROCm/tensorflow-upstream.git994.41 MB
62RUN |9 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 TF_REPO=https://github.com/ROCm/tensorflow-upstream.git26.26 MB
63ARG AINIC_BUNDLE_VERSION=1.117.5-a-770 B
64RUN |10 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 TF_REPO=https://github.com/ROCm/tensorflow-upstream.git2.66 MB
65RUN |10 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 TF_REPO=https://github.com/ROCm/tensorflow-upstream.git227 B
66ARG UCX_VERSION=1.18.00 B
67RUN |11 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 TF_REPO=https://github.com/ROCm/tensorflow-upstream.git86.44 MB
68ENV UCX_INSTALL_DIR=/workspace/ucx-1.18.0/install0 B
69ARG MPI_VERSION=4.1.60 B
70RUN |12 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 TF_REPO=https://github.com/ROCm/tensorflow-upstream.git3.84 MB
71RUN |12 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 TF_REPO=https://github.com/ROCm/tensorflow-upstream.git45.54 MB
72RUN |12 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 TF_REPO=https://github.com/ROCm/tensorflow-upstream.git442 B
73RUN |12 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 TF_REPO=https://github.com/ROCm/tensorflow-upstream.git32 B
74ARG GIT_COMMIT_TAG=5ffe026e7bb969c42ce9f3c9f5c4a3c3164c8a8c0 B
75ARG DOCKERFILE_PATH=jax_training_the_rock/docker/Dockerfile0 B
76RUN |14 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git MAXTEXT_BRANCH=release/v26.5 TF_REPO=https://github.com/ROCm/tensorflow-upstream.git17.94 KB
77COPY jax_training_the_rock/docker/Dockerfile /workspace/.manifest/Dockerfile # buildkit3.77 KB

Command

 ARG RELEASE