rocm/jax-training:maxtext-v26.6

Manifest digest

sha256:a71d8dbb045e890a981d916de10b11f65e43a10892a81a801194db602c8eb6a5

OS/ARCH

linux/amd64

Compressed size

11.19 GB

Last pushed

17 days by rocm

Type

Image

Manifest digest

sha256:a71d8dbb045e890a981d916de10b11f65e43a10892a81a801194db602c8eb6a5

Image Layers

Image Layers

1ARG RELEASE0 B

2ARG LAUNCHPAD_BUILD_ARCH0 B

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

4ADD file ... in / 28.37 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 update240.75 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

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

17RUN /bin/sh -c pip install32.84 MB

18ENV THE_ROCK_VERSION=7.14.00 B

19RUN /bin/sh -c python -m3.16 GB

20RUN /bin/sh -c rocm-sdk init3.16 GB

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

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

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

24ENV HIP_PLATFORM=amd0 B

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

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

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

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

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

30

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

31

ENV LD_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:/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/llvm/lib:/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib/host-math/lib:/opt/venv/lib/python3.12/site-packages/_rocm_sdk_devel/lib/rocm_sysdeps/lib

0 B

32

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

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

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

35RUN /bin/sh -c git clone45.54 MB

36ARG JAX_VERSION=0.11.00 B

37ENV JAX_VERSION=0.11.00 B

38ARG JAX_ROCM_VERSION=0.11.0.post10 B

39RUN |2 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 /bin/sh285.16 MB

40ARG TE_VERSION=2.17.0+rocm7.14.0.50a84ad0 B

41RUN |3 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad392.55 MB

42COPY /mount.gcsfuse /usr/sbin/mount.gcsfuse # buildkit18.24 MB

43ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B

44ENV NVTE_USE_HIPBLASLT=10 B

45ENV GPU_MAX_HW_QUEUES=20 B

46ENV HIP_FORCE_DEV_KERNARG=10 B

47ENV HSA_FORCE_FINE_GRAIN_PCIE=10 B

48ENV NVTE_FUSED_ATTN=10 B

49ENV NCCL_DEBUG=VERSION0 B

50ENV NVTE_CK_USES_BWD_V3=10 B

51ENV NVTE_CK_USES_FWD_V3=10 B

52ENV NVTE_CK_IS_V3_ATOMIC_FP32=10 B

53ENV NVTE_CK_HOW_V3_BF16_CVT=20 B

54ENV XLA_PYTHON_CLIENT_MEM_FRACTION=.90 B

55

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

56WORKDIR /workspace32 B

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

58ARG MAXTEXT_BRANCH=release/v26.60 B

59RUN |5 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad1.5 GB

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

61ARG TF_BRANCH=upstream-v2.21.00 B

62RUN |7 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad372.96 KB

63RUN |7 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad264.61 MB

64ARG PRIMUS_REPO=https://github.com/AMD-AGI/Primus.git0 B

65ARG PRIMUS_BRANCH=main0 B

66RUN |9 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad1.39 GB

67RUN |9 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad20.29 MB

68COPY /mount.gcsfuse /usr/sbin/mount.gcsfuse # buildkit18.24 MB

69COPY /gcsfuse /usr/bin/gcsfuse # buildkit24.85 MB

70RUN |9 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad25.15 MB

71ARG AINIC_BUNDLE_VERSION=1.117.5-a-770 B

72RUN |10 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad363.5 KB

73RUN |10 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad227 B

74ARG UCX_VERSION=1.18.00 B

75RUN |11 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad86.43 MB

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

77ARG MPI_VERSION=4.1.60 B

78RUN |12 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad3.84 MB

79RUN |12 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad539 B

80RUN |12 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad32 B

81ARG GIT_COMMIT_TAG=31a3a21d0a37cbd0b5de0342535f65d557ff77ba0 B

82ARG DOCKERFILE_PATH=jax_training_the_rock/docker/Dockerfile0 B

83RUN |14 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad15.21 KB

84COPY jax_training_the_rock/docker/Dockerfile /workspace/.manifest/docker-build-recipe.txt # buildkit9.95 KB

85ENV LD_LIBRARY_PATH=0 B

86ENV NVTE_FRAMEWORK=jax0 B

87RUN |14 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad189.46 MB

88RUN |14 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad44.1 MB

89RUN |14 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad30.25 MB

90RUN |14 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad12.62 KB

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

92ARG MAXDIFFUSION_COMMIT=68e069659f0af80694559e29939a17f879fe7f6a0 B

93RUN |16 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad263.89 MB

94RUN |16 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad3.18 KB

95RUN |16 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad13.36 KB

96RUN |16 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad6.26 KB

97RUN |16 JAX_VERSION=0.11.0 JAX_ROCM_VERSION=0.11.0.post1 TE_VERSION=2.17.0+rocm7.14.0.50a84ad32 B

Command

 ARG RELEASE