From cabef1bfe091a8e280ae5acc356869b3bcf041a8 Mon Sep 17 00:00:00 2001
From: Ankush <105823120+JustInCache@users.noreply.github.com>
Date: Tue, 30 Jun 2026 06:39:55 +0530
Subject: [PATCH] fix(docker): add ROCm GPU support via compose overlay (#630)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* fix(docker): add ROCm GPU support via compose overlay

Fixes #618. The Docker image installs CPU-only PyTorch from PyPI by
default, so even when users correctly pass /dev/kfd and /dev/dri device
nodes into the container, torch.cuda.is_available() returns False and
the GPU is reported as "None (CPU only)".

Changes:
- Dockerfile: add PYTORCH_VARIANT build arg (default: cpu). When set to
  "rocm", the ROCm-enabled PyTorch wheels are installed from the
  pytorch.org/whl/rocm6.3 index before requirements.txt runs, so pip
  sees the ROCm build as already satisfying the torch>=2.2.0 constraint
  and does not overwrite it with the CPU wheel. The render and video
  groups are created with parameterised GIDs (RENDER_GID / VIDEO_GID,
  defaulting to Ubuntu 22.04 values) and the voicebox user is added to
  both groups so it can open /dev/kfd and /dev/dri.

- docker-compose.rocm.yml: new compose overlay that wires everything
  together — PYTORCH_VARIANT=rocm build arg, /dev/kfd + /dev/dri device
  passthrough, group_add for render/video, HSA_OVERRIDE_GFX_VERSION
  (defaults to 11.0.0 for RDNA3/Strix Halo with a comment listing
  values for RDNA2/RDNA1/Vega), and PYTORCH_HIP_ALLOC_CONF for the
  memory allocator. Usage:
    docker compose -f docker-compose.yml -f docker-compose.rocm.yml up --build

- docker-compose.yml: add a comment pointing to the ROCm overlay.

The CPU default path is unchanged — no extra build time, no size increase.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(docker): address review comments on ROCm overlay

Two issues raised in PR review:

1. CodeRabbit: `docker compose up --build-arg` is not supported by the
   `up` subcommand. Replaced the GID override instructions with the
   correct env-var export pattern. Added RENDER_GID and VIDEO_GID to
   `build.args` using ${VAR:-default} interpolation so a single export
   covers both the Dockerfile group creation and the runtime group_add.
   Changed group_add entries from hardcoded strings to the same
   interpolated vars so host GIDs stay in sync end-to-end.

2. @Xarianne: ROCm 6.3 does not support RDNA 4 (RX 9000 series) cards.
   Added a ROCM_VERSION build arg (default 6.3) to both the Dockerfile
   and docker-compose.rocm.yml so users can set ROCM_VERSION=7.2 for
   RDNA 4 support without editing any files. Added RDNA 4 / 12.0.0 to
   the HSA_OVERRIDE_GFX_VERSION comment table.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 Dockerfile              | 48 ++++++++++++++++++++++++++++++-
 docker-compose.rocm.yml | 64 +++++++++++++++++++++++++++++++++++++++++
 docker-compose.yml      |  4 +++
 3 files changed, 115 insertions(+), 1 deletion(-)
 create mode 100644 docker-compose.rocm.yml

diff --git a/Dockerfile b/Dockerfile
index 980584d1..c7bb79e3 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,8 +1,15 @@
 # ============================================================
-# Voicebox — Local TTS Server with Web UI (CPU)
+# Voicebox — Local TTS Server with Web UI
 # 3-stage build: Frontend → Python deps → Runtime
+#
+# Build variants:
+#   CPU (default):  docker compose up --build
+#   ROCm (AMD GPU): docker compose -f docker-compose.yml -f docker-compose.rocm.yml up --build
 # ============================================================
 
+# Top-level ARG so it is visible to all stages.
+ARG PYTORCH_VARIANT=cpu
+
 # === Stage 1: Build frontend ===
 FROM oven/bun:1 AS frontend
 
@@ -24,6 +31,9 @@ RUN cd web && bunx --bun vite build
 # === Stage 2: Build Python dependencies ===
 FROM python:3.11-slim AS backend-builder
 
+# Re-declare ARG inside the stage (Docker scoping requirement).
+ARG PYTORCH_VARIANT=cpu
+
 WORKDIR /build
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
@@ -34,6 +44,21 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 RUN pip install --no-cache-dir --upgrade pip
 
 COPY backend/requirements.txt .
+
+# ROCm version to pull PyTorch wheels for. Default is 6.3 (supports RDNA1/2/3).
+# Set ROCM_VERSION=7.2 for RDNA 4 (RX 9000 series) support.
+ARG ROCM_VERSION=6.3
+
+# When building the ROCm variant, install the ROCm-enabled PyTorch wheels
+# first so that the subsequent requirements.txt install sees them as already
+# satisfying the torch/torchaudio constraints and leaves them in place.
+# The CPU path skips this step and installs torch from PyPI as before.
+RUN if [ "$PYTORCH_VARIANT" = "rocm" ]; then \
+      pip install --no-cache-dir --prefix=/install \
+        torch torchaudio \
+        --index-url "https://download.pytorch.org/whl/rocm${ROCM_VERSION}"; \
+    fi
+
 RUN pip install --no-cache-dir --prefix=/install -r requirements.txt
 RUN pip install --no-cache-dir --prefix=/install --no-deps chatterbox-tts
 RUN pip install --no-cache-dir --prefix=/install --no-deps hume-tada
@@ -44,10 +69,31 @@ RUN pip install --no-cache-dir --prefix=/install \
 # === Stage 3: Runtime ===
 FROM python:3.11-slim
 
+# Re-declare ARG inside the stage (Docker scoping requirement).
+ARG PYTORCH_VARIANT=cpu
+
+# ROCm device access requires the container user to belong to the render
+# and video groups. GIDs are parameterised to match the host; Ubuntu 22.04+
+# defaults are used here. Override via env vars (docker-compose.rocm.yml
+# passes them through automatically):
+#   export RENDER_GID=$(getent group render | cut -d: -f3)
+#   export VIDEO_GID=$(getent group video  | cut -d: -f3)
+ARG RENDER_GID=992
+ARG VIDEO_GID=44
+RUN if [ "$PYTORCH_VARIANT" = "rocm" ]; then \
+      groupadd -f -g ${RENDER_GID} render && \
+      groupadd -f -g ${VIDEO_GID}  video; \
+    fi
+
 # Create non-root user for security
 RUN groupadd -r voicebox && \
     useradd -r -g voicebox -m -s /bin/bash voicebox
 
+# ROCm: add voicebox user to render+video so it can open /dev/kfd and /dev/dri.
+RUN if [ "$PYTORCH_VARIANT" = "rocm" ]; then \
+      usermod -aG render,video voicebox; \
+    fi
+
 WORKDIR /app
 
 # Install only runtime system dependencies
diff --git a/docker-compose.rocm.yml b/docker-compose.rocm.yml
new file mode 100644
index 00000000..957c1d31
--- /dev/null
+++ b/docker-compose.rocm.yml
@@ -0,0 +1,64 @@
+# ROCm (AMD GPU) overlay for Voicebox
+#
+# Usage:
+#   docker compose -f docker-compose.yml -f docker-compose.rocm.yml up --build
+#
+# Prerequisites on the host:
+#   1. ROCm drivers installed — https://rocm.docs.amd.com/projects/install-on-linux
+#   2. Current user in the 'render' and 'video' groups:
+#        sudo usermod -aG render,video $USER   (then log out/in)
+#
+# If the render/video GIDs on your host differ from the Ubuntu 22.04 defaults
+# (render=992, video=44), export them before running compose so the build arg
+# and group_add values both stay in sync automatically:
+#   export RENDER_GID=$(getent group render | cut -d: -f3)
+#   export VIDEO_GID=$(getent group video  | cut -d: -f3)
+#   docker compose -f docker-compose.yml -f docker-compose.rocm.yml up --build
+#
+# ROCm version (ROCM_VERSION):
+#   Default is 6.3 (supports RDNA1/2/3). For RDNA 4 (RX 9000 series) use 7.2:
+#   export ROCM_VERSION=7.2
+#   docker compose -f docker-compose.yml -f docker-compose.rocm.yml up --build
+
+services:
+  voicebox:
+    build:
+      context: .
+      args:
+        PYTORCH_VARIANT: rocm
+        # These build args read from env vars so a single export covers both the
+        # Dockerfile group creation and the runtime group_add below.
+        RENDER_GID: ${RENDER_GID:-992}
+        VIDEO_GID: ${VIDEO_GID:-44}
+        ROCM_VERSION: ${ROCM_VERSION:-6.3}
+
+    # Pass the AMD GPU device nodes into the container.
+    # /dev/kfd  — ROCm compute interface (required for GPU inference)
+    # /dev/dri  — DRM render nodes (required for display/memory access)
+    devices:
+      - /dev/kfd
+      - /dev/dri
+
+    # Grant access to the render and video groups so the non-root user
+    # inside the container can open the GPU device nodes.
+    # These reference the same env vars used in build.args above so a
+    # single export keeps Dockerfile groups and runtime group_add in sync.
+    group_add:
+      - "${RENDER_GID:-992}"   # render
+      - "${VIDEO_GID:-44}"     # video
+
+    environment:
+      # HSA_OVERRIDE_GFX_VERSION forces the ROCm runtime to treat the GPU as a
+      # specific GFX version when auto-detection fails or the GPU is newer than
+      # the ROCm release. app.py sets 10.3.0 (RDNA2) by default; override here
+      # for your GPU family:
+      #   RDNA4 / RX 9000 series:              12.0.0  (requires ROCM_VERSION=7.2)
+      #   RDNA3 / RX 7000 series / Strix Halo: 11.0.0
+      #   RDNA2 / RX 6000 series:              10.3.0
+      #   RDNA1 / RX 5000 series:              10.1.0
+      #   Vega / GCN5:                          9.0.0
+      - HSA_OVERRIDE_GFX_VERSION=11.0.0
+
+      # Tune the ROCm memory allocator to reduce fragmentation during
+      # multi-engine inference (TTS + STT + LLM running concurrently).
+      - PYTORCH_HIP_ALLOC_CONF=garbage_collection_threshold:0.8,max_split_size_mb:512
diff --git a/docker-compose.yml b/docker-compose.yml
index 0e49d526..2c85fb61 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -1,3 +1,7 @@
+# Voicebox — CPU build (default)
+# For AMD ROCm GPU acceleration use the overlay:
+#   docker compose -f docker-compose.yml -f docker-compose.rocm.yml up --build
+
 services:
   voicebox:
     build: .