Sign inSign up

rocm/jax-training:maxtext-v25.9

Manifest digest

sha256:4bb16ab58279ef09cb7a5e362c38e3fe3f901de44d8dbac5d0cb3bac5686441e

OS/ARCH

linux/amd64

Compressed size

16.84 GB

Last pushed

10 months by rocmdev

Type

Image

Manifest digest

sha256:4bb16ab58279ef09cb7a5e362c38e3fe3f901de44d8dbac5d0cb3bac5686441e

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=22.040 B
5ADD file ... in / 28.17 MB
6CMD ["/bin/bash"]0 B
7RUN /bin/sh -c apt-get update54.18 MB
8RUN /bin/sh -c apt-get update11.93 MB
9ARG GPU_DEVICE_TARGETS=gfx900 gfx906 gfx908 gfx90a0 B
10ENV GPU_DEVICE_TARGETS=gfx900 gfx906 gfx908 gfx90a0 B
11ARG ROCM_VERSION=6.0.00 B
12ARG ROCM_PATH=/opt/rocm-7.00 B
13ENV ROCM_PATH=/opt/rocm-7.00 B
14ARG ROCM_BUILD_JOB0 B
15ARG ROCM_BUILD_NUM0 B
16RUN |5 GPU_DEVICE_TARGETS=gfx900 gfx906 gfx9085.07 GB
17ENV PATH=/opt/rocm-7.0/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
18ENV PATH=/root/bin:/root/.local/bin:/opt/rocm-7.0/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
19RUN |5 GPU_DEVICE_TARGETS=gfx900 gfx906 gfx90812.89 MB
20ARG PYTHON_VERSION=3.10.140 B
21RUN |6 GPU_DEVICE_TARGETS=gfx900 gfx906 gfx9087.66 MB
22ENV PYENV_ROOT=/pyenv0 B
23ENV PATH=/pyenv/shims:/pyenv/bin:/root/bin:/root/.local/bin:/opt/rocm-7.0/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
24RUN |6 GPU_DEVICE_TARGETS=gfx900 gfx906 gfx90855.75 MB
25RUN |6 GPU_DEVICE_TARGETS=gfx900 gfx906 gfx90898.11 MB
26ARG JAX_VERSION0 B
27ARG JAX_COMMIT0 B
28ARG XLA_COMMIT0 B
29ARG JAX_USE_CLANG0 B
30LABEL com.amdgpu.rocm_version=7.0 com.amdgpu.python_version=3.10 com.amdgpu.jax_version=0.6.0 com.amdgpu.jax_commit=0 B
31RUN |4 JAX_VERSION=0.6.0 JAX_COMMIT= XLA_COMMIT=105 B
32RUN |4 JAX_VERSION=0.6.0 JAX_COMMIT= XLA_COMMIT=257.67 MB
33ENV ROCM_PATH=/opt/rocm-7.0.00 B
34ENV PATH=/opt/rocm-7.0.0/opencl/bin:/opt/rocm-7.0.0/bin:/opt/rocm-7.0.0/hcc/bin:/opt/rocm-7.0.0//bin:/pyenv/shims:/pyenv/bin:/root/bin:/root/.local/bin:/opt/rocm-7.0/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin0 B
35RUN /bin/sh -c apt update26.15 MB
36RUN /bin/sh -c pip install51.89 MB
37ARG HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git0 B
38ARG HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad0 B
39ENV PYTORCH_ROCM_ARCH=gfx942;gfx9500 B
40ENV Python_ROOT=/pyenv0 B
41RUN |2 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad /bin/sh2.7 GB
42RUN |2 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad /bin/sh182 B
43COPY /opt/MIOpen/build/*.deb build/ # buildkit217.39 MB
44RUN |2 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad /bin/sh199.46 KB
45RUN |2 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad /bin/sh302.89 MB
46RUN |2 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad /bin/sh197 B
47ENV JAX_BRANCH=rocm-jaxlib-v0.6.00 B
48ENV XLA_BRANCH=rocm-jaxlib-v0.6.00 B
49ENV JAX_COMMIT=4528432beaa89bc33268b7fbf27410f37752909e0 B
50ENV XLA_COMMIT=fe042519c9af6112d50a46670301ef9a53ba3bde0 B
51WORKDIR /opt32 B
52RUN |2 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad /bin/sh312 MB
53RUN |2 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad /bin/sh2.98 GB
54WORKDIR /workspace99 B
55ARG TE_REPO=https://github.com/ROCm/TransformerEngine.git0 B
56ARG TE_BRANCH=c91bac54e9ff4992cdf2dbee9218d039986ebcd70 B
57ENV USE_ROCM=10 B
58ENV HIP_PATH=/opt/rocm0 B
59ENV NVTE_FRAMEWORK=jax0 B
60ENV CMAKE_BUILD_PARALLEL_LEVEL=640 B
61ENV NVTE_ROCM_ARCH=gfx942;gfx9500 B
62ENV NVTE_USE_ROCM=10 B
63ENV NVTE_FUSED_ATTN_AOTRITON=00 B
64ENV PYTHONPATH=/home/TransformerEngine/3rdparty/hipify_torch0 B
65RUN |4 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad TE_REPO=https://github.com/ROCm/TransformerEngine.git2.48 GB
66ARG JAX_VERSION=0.6.20 B
67COPY jax_training/docker/google_forced_balance.patch /workspace/google_forced_balance.patch # buildkit1.09 KB
68COPY jax_training/docker/device-info.patch /workspace/device-info.patch # buildkit1.69 KB
69COPY jax_training/docker/max_to_keep.patch /workspace/max_to_keep.patch # buildkit1 KB
70COPY jax_training/docker/attention_mask.patch /workspace/attention_mask.patch # buildkit1.32 KB
71COPY jax_training/docker/quantization.patch /workspace/quantization.patch # buildkit451 B
72COPY jax_training/docker/hf_data_processing.patch /workspace/hf_data_processing.patch # buildkit894 B
73COPY jax_training/docker/CustomPackAndBatchOperation.py /workspace/CustomPackAndBatchOperation.py # buildkit2.87 KB
74RUN |5 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad TE_REPO=https://github.com/ROCm/TransformerEngine.git2.2 GB
75RUN |5 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad TE_REPO=https://github.com/ROCm/TransformerEngine.git129.48 KB
76RUN |5 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad TE_REPO=https://github.com/ROCm/TransformerEngine.git154 B
77RUN |5 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad TE_REPO=https://github.com/ROCm/TransformerEngine.git146 B
78RUN |5 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad TE_REPO=https://github.com/ROCm/TransformerEngine.git146 B
79RUN |5 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad TE_REPO=https://github.com/ROCm/TransformerEngine.git148 B
80RUN |5 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad TE_REPO=https://github.com/ROCm/TransformerEngine.git147 B
81RUN |5 HIPBLASLT_REPO=https://github.com/ROCm/rocm-libraries.git HIPBLASLT_BRANCH=d24f36139c99a91d8cdd15680f3818dc416a99ad TE_REPO=https://github.com/ROCm/TransformerEngine.git152 B
82ENV NVTE_ALLOW_NONDETERMINISTIC_ALGO=10 B
83ENV NVTE_USE_HIPBLASLT=10 B
84ENV GPU_MAX_HW_QUEUES=20 B
85ENV HIP_FORCE_DEV_KERNARG=10 B
86ENV HSA_FORCE_FINE_GRAIN_PCIE=10 B
87ENV NVTE_FUSED_ATTN=10 B
88ENV NCCL_DEBUG=VERSION0 B
89ENV NVTE_CK_USES_BWD_V3=10 B
90ENV NVTE_CK_USES_FWD_V3=10 B
91ENV NVTE_CK_IS_V3_ATOMIC_FP32=00 B
92ENV NVTE_CK_HOW_V3_BF16_CVT=20 B
93ENV NVTE_FUSED_ATTN_CK=10 B
94ENV NVTE_FUSED_ATTN_AOTRITON=00 B
95ENV XLA_PYTHON_CLIENT_MEM_FRACTION=.970 B
96ENV XLA_FLAGS=--xla_gpu_memory_limit_slop_factor=95 --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592 --xla_gpu_graph_level=0 --xla_gpu_enable_latency_hiding_scheduler=True0 B

Command

 ARG RELEASE