rocm/jax-training:maxtext-v25.11-jax08-EarlyAccess

Manifest digest

sha256:0ee8405a744011cf209b427c8ada9f8a0e3a943ec6447229341451d37260f56c

OS/ARCH

linux/amd64

Compressed size

17.95 GB

Last pushed

8 months by rocmdev

Type

Image

Manifest digest

sha256:0ee8405a744011cf209b427c8ada9f8a0e3a943ec6447229341451d37260f56c

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=24.040 B

5ADD file ... in / 28.35 MB

6CMD ["/bin/bash"]0 B

7RUN /bin/sh -c apt-get update175.86 MB

8RUN /bin/sh -c apt-get update1.63 MB

9ARG GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a gfx9420 B

10ENV GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a gfx9420 B

11ARG ROCM_VERSION=7.1.00 B

12ARG ROCM_PATH=/opt/rocm-7.1.00 B

13ENV ROCM_PATH=/opt/rocm-7.1.00 B

14ARG ROCM_BUILD_JOB=0 B

15ARG ROCM_BUILD_NUM=0 B

16ARG PLUGIN_NAMESPACE=70 B

17ARG THEROCK_PATH0 B

18RUN |7 GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a5.7 GB

19ENV PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin:/opt/rocm/llvm/bin0 B

20ENV ln=-s /opt/rocm-/lib/librocblas.so.5 /opt/rocm-/lib/librocblas.so.40 B

21ENV LD_LIBRARY_PATH=/opt/rocm-7.1.0/lib:0 B

22ENV LLVM_PATH=/opt/rocm-7.1.0/llvm0 B

23ENV HCC_HOME=/opt/rocm-7.1.0/hcc0 B

24ENV HIP_PATH=/opt/rocm-7.1.0/0 B

25ENV OPENCL_ROOT=/opt/rocm-7.1.0/opencl0 B

26

ENV PATH=/opt/rocm-7.1.0/hcc/bin:/opt/rocm-7.1.0//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin:/opt/rocm/llvm/bin

0 B

27

ENV PATH=/opt/rocm-7.1.0/bin:/opt/rocm-7.1.0/hcc/bin:/opt/rocm-7.1.0//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin:/opt/rocm/llvm/bin

0 B

28

ENV PATH=/opt/rocm-7.1.0/opencl/bin:/opt/rocm-7.1.0/bin:/opt/rocm-7.1.0/hcc/bin:/opt/rocm-7.1.0//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin:/opt/rocm/llvm/bin

0 B

29

ENV PATH=/root/bin:/root/.local/bin:/opt/rocm-7.1.0/opencl/bin:/opt/rocm-7.1.0/bin:/opt/rocm-7.1.0/hcc/bin:/opt/rocm-7.1.0//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin:/opt/rocm/llvm/bin

0 B

30RUN |7 GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a98.54 MB

31RUN |7 GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a92.25 MB

32ARG JAX_VERSION0 B

33ARG XLA_COMMIT=2e5971172f0c0d8f3a4c37b55cc1b53b4ec37cae0 B

34ARG JAX_COMMIT=5712de44e97c455faed1fd45532e821ca66d025a0 B

35ARG ROCM_JAX_COMMIT=f10dbb67e9d738e2dd01b0d916afe4db2473bdd70 B

36

LABEL com.amdgpu.rocm_version=7.1.0 com.amdgpu.python_version=3.12 com.amdgpu.jax_version= com.amdgpu.jax_commit=5712de44e97c455faed1fd45532e821ca66d025a

0 B

37RUN |11 GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a157.22 MB

38ENV ROCM_PATH=/opt/rocm-7.1.00 B

39

ENV PATH=/opt/rocm-7.1.0/opencl/bin:/opt/rocm-7.1.0/bin:/opt/rocm-7.1.0/hcc/bin:/opt/rocm-7.1.0//bin:/root/bin:/root/.local/bin:/opt/rocm-7.1.0/opencl/bin:/opt/rocm-7.1.0/bin:/opt/rocm-7.1.0/hcc/bin:/opt/rocm-7.1.0//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin:/opt/rocm/llvm/bin

0 B

40ENV MAX_JOBS=1280 B

41RUN /bin/sh -c echo 'deb231 B

42RUN /bin/sh -c wget -O4.89 KB

43RUN /bin/sh -c sudo apt-get13.51 MB

44RUN /bin/sh -c sudo apt-get218.43 MB

45RUN /bin/sh -c apt update1.62 GB

46RUN /bin/sh -c python -m4 MB

47

ENV PATH=/opt/venv/bin:/opt/rocm-7.1.0/opencl/bin:/opt/rocm-7.1.0/bin:/opt/rocm-7.1.0/hcc/bin:/opt/rocm-7.1.0//bin:/root/bin:/root/.local/bin:/opt/rocm-7.1.0/opencl/bin:/opt/rocm-7.1.0/bin:/opt/rocm-7.1.0/hcc/bin:/opt/rocm-7.1.0//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/opt/rocm/bin:/opt/rocm/llvm/bin

0 B

48RUN /bin/sh -c pip install5.13 MB

49RUN /bin/sh -c pip uninstall259.9 KB

50RUN /bin/sh -c pip install52.73 MB

51RUN /bin/sh -c cd /opt/rocm/share/amd_smi3.49 MB

52WORKDIR /workspace99 B

53RUN /bin/sh -c pip install368.9 KB

54ARG ROCM_REPO=https://github.com/ROCm/rocm-systems.git0 B

55ARG ROCM_BRANCH=release/rocm-rel-7.10 B

56ENV CLR_DIR=/workspace/rocm-systems/projects/clr0 B

57ENV HIP_DIR=/workspace/rocm-systems/projects/hip0 B

58RUN |2 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 /bin/sh13.9 MB

59ARG ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git0 B

60ARG ROCPROF_BRANCH=ed42157c3145a2f10f3c95c6059efd6f367064190 B

61

RUN |4 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

14.06 MB

62ARG HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git0 B

63ENV HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git0 B

64ARG HIPBLASLT_BRANCH=0d52365d1dd659236c9632d0756ebf9cdf4c38d70 B

65ENV HIPBLASLT_BRANCH=0d52365d1dd659236c9632d0756ebf9cdf4c38d70 B

66ENV PYTORCH_ROCM_ARCH=gfx942;gfx9500 B

67ENV Python_ROOT=/pyenv0 B

68

RUN |6 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

176 B

69

RUN |6 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

1.52 GB

70

RUN |6 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

186 B

71

RUN |6 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

172 B

72

RUN |6 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

94.66 KB

73

RUN |6 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

94.67 KB

74

RUN |6 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

172 B

75WORKDIR /workspace32 B

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

77ARG MAXTEXT_BRANCH=genesu/add-jax-training-patches0 B

78

RUN |8 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

1.58 GB

79ENV ROCM_JAX_REPO=https://github.com/ROCm/rocm-jax.git0 B

80ENV ROCM_JAX_BRANCH=be6011b1b9d3e61e103fcf8700f54f9742d05fb50 B

81ARG JAX_VERSION=0.8.00 B

82ENV JAX_VERSION=0.8.00 B

83

RUN |9 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

5.86 GB

84

RUN |9 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

164.45 MB

85ENV XLA_HOME=/workspace/rocm-jax/xla0 B

86ARG TE_REPO=https://github.com/ROCm/TransformerEngine.git0 B

87ARG TE_BRANCH=9ce0ac0122d3cd04900cc719f68a1627325f0f120 B

88ENV USE_ROCM=10 B

89ENV HIP_PATH=/opt/rocm0 B

90ENV NVTE_FRAMEWORK=jax0 B

91ENV CMAKE_BUILD_PARALLEL_LEVEL=640 B

92ENV NVTE_ROCM_ARCH=gfx942;gfx9500 B

93ENV NVTE_USE_ROCM=10 B

94ENV PYTHONPATH=/home/TransformerEngine/3rdparty/hipify_torch0 B

95

RUN |11 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

174 B

96

RUN |11 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

190 B

97

RUN |11 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

668.93 MB

98

RUN |11 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

142 B

99ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B

100ENV NVTE_USE_HIPBLASLT=10 B

101ENV GPU_MAX_HW_QUEUES=20 B

102ENV HIP_FORCE_DEV_KERNARG=10 B

103ENV HSA_FORCE_FINE_GRAIN_PCIE=10 B

104ENV NVTE_FUSED_ATTN=10 B

105ENV NCCL_DEBUG=VERSION0 B

106ENV NVTE_CK_USES_BWD_V3=10 B

107ENV NVTE_CK_USES_FWD_V3=10 B

108ENV NVTE_CK_IS_V3_ATOMIC_FP32=00 B

109ENV NVTE_CK_HOW_V3_BF16_CVT=20 B

110ENV XLA_PYTHON_CLIENT_MEM_FRACTION=.970 B

111

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

112

RUN |11 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

129 B

113ARG GIT_COMMIT_TAG=16a78b68f79b8be8f0311cc7c47955303868d2380 B

114

RUN |12 ROCM_REPO=https://github.com/ROCm/rocm-systems.git ROCM_BRANCH=release/rocm-rel-7.1 ROCPROF_REPO=https://github.com/ROCm/rocm-systems.git

185 B

115bash1.55 MB

Command

 ARG RELEASE

Help

System theme

Docker Suite

rocm/jax-training:maxtext-v25.11-jax08-EarlyAccess