rocm/jax-training:maxtext-v25.11

Manifest digest

sha256:18e4d8f0b8ce7a7422c58046940dd5f32249960449fca09a562b65fb8eb1562a

OS/ARCH

linux/amd64

Compressed size

23.86 GB

Last pushed

8 months by rocmdev

Type

Image

Manifest digest

sha256:18e4d8f0b8ce7a7422c58046940dd5f32249960449fca09a562b65fb8eb1562a

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 'deb233 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.92 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=cbfc63b3f596a42eceeaba19716b0a1e8cb2e5870 B

65ENV HIPBLASLT_BRANCH=cbfc63b3f596a42eceeaba19716b0a1e8cb2e5870 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.53 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

185 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.7.10 B

82ENV JAX_VERSION=0.7.10 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

166.9 MB

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

5.86 GB

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

86ARG TE_BRANCH=0018b1bce0240ad53a6ac24a3b8052d9f067d34b0 B

87WORKDIR /workspace32 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

173 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

192 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

98ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B

99ENV NVTE_USE_HIPBLASLT=10 B

100ENV GPU_MAX_HW_QUEUES=20 B

101ENV HIP_FORCE_DEV_KERNARG=10 B

102ENV HSA_FORCE_FINE_GRAIN_PCIE=10 B

103ENV NVTE_FUSED_ATTN=10 B

104ENV NCCL_DEBUG=VERSION0 B

105ENV NVTE_CK_USES_BWD_V3=10 B

106ENV NVTE_CK_USES_FWD_V3=10 B

107ENV NVTE_CK_IS_V3_ATOMIC_FP32=00 B

108ENV NVTE_CK_HOW_V3_BF16_CVT=20 B

109ENV XLA_PYTHON_CLIENT_MEM_FRACTION=.970 B

110

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

111COPY jax_training/docker/google_forced_balance.patch /workspace/google_forced_balance.patch # buildkit1.09 KB

112COPY jax_training/docker/device-info.patch /workspace/device-info.patch # buildkit1.69 KB

113COPY jax_training/docker/max_to_keep.patch /workspace/max_to_keep.patch # buildkit1 KB

114COPY jax_training/docker/attention_mask.patch /workspace/attention_mask.patch # buildkit1.32 KB

115COPY jax_training/docker/quantization.patch /workspace/quantization.patch # buildkit452 B

116COPY jax_training/docker/hf_data_processing.patch /workspace/hf_data_processing.patch # buildkit895 B

117COPY jax_training/docker/CustomPackAndBatchOperation.py /workspace/CustomPackAndBatchOperation.py # buildkit2.87 KB

118

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

2.64 MB

119

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

396.67 MB

120

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

155 B

121

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

146 B

122

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

146 B

123

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

148 B

124

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

147 B

125

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

151 B

126

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

4.96 GB

127

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

32 B

128

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

131 B

129ARG GIT_COMMIT_TAG=16a78b68f79b8be8f0311cc7c47955303868d2380 B

130

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

184 B

131bash15.48 KB

132bash566.29 MB

133RUN /bin/sh -c rm -f258 B

134RUN /bin/sh -c rm -f389 B

Command

 ARG RELEASE

Help

System theme

Docker Suite

rocm/jax-training:maxtext-v25.11