fix: ROCm setup for Linux AMD GPUs (#817)

* Fix ROCm setup for Linux AMD GPUs

- Ensure Docker ROCm builds resolve PyTorch packages from the ROCm wheel index so later dependency installs do not replace them with CUDA wheels.
- Move ROCm device group handling to a runtime entrypoint that joins the groups owning /dev/kfd and /dev/dri, avoiding distro-specific render/video GID defaults.
- Leave HSA_OVERRIDE_GFX_VERSION unset by default in the ROCm compose overlay so newer RDNA GPUs can use native ROCm detection.
- Add Linux GPU detection to the Unix setup recipe so AMD systems install ROCm torch wheels and NVIDIA systems install CUDA wheels before backend dependencies.

* docs(changelog): add Linux ROCm setup entry

* fix(setup): pin ROCm torch wheels and prefer NVIDIA over amdgpu

- Install torch/torchaudio from the ROCm index only, before the pooled
  requirements install, so a plain PyPI (CUDA) wheel can't outrank +rocm
- Detect NVIDIA before AMD and gate ROCm on /dev/kfd, so hybrid
  AMD+NVIDIA hosts get CUDA instead of ROCm

Author: 1337hero

CHANGELOG.md

@@ -5,6 +5,17 @@
 
 # Changelog
 
+## [Unreleased]
+
+### Linux
+
+- **ROCm setup works on Linux AMD systems.** Docker ROCm builds now keep PyTorch
+  on the ROCm wheel index during dependency installation, so later installs do
+  not replace it with CUDA wheels. The ROCm compose overlay no longer assumes
+  Ubuntu render/video group IDs; the container joins the groups that own the GPU
+  device nodes at startup. Native Linux setup now picks ROCm wheels for AMD GPUs
+  and CUDA wheels for NVIDIA GPUs before installing backend dependencies.
+
 ## [0.5.0] - 2026-04-22
 
 **The Capture release.** Voicebox stops being just a voice-cloning studio and becomes a full AI voice studio. Hold a key anywhere on your machine, speak, release — the transcript lands in the focused text field. Flip the primitive around and any MCP-aware agent — Claude Code, Cursor, Spacebot — speaks back through an on-screen pill in one of your cloned voices. A local LLM sits between the two, so transcripts come out clean and voice profiles can carry a personality that reshapes what the agent says before it gets spoken.

Dockerfile

@@ -45,18 +45,16 @@ 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.
+# ROCm wheel index. Default 6.3 (RDNA1/2/3); set ROCM_VERSION=7.2 for RDNA4.
 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.
+# For ROCm, make the PyTorch ROCm index primary so every install below resolves
+# torch to ROCm wheels instead of the default CUDA build.
 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}"; \
+        --index-url "https://download.pytorch.org/whl/rocm${ROCM_VERSION}" \
+        torch torchaudio && \
+      printf '[global]\nindex-url = https://download.pytorch.org/whl/rocm%s\nextra-index-url = https://pypi.org/simple\n' "$ROCM_VERSION" > /etc/pip.conf; \
     fi
 
 RUN pip install --no-cache-dir --prefix=/install -r requirements.txt
@@ -69,37 +67,17 @@ 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
+# Create non-root user; the entrypoint joins GPU device groups at runtime.
 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
+# Install only runtime system dependencies (gosu drops root in the entrypoint)
 RUN apt-get update && apt-get install -y --no-install-recommends \
     ffmpeg \
     curl \
+    gosu \
     && rm -rf /var/lib/apt/lists/*
 
 # Copy installed Python packages from builder stage
@@ -115,15 +93,14 @@ COPY --from=frontend --chown=voicebox:voicebox /build/web/dist /app/frontend/
 RUN mkdir -p /app/data/generations /app/data/profiles /app/data/cache \
     && chown -R voicebox:voicebox /app/data
 
-# Switch to non-root user
-USER voicebox
-
 # Expose the API port
 EXPOSE 17493
 
 # Health check — auto-restart if the server hangs
 HEALTHCHECK --interval=30s --timeout=10s --retries=3 --start-period=60s \
     CMD curl -f http://localhost:17493/health || exit 1
 
-# Start the FastAPI server
+# Entrypoint joins GPU groups then drops to the voicebox user
+COPY --chmod=755 scripts/rocm-entrypoint.sh /usr/local/bin/entrypoint.sh
+ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]
 CMD ["uvicorn", "backend.main:app", "--host", "0.0.0.0", "--port", "17493"]

docker-compose.rocm.yml

@@ -1,64 +1,36 @@
+---
 # 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
+# Requires ROCm drivers on the host:
+# https://rocm.docs.amd.com/projects/install-on-linux
+# RDNA4 (RX 9000): export ROCM_VERSION=7.2 (default 6.3 covers RDNA1-3).
 
 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)
+      #   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
+      - HSA_OVERRIDE_GFX_VERSION=${HSA_OVERRIDE_GFX_VERSION:-}
 
-      # Tune the ROCm memory allocator to reduce fragmentation during
-      # multi-engine inference (TTS + STT + LLM running concurrently).
+      # Tune the ROCm memory allocator
       - PYTORCH_HIP_ALLOC_CONF=garbage_collection_threshold:0.8,max_split_size_mb:512

justfile

@@ -43,6 +43,26 @@ setup-python:
     fi
     echo "Installing Python dependencies..."
     {{ pip }} install --upgrade pip -q
+    if [ "$(uname)" = "Linux" ]; then
+        torch_index=""
+        if [ -e /proc/driver/nvidia/version ] || [ -d /sys/module/nvidia ]; then
+            echo "Detected NVIDIA GPU — installing CUDA PyTorch..."
+            torch_index="https://download.pytorch.org/whl/cu128"
+        elif [ -e /dev/kfd ]; then
+            if [ -n "${VOICEBOX_ROCM_VERSION:-}" ]; then
+                rocm_ver="$VOICEBOX_ROCM_VERSION"
+            elif lspci 2>/dev/null | grep -qi "Navi 4"; then
+                rocm_ver=7.2
+            else
+                rocm_ver=6.3
+            fi
+            echo "Detected AMD GPU — installing ROCm PyTorch (rocm${rocm_ver})..."
+            torch_index="https://download.pytorch.org/whl/rocm${rocm_ver}"
+        fi
+        if [ -n "$torch_index" ]; then
+            {{ pip }} install torch torchaudio --index-url "$torch_index"
+        fi
+    fi
     {{ pip }} install -r {{ backend_dir }}/requirements.txt
     # Chatterbox pins numpy<1.26 / torch==2.6 which break on Python 3.12+
     {{ pip }} install --no-deps chatterbox-tts

scripts/rocm-entrypoint.sh

@@ -0,0 +1,15 @@
+#!/bin/sh
+set -e
+# Join whatever groups own the mounted GPU nodes so /dev/kfd and /dev/dri work
+# on any host (no RENDER_GID/VIDEO_GID needed), then drop to the app user.
+for dev in /dev/kfd /dev/dri/render*; do
+    [ -e "$dev" ] || continue
+    gid=$(stat -c %g "$dev")
+    grp=$(getent group "$gid" | cut -d: -f1)
+    [ -n "$grp" ] || {
+        grp="gpu$gid"
+        groupadd -g "$gid" "$grp"
+    }
+    usermod -aG "$grp" voicebox
+done
+exec gosu voicebox "$@"