Skip to content

Update PyTorch installation for ROCm in Dockerfile #1019

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
boessu wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Update PyTorch installation for ROCm in Dockerfile #1019

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 32 additions & 2 deletions Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,11 +52,12 @@ COPY backend/requirements.txt .
ARG ROCM_VERSION=6.3

# For ROCm, make the PyTorch ROCm index primary so every install below resolves
# torch to ROCm wheels instead of the default CUDA build.
# torch to ROCm wheels instead of the default CUDA build. Fix it to torch 2.7.1 as
# it is actually the only working version for rocm.
RUN if [ "$PYTORCH_VARIANT" = "rocm" ]; then \
pip install --no-cache-dir --prefix=/install \
--index-url "https://download.pytorch.org/whl/rocm${ROCM_VERSION}" \
torch torchaudio && \
torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 && \
Comment on lines 57 to +60

@coderabbitai coderabbitai Bot Aug 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/usr/bin/env bash
set -euo pipefail

for item in \
  'torch:torch-2.7.1' \
  'torchvision:torchvision-0.22.1' \
  'torchaudio:torchaudio-2.7.1'
do
  package="${item%%:*}"
  wheel="${item#*:}"

  if curl -fsSL "https://download.pytorch.org/whl/rocm7.2/${package}/" | grep -Fq "$wheel"; then
    echo "Found ${wheel} on rocm7.2"
  else
    echo "Missing ${wheel} on rocm7.2" >&2
    exit 1
  fi
done

Repository: jamiepine/voicebox

Length of output: 187

🏁 Script executed:

#!/usr/bin/env bash
set -euo pipefail

printf 'Dockerfile relevant lines:\n'
sed -n '45,105p' Dockerfile | cat -n

printf '\nROCm index package listings:\n'
for idx in rocm6.3 rocm7.2; do
  printf '\nINDEX=%s\n' "$idx"
  for pkg in torch torchvision torchaudio; do
    idx_base="https://download.pytorch.org/whl/${idx}"
    html="$(curl -fsSL "${idx_base}/${pkg}/" || true)"
    if [ -z "${html}" ]; then
      printf 'index unavailable for %s\n' "$pkg"
      continue
    fi
    printf '%s matches count: ' "$pkg"
    printf '%s\n' "${html}" | grep -oE "<a href=\"${pkg}(-[0-9][^\"]*)?\.whl" | wc -l || true
    printf 'sample refs:\n'
    printf '%s\n' "${html}" | grep -oE "href=\"${pkg}[^\"\"]*.whl" | sed 's/^/  /' | head -30
  done
done

Repository: jamiepine/voicebox

Length of output: 3718

Keep the ROCm index and package tuple compatible.

ROCM_VERSION=7.2 is documented and used here, but the rocm7.2 index does not provide torch-2.7.1; the first ROCm install then fails before the image is built. Add the missing PyTorch packages to rocm7.2 on the PyTorch wheel index, or reject unsupported ROCM_VERSION values.

Also applies to: 50-53

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@Dockerfile` around lines 57 - 60, Update the ROCm installation branch around
PYTORCH_VARIANT and ROCM_VERSION so the selected PyTorch package tuple is
available from the configured ROCm wheel index, including the documented
ROCM_VERSION=7.2 case. Add the missing compatible packages to the rocm7.2 index
or validate and reject unsupported ROCM_VERSION values before pip install;
preserve the existing installation behavior for supported variants.

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

Expand All @@ -66,6 +67,35 @@ RUN pip install --no-cache-dir --prefix=/install --no-deps hume-tada
RUN pip install --no-cache-dir --prefix=/install \
git+https://github.com/QwenLM/Qwen3-TTS.git

# Qwen3-TTS's setup.py pulls in plain PyPI "triton" (CUDA-oriented) as an
# unconstrained dependency, alongside its nvidia-cuda-*/cuda-bindings/
# cuda-toolkit sub-deps. This shadows the correct pytorch-triton-rocm
# that shipped with our pinned ROCm torch, and breaks at runtime with
# "libcudart.so.13: cannot open shared object file" since no real NVIDIA
# driver is present. Strip the stray CUDA packages on ROCm builds.
RUN if [ "$PYTORCH_VARIANT" = "rocm" ]; then \
SITE=/install/lib/python3.11/site-packages && \
rm -rf "$SITE"/triton "$SITE"/triton-*.dist-info \
"$SITE"/nvidia "$SITE"/nvidia_*.dist-info \
"$SITE"/cuda_bindings "$SITE"/cuda_bindings-*.dist-info \
"$SITE"/cuda_pathfinder "$SITE"/cuda_pathfinder-*.dist-info \
"$SITE"/cuda_toolkit "$SITE"/cuda_toolkit-*.dist-info \
"$SITE"/torch "$SITE"/torch-*.dist-info \
"$SITE"/torchaudio "$SITE"/torchaudio-*.dist-info \
"$SITE"/torchvision "$SITE"/torchvision-*.dist-info \
"$SITE"/functorch \
"$SITE"/*.dist-info/../nvidia* 2>/dev/null; \
Comment on lines +78 to +87

@coderabbitai coderabbitai Bot Aug 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/usr/bin/env bash
set -euo pipefail

index='https://download.pytorch.org/whl/rocm6.3/pytorch-triton-rocm/'
html="$(curl -fsSL "$index")"

if ! grep -Eq 'pytorch_triton_rocm-3\.3\.1.*cp311.*\.whl' <<<"$html"; then
  echo 'Expected cp311 pytorch-triton-rocm 3.3.1 wheel was not found.' >&2
  exit 1
fi

echo 'Verify that the final Dockerfile install does not use --no-deps after deleting triton.'

Repository: jamiepine/voicebox

Length of output: 243

🏁 Script executed:

#!/usr/bin/env bash
set -euo pipefail

echo 'Files named Dockerfile:'
git ls-files | grep -E '(^|/)Dockerfile$' || true

echo
echo 'Relevant Dockerfile excerpts:'
if [ -f Dockerfile ]; then
  sed -n '1,130p' Dockerfile | cat -n
fi

echo
echo 'Search for rocm/pytorch/triton deps:'
rg -n "rocm|pytorch|triton|PIP_EXTRA_INDEX_URL|--no-deps|cuda" Dockerfile .github README.md pyproject.toml requirements*.txt setup.cfg 2>/dev/null || true

Repository: jamiepine/voicebox

Length of output: 11897

🏁 Script executed:

#!/usr/bin/env bash
set -euo pipefail

rocm="6.3"
base="https://download.pytorch.org/whl/rocm${rocm}"

echo "== PyPI triton candidates =="
python3 - <<'PY'
import urllib.request, re
url='https://pypi.org/simple/triton/'
html=urllib.request.urlopen(url, timeout=20).read().decode()
for x in sorted(set(re.findall(r'href="[^"]*.whl[^"]*"', html))):
    name=''.join(x.split('"')[1:-1].split('/')[-1].split('.'))
    print(name)
PY | grep -E 'triton-(2|3|4)\.[0-9]+\.[0-9]+-cp311-cp311-linux_x86_64.whl' || true

echo
echo "== ROCm Torch package candidates =="
python3 - <<'PY'
import urllib.request, re
url='https://download.pytorch.org/whl/rocm6.3/triton/'
try:
    html=urllib.request.urlopen(url, timeout=20).read().decode()
    for x in sorted(set(re.findall(r'href="[^"]*\.whl[^"]*"', html))):
        print(''.join(x.split('"')[1:-1].split('/')[-1]))
except Exception as e:
    print('ERROR torch-triton-index:', repr(e))
PY | grep -E 'pytorch_triton_rocm-3\.[23][0-9]*\.[0-9].*cp311.*\.whl' || true

echo
echo "== pypa packaging metadata for torch packages =="
python3 - <<'PY'
from urllib.request import urlopen
from email.parser import Parser
import zipfile, tempfile, os

packages = [
    'torch',
    'pytorch-triton-rocm',
    # direct simple index links are enough if wheel metadata exists via package API
]
pkg = 'pytorch-triton-rocm'
json_url = f'https://pypi.org/pypi/{pkg}/json'
html = urlopen('https://pypi.org/simple/pytest/triton/', timeout=20).read().decode()
text=json = None
try:
    import json
    text = urlopen(json_url, timeout=20).read().decode()
    j = json.loads(text)
    print('version:', j['info']['version'])
    print('requires_dist:')
    for req in j['info'].get('requires_dist') or []:
        if 'triton' in req.lower() or 'torch' in req.lower():
            print(' -', req)
except Exception as e:
    print('json error:', repr(e))
PY

Repository: jamiepine/voicebox

Length of output: 344

🏁 Script executed:

#!/usr/bin/env bash
set -euo pipefail

rocm="6.3"

echo "== triton simple index =="
python3 - <<'PY'
import urllib.request, re
url='https://pypi.org/simple/triton/'
html=urllib.request.urlopen(url, timeout=20).read().decode()
for x in sorted(set(re.findall(r'href="[^"]*\.whl[^"]*"', html))):
    print(x.split('/')[-1])
PY
echo

echo "== pytorch-triton-rocm simple index cp311 =="
python3 - <<'PY'
import urllib.request, re
url='https://download.pytorch.org/whl/rocm6.3/pytorch-triton-rocm/'
html=urllib.request.urlopen(url, timeout=20).read().decode()
for x in sorted(set(re.findall(r'href="[^"]*\.whl[^"]*"', html))):
    w=x.split('/')[-1]
    if 'cp311' in w:
        print(w)
PY
echo

echo "== pytorch-triton-rocm metadata =="
python3 - <<'PY'
import urllib.request, json
url='https://pypi.org/pypi/pytorch-triton-rocm/json'
data=json.load(urllib.request.urlopen(url, timeout=20))
print('version:', data['info']['version'])
for req in data['info'].get('requires_dist') or []:
    print(' -', req)
PY

Repository: jamiepine/voicebox

Length of output: 2468

🌐 Web query:

pypi simple triton packages triton 3.4.0 pytorch triton rocm

💡 Result:

Triton 3.4.0 is an official release that was published on July 30, 2025 [1][2]. It is available on PyPI [3][2] and supported within the PyTorch ecosystem, particularly for ROCm environments [4][5][6]. For ROCm environments, the package used is typically pytorch-triton-rocm [6], though wheels labeled specifically with triton and a +rocm suffix are also distributed for specific ROCm releases [5][7]. Key details regarding Triton 3.4.0 and ROCm: 1. Availability: The standard triton 3.4.0 package can be found on PyPI [3][2]. 2. ROCm Integration: AMD provides specific instructions for installing pytorch_triton_rocm or triton (with ROCm-specific suffixes) alongside PyTorch for ROCm [4][5][7]. These wheels are hosted on AMD's repositories (e.g., repo.radeon.com) and are often required to ensure compatibility with specific ROCm stack versions [4][7]. 3. PyTorch Foundation: Wheels for pytorch-triton-rocm are also available via the PyTorch download index [6], though AMD documentation recommends prioritizing wheels from their official repository (repo.radeon.com) for tested compatibility with their ROCm release versions [7]. 4. Installation Note: When installing these packages, particularly in virtual environments, it may be necessary to use flags like --break-system-packages if installing globally or in specific managed environments [7]. Additionally, enabling certain experimental features may require setting environment variables like export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 [7]. You should check your specific ROCm version documentation (e.g., ROCm 6.4.4 or 7.1) to identify the exact wheel compatible with your system [4][5].

Citations:

🌐 Web query:

github jamiepine voicebox Dockerfile --no-deps torch reinstall

💡 Result:

The Dockerfile for the jamiepine/voicebox repository handles PyTorch installation and dependency management through a multi-stage build process to ensure compatibility and performance. When building the Docker image, the process includes specific logic to manage PyTorch variants (CPU or ROCm) [1]. For ROCm builds, the Dockerfile specifically targets the ROCm-optimized PyTorch wheel index to ensure the environment uses the correct hardware-accelerated drivers instead of default builds [1][2]. The use of --no-deps in the Dockerfile (and in the repository's broader setup scripts) is primarily employed when installing specific TTS engine packages—such as chatterbox-tts and hume-tada—to avoid dependency conflicts [1][3][4]. These packages often pin older or conflicting versions of core libraries like torch, numpy, or transformers that would otherwise break the project's established dependency stack [3][4]. In the context of the repository's backend build utilities (specifically backend/build_binary.py), --force-reinstall and --no-deps are used to programmatically swap or ensure the correct PyTorch runtime is installed depending on the detected hardware environment, such as when shifting between CPU, CUDA, or ROCm variants during the binary creation process [5]. In summary, you do not typically "reinstall" torch manually unless you are modifying the build process or troubleshooting a specific environment issue [6]. If you encounter dependency issues, the repository recommends following the documented setup procedures (using just setup or the provided Dockerfile instructions), which are designed to correctly resolve these complex dependency constraints [7][8][4].

Citations:

🌐 Web query:

pip install pytorch-triton-rocm requires torch version

💡 Result:

The package pytorch-triton-rocm is tightly coupled to specific PyTorch releases and ROCm versions. Rather than having a single, universal version requirement, compatibility is maintained by ensuring that the specific build of pytorch-triton-rocm matches the ROCm-enabled PyTorch wheel you are installing [1][2][3]. When using ROCm-enabled PyTorch, you must install the version of pytorch-triton-rocm (or triton) that is specifically built for that same PyTorch and ROCm release [4][3][5]. Key guidelines for successful installation: 1. Use Matching Wheels: AMD recommends installing the PyTorch, Torchvision, Torchaudio, and Triton wheels together as a set from the official AMD ROCm repositories (repo.radeon.com) [4][6]. These sets are pre-tested to ensure the triton version is compatible with the torch version [4][7]. 2. Avoid Version Mismatches: Installing a mismatched pytorch-triton-rocm version often leads to runtime errors, such as InductorErrors or library version mismatches (e.g., attempting to use a ROCm 7-based triton with a ROCm 6.4-based PyTorch) [1]. 3. Installation Method: For most users, the most reliable method is to use the specific install commands provided in the official ROCm documentation for your target ROCm and PyTorch version [4][2][3][6]. These commands typically bundle the correct versions to avoid dependency conflicts [8]. 4. Check Official Sources: Because these packages are frequently updated, always consult the current AMD ROCm installation documentation for the recommended pip install commands rather than relying on standard pip behavior, which may not resolve the correct ROCm-specific dependencies [4][5]. If you are using nightly or test builds, ensure your --index-url points to the specific ROCm-enabled repository corresponding to your target ROCm version (e.g., https://download.pytorch.org/whl/rocmX.Y/) to ensure pip pulls the correct dependent packages [8][9][10].

Citations:

Restore pytorch-triton-rocm after deleting triton.

Lines 78-79 remove the ROCm Triton module that was installed with PyTorch 2.7.1. Line 95 reinstalls torch/torchaudio/torchvision with --no-deps, so the matching ROCm Triton dependency is not restored. Remove --no-deps for this re-pin while PIP_EXTRA_INDEX_URL remains disabled, or explicitly reinstall the matching pytorch-triton-rocm wheel.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@Dockerfile` around lines 78 - 87, Update the PyTorch re-pin command in the
Dockerfile to restore the matching pytorch-triton-rocm dependency after the
earlier triton removal. Remove --no-deps from the torch/torchaudio/torchvision
reinstall while PIP_EXTRA_INDEX_URL remains disabled, or explicitly install the
matching pytorch-triton-rocm wheel.

true; \
fi

# Re-pin torch/torchaudio to the ROCm build. requirements.txt (unpinned
# torch/torchvision) can pull the default CUDA wheels from pypi.org via
# the extra-index-url fallback and silently clobber the ROCm install above.
RUN if [ "$PYTORCH_VARIANT" = "rocm" ]; then \
PIP_EXTRA_INDEX_URL= pip install --no-cache-dir --prefix=/install --force-reinstall --no-deps \
--index-url "https://download.pytorch.org/whl/rocm${ROCM_VERSION}" \
torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1; \
fi

# === Stage 3: Runtime ===
FROM python:3.11-slim
Expand Down