rocm/jax-training:maxtext-v26.1

Manifest digest

sha256:901083bde353fe6362ada3036e452c792b2c96124e5900f4e9b5946c02ff9d6a

OS/ARCH

linux/amd64

Compressed size

21.68 GB

Last pushed

6 months by rocmdev

Type

Image

Manifest digest

sha256:901083bde353fe6362ada3036e452c792b2c96124e5900f4e9b5946c02ff9d6a

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 / 29.18 MB

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

7RUN /bin/sh -c apt-get update176.23 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.10 B

12ARG ROCM_PATH=/opt/rocm-7.1.10 B

13ENV ROCM_PATH=/opt/rocm-7.1.10 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 gfx90a6.06 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.1/lib:0 B

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

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

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

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

26

ENV PATH=/opt/rocm-7.1.1/hcc/bin:/opt/rocm-7.1.1//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.1/bin:/opt/rocm-7.1.1/hcc/bin:/opt/rocm-7.1.1//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.1/opencl/bin:/opt/rocm-7.1.1/bin:/opt/rocm-7.1.1/hcc/bin:/opt/rocm-7.1.1//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.1/opencl/bin:/opt/rocm-7.1.1/bin:/opt/rocm-7.1.1/hcc/bin:/opt/rocm-7.1.1//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.46 MB

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

32ARG JAX_VERSION0 B

33ARG XLA_COMMIT=4178517b7d7e5e3318ea4b701b01dc06847b73b00 B

34ARG JAX_COMMIT=5712de44e97c455faed1fd45532e821ca66d025a0 B

35ARG ROCM_JAX_COMMIT=be6011b1b9d3e61e103fcf8700f54f9742d05fb50 B

36

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

0 B

37RUN |11 GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a173.23 MB

38ENV ROCM_PATH=/opt/rocm-7.1.10 B

39

ENV PATH=/opt/rocm-7.1.1/opencl/bin:/opt/rocm-7.1.1/bin:/opt/rocm-7.1.1/hcc/bin:/opt/rocm-7.1.1//bin:/root/bin:/root/.local/bin:/opt/rocm-7.1.1/opencl/bin:/opt/rocm-7.1.1/bin:/opt/rocm-7.1.1/hcc/bin:/opt/rocm-7.1.1//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 apt update1.63 GB

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

43

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

0 B

44RUN /bin/sh -c pip install58.12 MB

45WORKDIR /workspace99 B

46ARG ROCR_REPO=https://github.com/ROCm/rocm-systems.git0 B

47ARG ROCR_BRANCH=release/rocm-rel-7.1.1.10 B

48ARG ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac669230 B

49

RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

176 B

50

RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

188 B

51

RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

1.44 MB

52

RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

206 B

53

RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

194 B

54

RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

233 B

55

RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

4.85 KB

56

RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

17.83 KB

57

RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

217.23 MB

58

RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

2.38 MB

59WORKDIR /workspace32 B

60

RUN |3 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

368.87 KB

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

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

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

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

65

RUN |5 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

13.9 MB

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

67ARG ROCPROF_BRANCH=ed42157c3145a2f10f3c95c6059efd6f367064190 B

68

RUN |7 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

14.06 MB

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

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

71ARG HIPBLASLT_BRANCH=fe818bf7619575499cb72225ced3e255aa08cbb50 B

72ENV HIPBLASLT_BRANCH=fe818bf7619575499cb72225ced3e255aa08cbb50 B

73ENV PYTORCH_ROCM_ARCH=gfx942;gfx9500 B

74ENV ROCM_AMDGPU_TARGETS=gfx942,gfx9500 B

75ENV Python_ROOT=/pyenv0 B

76

RUN |9 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

176 B

77

RUN |9 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

1.54 GB

78

RUN |9 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

187 B

79

RUN |9 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

172 B

80

RUN |9 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

94.81 KB

81

RUN |9 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

94.81 KB

82

RUN |9 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

162 B

83WORKDIR /workspace32 B

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

85ARG MAXTEXT_BRANCH=yeandy/update-patches-scaling-patch-v2-checkpoint-restore0 B

86

RUN |11 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

1.65 GB

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

88ENV ROCM_JAX_BRANCH=6251484e78c34751b2098cc1026adb79d4d51fad0 B

89ARG JAX_VERSION=0.8.20 B

90ENV JAX_VERSION=0.8.20 B

91ENV JAX_COMMIT=cef02f3ae4abf69862294cca9370c721eb7eb2b40 B

92ENV XLA_COMMIT=e301b64cf5ca2f4d356fcb825b25feeee31897330 B

93ENV XLA_BRANCH=rocm-jaxlib-v0.8.20 B

94ENV JAX_BRANCH=rocm-jaxlib-v0.8.20 B

95WORKDIR /opt32 B

96

RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

356.93 MB

97

RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

170.24 MB

98

RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

14.29 MB

99

RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

165.49 MB

100

RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

976.63 KB

101

RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

5.59 GB

102WORKDIR /opt/jax32 B

103

RUN |12 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

2.12 GB

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

105ARG TE_BRANCH=aec00a7f7ce1273ad1366dbf6d7074c9904999940 B

106WORKDIR /workspace32 B

107ENV USE_ROCM=10 B

108ENV HIP_PATH=/opt/rocm0 B

109ENV NVTE_FRAMEWORK=jax0 B

110ENV CMAKE_BUILD_PARALLEL_LEVEL=640 B

111ENV NVTE_ROCM_ARCH=gfx942;gfx9500 B

112ENV NVTE_USE_ROCM=10 B

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

114

ENV NVTE_AITER_PREBUILT_BASE_URL=https://compute-artifactory.amd.com:5000/artifactory/rocm-generic-local/te-ci/aiter-prebuilts

0 B

115

RUN |14 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

571.94 KB

116

RUN |14 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

872.13 MB

117ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B

118ENV NVTE_USE_HIPBLASLT=10 B

119ENV GPU_MAX_HW_QUEUES=20 B

120ENV HIP_FORCE_DEV_KERNARG=10 B

121ENV HSA_FORCE_FINE_GRAIN_PCIE=10 B

122ENV NVTE_FUSED_ATTN=10 B

123ENV NCCL_DEBUG=VERSION0 B

124ENV NVTE_CK_USES_BWD_V3=10 B

125ENV NVTE_CK_USES_FWD_V3=10 B

126ENV NVTE_CK_IS_V3_ATOMIC_FP32=00 B

127ENV NVTE_CK_HOW_V3_BF16_CVT=20 B

128ENV XLA_PYTHON_CLIENT_MEM_FRACTION=.970 B

129

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

130

RUN |14 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

1.55 MB

131

RUN |14 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

949 B

132WORKDIR /opt/jax32 B

133

RUN |14 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

89.44 MB

134WORKDIR /workspace32 B

135

RUN |14 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

68.61 MB

136ARG AINIC_BUNDLE_VERSION=1.117.5-a-380 B

137

RUN |15 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

9.92 MB

138

RUN |15 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

227 B

139ARG UCX_VERSION=1.18.00 B

140

RUN |16 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

86.44 MB

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

142ARG MPI_VERSION=4.1.60 B

143

RUN |17 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

3.84 MB

144ARG RCCL_REPO=https://github.com/ROCm/rccl.git0 B

145ARG RCCL_BRANCH=rocm-7.1.10 B

146

RUN |19 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

302.74 MB

147ARG ANP_REPO=https://github.com/rocm/amd-anp.git0 B

148ARG ANP_BRANCH=tags/v1.3.00 B

149ENV RCCL_HOME=/workspace/rccl0 B

150ENV MPI_PATH=/ompi-4.1.6/install0 B

151

RUN |21 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

107.31 MB

152

RUN |21 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

516.17 KB

153

RUN |21 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

131 B

154

RUN |21 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

126 B

155

RUN |21 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

126 B

156

RUN |21 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

128 B

157ARG GIT_COMMIT_TAG=e5be0ef7ea8210e52a702430bc62fb90fd5a73880 B

158

RUN |22 ROCR_REPO=https://github.com/ROCm/rocm-systems.git ROCR_BRANCH=release/rocm-rel-7.1.1.1 ROCR_CP_COMMIT=9baa65a8b71f4a0d99e7fdee4ff00a7c5ac66923

185 B

159bash351.24 KB

160bash31.55 MB

Command

 ARG RELEASE

Help

System theme

Docker Suite

rocm/jax-training:maxtext-v26.1