triton-torch-hip.Dockerfile

# Build: docker build -f triton-torch-hip.Dockerfile -t triton-torch-hip .
# Run:   docker run -it --rm --device /dev/kfd --device /dev/dri triton-torch-hip

FROM ubuntu:22.04

ARG INSTALL_TORCH=TRUE
ARG TORCH_VERSION=2.4.0.dev20240530
ARG ROCM_VERSION=6.1.1

# Setup ROCm package signing key
RUN apt-get update && apt-get install -y curl wget gnupg2 && \
    curl https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor | tee /etc/apt/keyrings/rocm.gpg > /dev/null

# Register amdgpu kernel-mode driver
RUN sh -c "echo deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/amdgpu/${ROCM_VERSION}/ubuntu jammy main > /etc/apt/sources.list.d/amdgpu.list" && \
    sh -c "echo deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${ROCM_VERSION} jammy main > /etc/apt/sources.list.d/rocm.list" && \
    sh -c "echo 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' > /etc/apt/preferences.d/rocm-pin-600"

# Install base development packages
RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \
      git vim llvm clang lld ccache clang-format less \
      python3 python3-dev python3-pip

# ~1.3GB till here

# Install ROCm information tools
#
# Note that we only install the base system information tools here;
# The full ROCm stack, even the HIP runtime, is not installed--they
# are packaged into the PyTorch pip packages and we use from there.
#
# To install the HIP runtime, add the hip-runtime-amd package here.
# It pulls in compiler packages like rocm-llvm and hipcc.
RUN DEBIAN_FRONTEND=noninteractive apt-get install -y \
      rocminfo amd-smi-lib rocm-smi-lib && \
    apt-get clean && rm -rf /var/lib/apt/lists/*

# ~1.3GB till here

# Install Triton development pip packages
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 300 && \
    pip install --upgrade pip && \
    pip install numpy scipy pandas matplotlib cmake ninja wheel lit pytest pytest-xdist

# Install torch pip packages
RUN if [ "${INSTALL_TORCH}" = "TRUE" ]; then \
    pip install --pre torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/nightly/rocm6.1; \
    fi

# ~20GB till here with torch

ENV PATH="${PATH}:/opt/rocm/bin"