|
| 1 | +ARG BASE_IMAGE=rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.12_pytorch_release_2.9.1 |
| 2 | + |
| 3 | +FROM ${BASE_IMAGE} |
| 4 | + |
| 5 | +# OS: |
| 6 | +# - Ubuntu: 24.04 |
| 7 | +# - Docker Client: 20.10.8 |
| 8 | +# ROCm: |
| 9 | +# - ROCm: 7.2 |
| 10 | +# Lib: |
| 11 | +# - torch: 2.9.1 |
| 12 | +# - rccl: release/rocm-rel-7.2 |
| 13 | +# - hipblaslt: release-staging/rocm-rel-7.2 |
| 14 | +# - rocblas: release-staging/rocm-rel-7.2 |
| 15 | +# - openmpi: 4.1.x |
| 16 | +# Intel: |
| 17 | +# - mlc: v3.12 |
| 18 | + |
| 19 | +LABEL maintainer="SuperBench" |
| 20 | + |
| 21 | +ENV DEBIAN_FRONTEND=noninteractive |
| 22 | +RUN apt-get update && \ |
| 23 | + apt-get -q install -y --no-install-recommends \ |
| 24 | + autoconf \ |
| 25 | + automake \ |
| 26 | + bc \ |
| 27 | + build-essential \ |
| 28 | + curl \ |
| 29 | + dmidecode \ |
| 30 | + git \ |
| 31 | + hipify-clang \ |
| 32 | + iproute2 \ |
| 33 | + jq \ |
| 34 | + libaio-dev \ |
| 35 | + libboost-program-options-dev \ |
| 36 | + libcap2 \ |
| 37 | + libcurl4-openssl-dev \ |
| 38 | + libnuma-dev \ |
| 39 | + libpci-dev \ |
| 40 | + libssl-dev \ |
| 41 | + libtinfo6 \ |
| 42 | + libtool \ |
| 43 | + lshw \ |
| 44 | + net-tools \ |
| 45 | + numactl \ |
| 46 | + openssh-client \ |
| 47 | + openssh-server \ |
| 48 | + pciutils \ |
| 49 | + python3-mpi4py \ |
| 50 | + rsync \ |
| 51 | + sudo \ |
| 52 | + util-linux \ |
| 53 | + vim \ |
| 54 | + wget \ |
| 55 | + && \ |
| 56 | + rm -rf /tmp/* |
| 57 | + |
| 58 | +ARG NUM_MAKE_JOBS=64 |
| 59 | + |
| 60 | +# Install CMake via apt if not already present (Ubuntu 24.04 provides >= 3.28) |
| 61 | +RUN if ! command -v cmake >/dev/null 2>&1; then \ |
| 62 | + apt-get update && apt-get install -y --no-install-recommends cmake; \ |
| 63 | + fi && \ |
| 64 | + echo "CMake version: $(cmake --version | head -1)" |
| 65 | + |
| 66 | +# Install Docker |
| 67 | +ENV DOCKER_VERSION=20.10.8 |
| 68 | +RUN cd /tmp && \ |
| 69 | + wget -q https://download.docker.com/linux/static/stable/x86_64/docker-${DOCKER_VERSION}.tgz -O docker.tgz && \ |
| 70 | + tar --extract --file docker.tgz --strip-components 1 --directory /usr/local/bin/ && \ |
| 71 | + rm docker.tgz |
| 72 | + |
| 73 | +# Update system config |
| 74 | +RUN mkdir -p /root/.ssh && \ |
| 75 | + touch /root/.ssh/authorized_keys && \ |
| 76 | + mkdir -p /var/run/sshd && \ |
| 77 | + sed -i "s/[# ]*PermitRootLogin prohibit-password/PermitRootLogin yes/" /etc/ssh/sshd_config && \ |
| 78 | + sed -i "s/[# ]*PermitUserEnvironment no/PermitUserEnvironment yes/" /etc/ssh/sshd_config && \ |
| 79 | + sed -i "s/[# ]*Port.*/Port 22/" /etc/ssh/sshd_config && \ |
| 80 | + echo "* soft nofile 1048576\n* hard nofile 1048576" >> /etc/security/limits.conf && \ |
| 81 | + echo "root soft nofile 1048576\nroot hard nofile 1048576" >> /etc/security/limits.conf |
| 82 | + |
| 83 | + |
| 84 | +# Set Ubuntu version |
| 85 | +ENV UBUNTU_VERSION=24.04 |
| 86 | + |
| 87 | +# Install OFED |
| 88 | +ENV OFED_VERSION=24.10-1.1.4.0 |
| 89 | +RUN if ! command -v ofed_info >/dev/null 2>&1; then \ |
| 90 | + echo "OFED not found. Installing OFED..."; \ |
| 91 | + cd /tmp && \ |
| 92 | + wget -q http://content.mellanox.com/ofed/MLNX_OFED-${OFED_VERSION}/MLNX_OFED_LINUX-${OFED_VERSION}-ubuntu${UBUNTU_VERSION}-x86_64.tgz && \ |
| 93 | + tar xzf MLNX_OFED_LINUX-${OFED_VERSION}-ubuntu${UBUNTU_VERSION}-x86_64.tgz && \ |
| 94 | + PATH=/usr/bin:${PATH} MLNX_OFED_LINUX-${OFED_VERSION}-ubuntu${UBUNTU_VERSION}-x86_64/mlnxofedinstall --user-space-only --without-fw-update --force --all && \ |
| 95 | + rm -rf MLNX_OFED_LINUX-${OFED_VERSION}* ; \ |
| 96 | + fi |
| 97 | + |
| 98 | +ENV ROCM_PATH=/opt/rocm |
| 99 | + |
| 100 | +# CMake 4.0+ dropped compat for cmake_minimum_required < 3.5. |
| 101 | +# Set globally so all subprojects (RCCL, mscclpp, etc.) work. |
| 102 | +ENV CMAKE_POLICY_VERSION_MINIMUM=3.5 |
| 103 | + |
| 104 | +# Install OpenMPI |
| 105 | +ENV OPENMPI_VERSION=4.1.x |
| 106 | +ENV MPI_HOME=/usr/local/mpi |
| 107 | +RUN cd /tmp && \ |
| 108 | + git clone --recursive https://github.com/open-mpi/ompi.git -b v${OPENMPI_VERSION} && \ |
| 109 | + cd ompi && \ |
| 110 | + ./autogen.pl && \ |
| 111 | + mkdir build && \ |
| 112 | + cd build && \ |
| 113 | + ../configure --prefix=/usr/local/mpi --enable-orterun-prefix-by-default --enable-mpirun-prefix-by-default --enable-prte-prefix-by-default --with-rocm=/opt/rocm && \ |
| 114 | + make -j $(nproc) && \ |
| 115 | + make -j $(nproc) install && \ |
| 116 | + ldconfig && \ |
| 117 | + cd / && \ |
| 118 | + rm -rf /tmp/openmpi-${OPENMPI_VERSION}* |
| 119 | + |
| 120 | +# Install Intel MLC |
| 121 | +RUN cd /tmp && \ |
| 122 | + wget -q https://downloadmirror.intel.com/866182/mlc_v3.12.tgz -O mlc.tgz && \ |
| 123 | + tar xzf mlc.tgz Linux/mlc && \ |
| 124 | + cp ./Linux/mlc /usr/local/bin/ && \ |
| 125 | + rm -rf ./Linux mlc.tgz |
| 126 | + |
| 127 | +# Install RCCL |
| 128 | +RUN cd /opt/ && \ |
| 129 | + git clone -b release/rocm-rel-7.2 https://github.com/ROCmSoftwarePlatform/rccl.git && \ |
| 130 | + cd rccl && \ |
| 131 | + mkdir build && \ |
| 132 | + cd build && \ |
| 133 | + CXX=/opt/rocm/bin/hipcc cmake -DHIP_COMPILER=clang -DCMAKE_BUILD_TYPE=Release -DCMAKE_VERBOSE_MAKEFILE=1 \ |
| 134 | + -DCMAKE_PREFIX_PATH="${ROCM_PATH}/hsa;${ROCM_PATH}/hip;${ROCM_PATH}/share/rocm/cmake/;${ROCM_PATH}" \ |
| 135 | + .. && \ |
| 136 | + make -j${NUM_MAKE_JOBS} |
| 137 | + |
| 138 | +# Install AMD SMI Python Library |
| 139 | +RUN apt install amd-smi-lib -y && \ |
| 140 | + cd /opt/rocm/share/amd_smi && \ |
| 141 | + python3 -m pip install . |
| 142 | + |
| 143 | +# Do NOT LD_PRELOAD librccl.so — causes segfaults on process exit due to |
| 144 | +# HIP static object teardown order. Use LD_LIBRARY_PATH instead. |
| 145 | +# Do NOT put /usr/lib/x86_64-linux-gnu/ before /opt/rocm/lib — OFED installs |
| 146 | +# an old libhsa-runtime64.so there that conflicts with ROCm's version. |
| 147 | +ENV PATH="/usr/local/mpi/bin:/opt/superbench/bin:/usr/local/bin/:/opt/rocm/hip/bin/:/opt/rocm/bin/:${PATH}" \ |
| 148 | + LD_LIBRARY_PATH="/opt/rccl/build:/usr/local/mpi/lib:/opt/rocm/lib:/usr/local/lib/:${LD_LIBRARY_PATH}" \ |
| 149 | + SB_HOME=/opt/superbench \ |
| 150 | + SB_MICRO_PATH=/opt/superbench \ |
| 151 | + ANSIBLE_DEPRECATION_WARNINGS=FALSE \ |
| 152 | + ANSIBLE_COLLECTIONS_PATH=/usr/share/ansible/collections |
| 153 | + |
| 154 | +RUN echo PATH="$PATH" > /etc/environment && \ |
| 155 | + echo LD_LIBRARY_PATH="$LD_LIBRARY_PATH" >> /etc/environment && \ |
| 156 | + echo SB_MICRO_PATH="$SB_MICRO_PATH" >> /etc/environment |
| 157 | + |
| 158 | +RUN apt install rocm-cmake -y && \ |
| 159 | + python3 -m pip install --upgrade pip wheel "setuptools>=69.0" |
| 160 | + |
| 161 | +WORKDIR ${SB_HOME} |
| 162 | + |
| 163 | +ADD third_party third_party |
| 164 | +# perftest_rocm6.patch changes are already upstream in the submodule version |
| 165 | +# rocm_megatron_lm: broken upstream (pretrain_deepseek.py missing in rocm_dev branch) |
| 166 | +# apex_rocm: skipped — all imports guarded, PyTorch 2.9 has native fused optimizers/AMP. |
| 167 | +RUN make RCCL_HOME=/opt/rccl/build/ ROCBLAS_BRANCH=release/rocm-rel-7.2 HIPBLASLT_BRANCH=release/rocm-rel-7.2 ROCM_VER=rocm-5.5.0 -C third_party rocm -o cpu_hpl -o cpu_stream -o megatron_lm -o rocm_hipblaslt -o rocm_megatron_lm -o apex_rocm |
| 168 | +# Build hipblaslt-bench using our standalone CMake. The upstream 7.2 build |
| 169 | +# system pulls in AMD-internal "origami" headers and a new tensilelite-host |
| 170 | +# C++ library that is broken under our system (link to system libhipblaslt |
| 171 | +# at runtime instead of building the full library). Our standalone CMake |
| 172 | +# compiles only the bench sources against system hipblaslt + LAPACK. |
| 173 | +COPY dockerfile/etc/hipblaslt-bench-standalone.cmake /tmp/hipblaslt-bench-standalone.cmake |
| 174 | +RUN cd third_party && \ |
| 175 | + git clone --depth 1 -b release/rocm-rel-7.2 https://github.com/ROCmSoftwarePlatform/hipBLASLt.git && \ |
| 176 | + cp /tmp/hipblaslt-bench-standalone.cmake hipBLASLt/CMakeLists.txt && \ |
| 177 | + cd hipBLASLt && \ |
| 178 | + # Pre-build cblas/lapack into /usr/local using the upstream deps script. |
| 179 | + # The deps superbuild builds but does not install lapack; install both |
| 180 | + # gtest and lapack explicitly so find_package(BLAS)/find_package(LAPACK) |
| 181 | + # can locate /usr/local/lib/{liblapack.a,libcblas.a,libblas.a}. |
| 182 | + mkdir -p deps/build && cd deps/build && \ |
| 183 | + CMAKE_POLICY_VERSION_MINIMUM=3.5 cmake .. && \ |
| 184 | + cmake --build . -j$(nproc) --target lapack && \ |
| 185 | + cmake --build lapack/src/lapack-build -j$(nproc) --target install && \ |
| 186 | + cd ../.. && \ |
| 187 | + mkdir -p build && cd build && \ |
| 188 | + CMAKE_POLICY_VERSION_MINIMUM= cmake \ |
| 189 | + -DCMAKE_CXX_COMPILER=/opt/rocm/llvm/bin/clang++ \ |
| 190 | + -DCMAKE_HIP_COMPILER=/opt/rocm/llvm/bin/clang++ \ |
| 191 | + -DCMAKE_HIP_ARCHITECTURES=gfx942 \ |
| 192 | + -DCMAKE_PREFIX_PATH="/opt/rocm;/usr/local" \ |
| 193 | + -DBLAS_LIBRARIES=/usr/local/lib/libblas.a \ |
| 194 | + -DLAPACK_LIBRARIES=/usr/local/lib/liblapack.a \ |
| 195 | + -DCMAKE_BUILD_TYPE=Release \ |
| 196 | + .. && \ |
| 197 | + make -j$(nproc) hipblaslt-bench && \ |
| 198 | + cp -v hipblaslt-bench /opt/superbench/bin/ |
| 199 | +RUN cd third_party/Megatron/Megatron-DeepSpeed && \ |
| 200 | + git apply ../megatron_deepspeed_rocm6.patch |
| 201 | + |
| 202 | +# Install TransformerEngine — use AOTriton-only fused attention to avoid the |
| 203 | +# CK + AITER build chain (validated working on ROCm 7.0; same TE upstream). |
| 204 | +# onnxscript is now an unconditional import in TE main (export.py); install it |
| 205 | +# explicitly because TE's setup.py does not list it as a dependency. |
| 206 | +RUN python3 -m pip install onnxscript && \ |
| 207 | + git clone --recursive https://github.com/ROCm/TransformerEngine.git && \ |
| 208 | + cd TransformerEngine && \ |
| 209 | + NVTE_FRAMEWORK=pytorch \ |
| 210 | + NVTE_FUSED_ATTN_CK=0 \ |
| 211 | + NVTE_FUSED_ATTN_AOTRITON=1 \ |
| 212 | + NVTE_ROCM_ARCH="gfx942;gfx950" \ |
| 213 | + python3 setup.py install |
| 214 | +RUN python3 -c "import transformer_engine.pytorch; print('TE installed successfully')" |
| 215 | + |
| 216 | +ADD . . |
| 217 | +ENV USE_HIP_DATATYPE=1 |
| 218 | +ENV USE_HIPBLAS_COMPUTETYPE=1 |
| 219 | +RUN python3 -m pip install .[amdworker] && \ |
| 220 | + CXX=/opt/rocm/bin/hipcc make cppbuild && \ |
| 221 | + make postinstall |
| 222 | + |
| 223 | +# Fix stale hypothesis plugin from base image (imports removed pkg_resources) |
| 224 | +# and add test dependencies missing from the base image. |
| 225 | +RUN python3 -m pip install --upgrade hypothesis setuptools pytest-timeout vcrpy |
0 commit comments