ROCM Setup Steps

build-pytorch.sh

#!/bin/sh

set -e

mkdir -p pytorch-examples pytorch-cache
docker build pytorch-examples \
     --rm \
     -t rocm-examples-pytorch \
     -f post-rocm-python-ubuntu.Dockerfile

build-rocm.sh

#!/bin/sh

set -e

LOCAL_GID=$(getent group render  | cut -d: -f3)

mkdir -p rocm-examples
exec docker build rocm-examples \
        --build-arg GID="${LOCAL_GID}" \
        --rm \
        -t rocm-examples \
        -f hip-libraries-rocm-ubuntu.Dockerfile

hip-libraries-rocm-ubuntu.Dockerfile

# Ubuntu based docker image
FROM ubuntu:20.04

# Base packages that are required for the installation
RUN export DEBIAN_FRONTEND=noninteractive; \
    apt-get update -qq \
    && apt-get install --no-install-recommends -y \
        ca-certificates \
        git \
        locales-all \
        make \
        python3 \
        python3-venv \
        python3-dev \
        ssh \
        sudo \
        wget \
        pkg-config \
        glslang-tools \
        libvulkan-dev \
        vulkan-validationlayers \
        libglfw3-dev \
        neovim \
    && rm -rf /var/lib/apt/lists/*

ENV LANG en_US.utf8

# Install ROCM HIP and libraries using the installer script
RUN export DEBIAN_FRONTEND=noninteractive; \
    wget https://repo.radeon.com/amdgpu-install/5.4.3/ubuntu/focal/amdgpu-install_5.4.50403-1_all.deb \
    && apt-get update -qq \
    && apt-get install -y ./amdgpu-install_5.4.50403-1_all.deb \
    && rm ./amdgpu-install_5.4.50403-1_all.deb \
    && amdgpu-install -y --usecase=hiplibsdk --no-dkms \
    && apt-get install -y libnuma-dev \
    && rm -rf /var/lib/apt/lists/*

# Install CMake
RUN wget https://github.com/Kitware/CMake/releases/download/v3.21.7/cmake-3.21.7-linux-x86_64.sh \
    && mkdir /cmake \
    && sh cmake-3.21.7-linux-x86_64.sh --skip-license --prefix=/cmake \
    && rm cmake-3.21.7-linux-x86_64.sh

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

RUN echo "/opt/rocm/lib" >> /etc/ld.so.conf.d/rocm.conf \
    && ldconfig

# Use render group as an argument from user
ARG GID=109

# Add the render group and a user with sudo permissions for the container
RUN groupadd --system --gid ${GID} render \
    && useradd -Um -G sudo,video,render developer \
    && echo developer ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/developer \
    && chmod 0440 /etc/sudoers.d/developer

RUN mkdir /workspaces && chown developer:developer /workspaces
WORKDIR /workspaces
VOLUME /workspaces

USER developer

launch-pytorch.sh

#!/bin/sh
# shellcheck disable=SC2086

set -e

DOCKER_ARGS=${DOCKER_ARGS:--it}
CMD=${1:-bash}
exec docker run \
        --rm \
        ${DOCKER_ARGS} \
        --name rocm-examples-pytorch \
        -h rocm-examples-pytorch \
        --device /dev/kfd --device /dev/dri \
        -v "$(pwd -P)/pytorch-examples":/workspaces/pytorch-examples \
        -v "$(pwd -P)/pytorch-cache":/home/developer/.cache \
        rocm-examples-pytorch "$CMD"

launch-rocm.sh

#!/bin/sh
# shellcheck disable=SC2086

set -e

DOCKER_ARGS=${DOCKER_ARGS:--it}
CMD=${1:-bash}
exec docker run \
        --rm \
        ${DOCKER_ARGS} \
        --name rocm-examples \
        -h rocm-examples \
        --device /dev/kfd --device /dev/dri \
        -v "$(pwd -P)/rocm-examples":/workspaces/rocm-examples \
        rocm-examples "$CMD"

post-rocm-python-ubuntu.Dockerfile

# ROCm based docker image
FROM rocm-examples:latest

# Set Root User
USER root

# Create VEnv Directory
RUN mkdir -p /venv && chown developer:developer /venv

# Set User
USER developer

# Install VEnv and PyTorch
RUN  python3.8 -m venv /venv && \
     . /venv/bin/activate && \
     python3.8 -m pip install --upgrade \
         pip setuptools wheel six && \
     python3.8 -m pip install \
        --index-url https://download.pytorch.org/whl/rocm5.4.2 \
        --pre torch torchvision torchaudio pillow && \
    python3.8 -m pip cache purge

VOLUME /venv

CMD /venv/bin/python3.8

ROCM Setup Steps.md

ROCm Setup Steps

ROCm Docker Image

1. Install Docker and ensure you can run docker ps correctly, add yourself to the docker group if necessary.
2. Run build-rocm.sh builds a ROCm docker image for your Linux System. It is configured to use the render group configured in your Linux distribution and ensure that /dev/kfd and /dev/dri are writeable by the render group users. Ensure you are added to the render group if necessary.
3. Run launch-rocm.sh if you wish to only use ROCm with the docker image you built.

PyTorch Docker Image

1. Run build-pytorch.sh if you wish to build a PyTorch image for your Linux System. It is built as a separate docker image, on top of the ROCm docker image you built earlier.
2. Run launch-pytorch.sh if you wish to run PyTorch with the second image just now built.
3. Run source /venv/bin/activate within the container and you should be able to run all the cool PyTorch things you need.