Sign inSign up

genesu100/jax_test:jax_rocm7.0_jax_0.6.2_20251106

Manifest digest

sha256:8740fda4365f87b1e0f4561326c724d026706b7288b0f8ec8ad9aca46fa04e47

OS/ARCH

linux/amd64

Compressed size

18.65 GB

Last pushed

11 months by genesu100

Type

Image

Manifest digest

sha256:8740fda4365f87b1e0f4561326c724d026706b7288b0f8ec8ad9aca46fa04e47

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 update177.47 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=6.2.00 B
12ARG ROCM_PATH=/opt/rocm-7.0.20 B
13ENV ROCM_PATH=/opt/rocm-7.0.20 B
14ARG ROCM_BUILD_JOB0 B
15ARG ROCM_BUILD_NUM0 B
16RUN |5 GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a5.69 GB
17ENV HCC_HOME=/opt/rocm-7.0.2/hcc0 B
18ENV HIP_PATH=/opt/rocm-7.0.2/0 B
19ENV OPENCL_ROOT=/opt/rocm-7.0.2/opencl0 B
20ENV PATH=/opt/rocm-7.0.2/hcc/bin:/opt/rocm-7.0.2//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
21ENV PATH=/opt/rocm-7.0.2/bin:/opt/rocm-7.0.2/hcc/bin:/opt/rocm-7.0.2//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
22ENV PATH=/opt/rocm-7.0.2/opencl/bin:/opt/rocm-7.0.2/bin:/opt/rocm-7.0.2/hcc/bin:/opt/rocm-7.0.2//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
23ENV PATH=/root/bin:/root/.local/bin:/opt/rocm-7.0.2/opencl/bin:/opt/rocm-7.0.2/bin:/opt/rocm-7.0.2/hcc/bin:/opt/rocm-7.0.2//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
24RUN |5 GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a98.21 MB
25RUN |5 GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a97.33 MB
26COPY jax_rocm_plugin/third_party/jax/namespace.patch /tmp/namespace.patch # buildkit439 B
27RUN |5 GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a1.69 KB
28ARG JAX_VERSION0 B
29ARG JAX_COMMIT0 B
30ARG XLA_COMMIT0 B
31LABEL com.amdgpu.rocm_version=7.0.2 com.amdgpu.python_version=3.12 com.amdgpu.jax_version= com.amdgpu.jax_commit=0 B
32RUN |8 GPU_DEVICE_TARGETS=gfx906 gfx908 gfx90a131.36 MB
33/bin/bash195.73 MB
34ENV ROCM_PATH=/opt/rocm-7.0.20 B
35ENV PATH=/opt/rocm-7.0.2/opencl/bin:/opt/rocm-7.0.2/bin:/opt/rocm-7.0.2/hcc/bin:/opt/rocm-7.0.2//bin:/root/bin:/root/.local/bin:/opt/rocm-7.0.2/opencl/bin:/opt/rocm-7.0.2/bin:/opt/rocm-7.0.2/hcc/bin:/opt/rocm-7.0.2//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
36RUN /bin/sh -c apt update1.6 GB
37RUN /bin/sh -c python -m4 MB
38ENV PATH=/opt/venv/bin:/opt/rocm-7.0.2/opencl/bin:/opt/rocm-7.0.2/bin:/opt/rocm-7.0.2/hcc/bin:/opt/rocm-7.0.2//bin:/root/bin:/root/.local/bin:/opt/rocm-7.0.2/opencl/bin:/opt/rocm-7.0.2/bin:/opt/rocm-7.0.2/hcc/bin:/opt/rocm-7.0.2//bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
39RUN /bin/sh -c pip install5.13 MB
40RUN /bin/sh -c pip uninstall259.82 KB
41RUN /bin/sh -c pip install52.73 MB
42RUN /bin/sh -c cd /opt/rocm/share/amd_smi2.87 MB
43ARG HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git0 B
44ARG HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff090 B
45ENV PYTORCH_ROCM_ARCH=gfx942;gfx9500 B
46ENV Python_ROOT=/pyenv0 B
47RUN |2 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 /bin/sh167 B
48RUN |2 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 /bin/sh1.4 GB
49RUN |2 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 /bin/sh186 B
50RUN |2 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 /bin/sh165 B
51RUN |2 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 /bin/sh178.11 KB
52RUN |2 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 /bin/sh178.11 KB
53RUN |2 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 /bin/sh164 B
54WORKDIR /workspace99 B
55ARG MAXTEXT_REPO=https://github.com/ROCm/maxtext.git0 B
56ARG MAXTEXT_BRANCH=genesu/add-jax-training-patches0 B
57RUN |4 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git1.54 GB
58ARG JAX_VERSION=0.6.20 B
59ENV JAX_BRANCH=ae78eb36bcadd26bdacd62bca8b2045809cc93820 B
60ENV XLA_BRANCH=cfce2f10011aab8fc85a8419c2b3fd90454397770 B
61WORKDIR /opt32 B
62RUN |5 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git323.89 MB
63RUN |5 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git2.98 GB
64RUN |5 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git183.15 MB
65WORKDIR /workspace32 B
66ARG TE_REPO=https://github.com/ROCm/TransformerEngine.git0 B
67ARG TE_BRANCH=533bee5ae1e3be462a5a6ae9a5e97cc7d7327f670 B
68ENV USE_ROCM=10 B
69ENV HIP_PATH=/opt/rocm0 B
70ENV NVTE_FRAMEWORK=jax0 B
71ENV CMAKE_BUILD_PARALLEL_LEVEL=640 B
72ENV NVTE_ROCM_ARCH=gfx942;gfx9500 B
73ENV NVTE_USE_ROCM=10 B
74ENV NVTE_FUSED_ATTN_AOTRITON=00 B
75ENV PYTHONPATH=/home/TransformerEngine/3rdparty/hipify_torch0 B
76RUN |7 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git3.8 GB
77ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B
78ENV NVTE_USE_HIPBLASLT=10 B
79ENV GPU_MAX_HW_QUEUES=20 B
80ENV HIP_FORCE_DEV_KERNARG=10 B
81ENV HSA_FORCE_FINE_GRAIN_PCIE=10 B
82ENV NVTE_FUSED_ATTN=10 B
83ENV NCCL_DEBUG=VERSION0 B
84ENV NVTE_CK_USES_BWD_V3=10 B
85ENV NVTE_CK_USES_FWD_V3=10 B
86ENV NVTE_CK_IS_V3_ATOMIC_FP32=00 B
87ENV NVTE_CK_HOW_V3_BF16_CVT=20 B
88ENV NVTE_FUSED_ATTN_CK=10 B
89ENV NVTE_FUSED_ATTN_AOTRITON=00 B
90ENV XLA_PYTHON_CLIENT_MEM_FRACTION=.970 B
91ENV 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=85899345920 B
92COPY jax_training/docker/google_forced_balance.patch /workspace/google_forced_balance.patch # buildkit1.09 KB
93COPY jax_training/docker/device-info.patch /workspace/device-info.patch # buildkit1.7 KB
94COPY jax_training/docker/max_to_keep.patch /workspace/max_to_keep.patch # buildkit1 KB
95COPY jax_training/docker/attention_mask.patch /workspace/attention_mask.patch # buildkit1.32 KB
96COPY jax_training/docker/quantization.patch /workspace/quantization.patch # buildkit453 B
97COPY jax_training/docker/hf_data_processing.patch /workspace/hf_data_processing.patch # buildkit895 B
98COPY jax_training/docker/CustomPackAndBatchOperation.py /workspace/CustomPackAndBatchOperation.py # buildkit2.87 KB
99RUN |7 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git1.49 MB
100RUN |7 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git375.36 MB
101RUN |7 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git155 B
102RUN |7 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git146 B
103RUN |7 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git146 B
104RUN |7 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git148 B
105RUN |7 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git147 B
106RUN |7 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=be40066e77bb6a1f63cc2c7031e55a3a0e7dff09 MAXTEXT_REPO=https://github.com/ROCm/maxtext.git152 B

Command

 ARG RELEASE