rocm/jax-training:maxtext-v26.8-jax0.11.1-te2.18.0

Manifest digest

sha256:8267d6a4a016063575722b4113b8dd269b015ef07f81da1ccee4b2beb289c767

OS/ARCH

linux/amd64

Compressed size

11.13 GB

Last pushed

5 days by rocm

Type

Image

Manifest digest

sha256:8267d6a4a016063575722b4113b8dd269b015ef07f81da1ccee4b2beb289c767

Image Layers

Image Layers

1ARG RELEASE0 B

2ARG LAUNCHPAD_BUILD_ARCH0 B

3LABEL org.opencontainers.image.version=24.040 B

4ADD file ... in / 28.39 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

13ARG LIBC_SECURITY_REFRESH=2026-10-010 B

14RUN |1 LIBC_SECURITY_REFRESH=2026-10-01 /bin/sh -c239.22 MB

15COPY /uv /usr/local/bin/uv # buildkit20.96 MB

16RUN |1 LIBC_SECURITY_REFRESH=2026-10-01 /bin/sh -c9.64 KB

17ENV PATH=/opt/venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B

18ENV VIRTUAL_ENV=/opt/venv0 B

19COPY /dist/pip-*.whl /tmp/ # buildkit1.71 MB

20RUN |1 LIBC_SECURITY_REFRESH=2026-10-01 /bin/sh -c29.26 MB

21ENV THE_ROCK_VERSION=10.2.0a202609230 B

22RUN |1 LIBC_SECURITY_REFRESH=2026-10-01 /bin/sh -c3.72 GB

23RUN |1 LIBC_SECURITY_REFRESH=2026-10-01 /bin/sh -c3.72 GB

24COPY /bazelisk /usr/local/bin/bazel # buildkit5.67 MB

25ENV ROCM_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel0 B

26ENV ROCM_HOME=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel0 B

27ENV HIP_PLATFORM=amd0 B

28ENV HIP_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel0 B

29ENV HIP_CLANG_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/llvm/bin0 B

30ENV HIP_INCLUDE_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/include0 B

31ENV HIP_LIB_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib0 B

32ENV HIP_DEVICE_LIB_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib/llvm/amdgcn/bitcode0 B

33

ENV PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/bin:/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/llvm/bin:/opt/venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin

0 B

34

ENV LIBRARY_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib:/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib:/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib64

0 B

35ENV CPATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/include0 B

36ENV PKG_CONFIG_PATH=/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib/pkgconfig0 B

37ENV LD_LIBRARY_PATH=0 B

38ARG JAX_VERSION=0.11.10 B

39ENV JAX_VERSION=0.11.10 B

40ENV JAX_PJRT_VERSION=0.11.1+rocm10.2.0a202609230 B

41ENV JAX_PLUGIN_VERSION=0.11.1+rocm10.2.0a202609230 B

42RUN |2 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 /bin/sh275.89 MB

43ARG TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f0 B

44RUN |3 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f400.37 MB

45COPY /mount.gcsfuse /usr/sbin/mount.gcsfuse # buildkit18.92 MB

46ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B

47ENV NVTE_USE_HIPBLASLT=10 B

48ENV GPU_MAX_HW_QUEUES=20 B

49ENV HIP_FORCE_DEV_KERNARG=10 B

50ENV HSA_FORCE_FINE_GRAIN_PCIE=10 B

51ENV NVTE_FUSED_ATTN=10 B

52ENV NCCL_DEBUG=VERSION0 B

53ENV NVTE_CK_USES_BWD_V3=10 B

54ENV NVTE_CK_USES_FWD_V3=10 B

55ENV NVTE_CK_IS_V3_ATOMIC_FP32=10 B

56ENV NVTE_CK_HOW_V3_BF16_CVT=20 B

57ENV XLA_PYTHON_CLIENT_MEM_FRACTION=.90 B

58

ENV 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=8589934592

0 B

59WORKDIR /workspace32 B

60ARG MAXTEXT_REPO=https://github.com/ROCm/maxtext.git0 B

61ARG MAXTEXT_BRANCH=release/v26.80 B

62RUN |5 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f1.52 GB

63ARG TF_REPO=https://github.com/ROCm/tensorflow-upstream.git0 B

64ARG TF_BRANCH=upstream-v2.21.00 B

65RUN |7 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f505.87 KB

66COPY /909041e4802c4b9a2223ca04099f35bf1dbbd460.tar.gz /opt/bazel_distdir/ # buildkit243.86 MB

67RUN |7 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f254.42 MB

68RUN |7 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f21.31 MB

69COPY /mount.gcsfuse /usr/sbin/mount.gcsfuse # buildkit18.92 MB

70COPY /gcsfuse /usr/bin/gcsfuse # buildkit24.89 MB

71RUN |7 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f25.73 MB

72ARG AINIC_BUNDLE_VERSION=1.117.5-a-770 B

73RUN |8 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f362.92 KB

74RUN |8 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f226 B

75ARG UCX_VERSION=1.18.00 B

76RUN |9 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f86.43 MB

77ENV UCX_INSTALL_DIR=/workspace/ucx-1.18.0/install0 B

78ARG MPI_VERSION=4.1.60 B

79RUN |10 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f3.84 MB

80RUN |10 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f32 B

81ARG GIT_COMMIT_TAG=73710dda449ad7aaf12207a984f986a53d1234e20 B

82ARG DOCKERFILE_PATH=pipeline/jax_training/docker/Dockerfile0 B

83RUN |12 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f16 KB

84COPY pipeline/jax_training/docker/Dockerfile /workspace/.manifest/docker-build-recipe.txt # buildkit11.21 KB

85ENV NVTE_FRAMEWORK=jax0 B

86RUN |12 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f195.11 MB

87RUN |12 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f42.35 MB

88RUN |12 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f13.05 MB

89RUN |12 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f12.61 KB

90ARG MAXDIFFUSION_REPO=https://github.com/AI-Hypercomputer/maxdiffusion.git0 B

91ARG MAXDIFFUSION_COMMIT=68e069659f0af80694559e29939a17f879fe7f6a0 B

92RUN |14 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f266.49 MB

93RUN |14 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f3.17 KB

94RUN |14 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f13.36 KB

95RUN |14 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f6.27 KB

96RUN |14 LIBC_SECURITY_REFRESH=2026-10-01 JAX_VERSION=0.11.1 TE_VERSION=2.18.0.dev0+rocm10.2.0a20260929.03afd8f888 B

Command

 ARG RELEASE