Deploy your PyTorch app on the Framework desktop with ROCm in a Docker container¶

We, the NOLAI Tech Team, have bought 4 framework desktops. They are great since they have such a large amount of (shared) RAM and a powerful AMD GPU. But that same AMD GPU is also a bit of a headache when it comes to running PyTorch models. The CUDA ecosystem is so dominant that most PyTorch images and packages assume you have an NVIDIA GPU. But luckily AMD has a solution: AMD gives you a Docker image with ROCm and a version of PyTorch that should work as a drop-in replacement for the CUDA build.

The reality is that it is not so easy and you will probably fall in a series of traps. These traps are individually small and collectively very annoying. This post documents the ones we fell in so you don't have to find them yourself.

Before you start: AMD Container Toolkit¶

If you want Docker containers to actually see the GPU, you can do so with the AMD Container Toolkit. Think of it as the AMD equivalent of NVIDIA's container toolkit: it hooks into the Docker runtime and handles the GPU passthrough for you.

Once it's in place, you get runtime: amd in your Compose files and AMD_VISIBLE_DEVICES to control which GPUs are visible. Without it, the device flags in your docker run command do nothing useful. The Quick Start Guide gets you there in a few steps.

The base image¶

Everything starts from AMD's official ROCm PyTorch image:

FROM rocm/pytorch:rocm7.2.1_ubuntu22.04_py3.10_pytorch_release_2.9.1

This image ships with PyTorch, Triton, and the full ROCm stack already installed inside a virtual environment at /opt/venv. That is great, except it is also your first trap.

Trap 1: Don't bring your own PyTorch¶

If your project has a requirements.txt, it almost certainly lists torch, torchvision, or torchaudio. Installing those from PyPI will pull in the CUDA-flavored versions and silently overwrite the ROCm build.

Strip them out before installing anything:

RUN grep -E -v '^(torch|torchvision|torchaudio|triton)([=><\[@]|$)|^nvidia-' \
    requirements.txt > clean_reqs.txt

Same goes for nvidia-* packages. They will install fine and do absolutely nothing useful on an AMD GPU.

Trap 2: Use the right pip¶

The ROCm image does not use the system Python. It uses a virtual environment. If you call pip or python directly, you are talking to the system interpreter, not the one with ROCm PyTorch in it. Always use:

RUN /opt/venv/bin/pip install --no-cache-dir -r clean_reqs.txt

Trap 3: Enable experimental AOTriton¶

Add this environment variable:

ENV TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1

Without it, certain attention kernels will either fall back to a slow path or fail outright. It is experimental in name only at this point. Just turn it on and you will be fine.

Putting it together: a minimal app¶

For a straightforward FastAPI or similar app, the Dockerfile ends up looking like this:

FROM rocm/pytorch:rocm7.2.1_ubuntu22.04_py3.10_pytorch_release_2.9.1

WORKDIR /app
ENV TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1

COPY requirements.txt ./
RUN grep -E -v '^(torch|torchvision|torchaudio|triton)([=><\[@]|$)|^nvidia-' \
    requirements.txt > clean_reqs.txt && \
    /opt/venv/bin/pip install --no-cache-dir -r clean_reqs.txt

COPY . .
EXPOSE 8000
CMD ["/opt/venv/bin/uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

That covers maybe 80% of projects. If yours needs Flash Attention or CTranslate2 keep reading.

When it gets harder: Whisper on ROCm¶

We also run Whisper-WebUI on our Framework machines. That project adds three extra layers of pain.

CTranslate2. The PyPI release is CUDA-only. For ROCm you need to grab the wheel directly from the GitHub releases:

RUN curl -sL https://github.com/OpenNMT/CTranslate2/releases/download/v4.7.1/rocm-python-wheels-Linux.zip | \
    python3 -c "import sys,zipfile,io; z=zipfile.ZipFile(io.BytesIO(sys.stdin.buffer.read())); \
    [z.extract(n,'/tmp') for n in z.namelist() if 'cp310' in n]" && \
    /opt/venv/bin/pip install --no-cache-dir --force-reinstall /tmp/temp-linux/ctranslate2-*.whl

Not pretty. But it works.

Flash Attention 2. The CUDA wheel obviously does not apply here. For AMD you build from source using Triton, which is already bundled in the ROCm image:

RUN /opt/venv/bin/pip install --no-cache-dir ninja && \
    git clone --depth=1 https://github.com/Dao-AILab/flash-attention.git /tmp/flash-attention && \
    cd /tmp/flash-attention && \
    FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" /opt/venv/bin/python setup.py install && \
    rm -rf /tmp/flash-attention

The build takes a few minutes, but it is a one-time cost.

Running it with Docker Compose¶

With the AMD Container Toolkit in place, the Compose setup is clean. runtime: amd replaces the manual device flags, and AMD_VISIBLE_DEVICES=all handles GPU exposure:

services:
  whisper-webui:
    image: registry.science.ru.nl/nolai/ops/whisper-webui-rocm/whisper-webui:main
    container_name: whisper-webui
    restart: unless-stopped
    runtime: amd
    environment:
      - AMD_VISIBLE_DEVICES=all
      - TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
      - FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE
    volumes:
      - ./models:/Whisper-WebUI/models
      - ./outputs:/Whisper-WebUI/outputs
      - ./configs:/Whisper-WebUI/configs
    entrypoint: ["python", "app.py", "--server_port", "7860", "--server_name", "0.0.0.0"]

Is it worth it?¶

Mostly yes. Once the container builds, it runs reliably. PyTorch operations that expect CUDA work with the ROCm drop-in replacement, and the Framework desktops have enough memory to run reasonably large models locally without sending anything to a cloud.

The friction is almost entirely at build time. The traps above are non-obvious, the error messages are not always helpful, and the documentation assumes CUDA more often than not. But they are all solvable, and this post should get you past most of them.

Having a framework desktop with such a large amount of (V)RAM is a huge advantage for local experimentation. We can now easily finetune large models, run whatever model we want, and do it all without worrying about cloud costs.