Update PyTorch installation for ROCm in Dockerfile - #1019
Conversation
Fix PyTorch version for ROCm compatibility and remove stray CUDA packages.
📝 WalkthroughWalkthroughThe Dockerfile now pins ROCm PyTorch packages and removes conflicting CUDA-oriented packages before force-reinstalling the pinned ROCm wheels. ChangesROCm build dependency installation
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with 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.
Inline comments:
In `@Dockerfile`:
- Around line 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.
- Around line 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.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
| 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 && \ |
There was a problem hiding this comment.
🎯 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
doneRepository: 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
doneRepository: 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.
| 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; \ |
There was a problem hiding this comment.
🎯 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 || trueRepository: 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))
PYRepository: 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)
PYRepository: 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:
- 1: https://github.com/triton-lang/triton/releases/tag/v3.4.0
- 2: https://pypi.org/project/triton/
- 3: https://pypi.org/project/triton/3.4.0/
- 4: https://rocm.docs.amd.com/projects/radeon-ryzen/en/docs-6.4.4/docs/install/installrad/native%5Flinux/install-pytorch.html
- 5: https://rocm.docs.amd.com/projects/radeon-ryzen/en/docs-7.1/docs/install/installrad/native%5Flinux/install-pytorch.html
- 6: https://download.pytorch.org/whl/pytorch-triton-rocm/
- 7: https://rocm.docs.amd.com/projects/radeon-ryzen/en/docs-7.1/docs/install/installryz/native%5Flinux/install-pytorch.html
🌐 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:
- 1: https://github.com/jamiepine/voicebox/blob/main/Dockerfile
- 2: https://deepwiki.com/jamiepine/voicebox/6.4-docker-deployment
- 3: https://github.com/jamiepine/voicebox/blob/HEAD/docs/content/docs/developer/tts-engines.mdx
- 4: https://docs.voicebox.sh/developer/setup
- 5: https://github.com/jamiepine/voicebox/blob/f2cf2a72/backend/build_binary.py
- 6: https://github.com/jamiepine/voicebox/blob/f2cf2a72/docs/content/docs/overview/troubleshooting.mdx
- 7: https://github.com/jamiepine/voicebox
- 8: https://github.com/jamiepine/voicebox/blob/main/docs/content/docs/overview/docker.mdx
🌐 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:
- 1: Torch nightly rocm6.4 pytorch-triton-rocm requirement mismatch pytorch/pytorch#167411
- 2: https://rocm.docs.amd.com/projects/radeon-ryzen/en/docs-6.4.2/docs/install/native%5Flinux/install-pytorch.html
- 3: https://rocm.docs.amd.com/projects/radeon-ryzen/en/docs-6.4.4/docs/install/installrad/native_linux/install-pytorch.html
- 4: https://rocm.docs.amd.com/projects/radeon-ryzen/en/latest/docs/install/installryz/native_linux/install-pytorch.html
- 5: https://rocm.docs.amd.com/projects/radeon-ryzen/en/docs-7.2/docs/install/installrad/native_linux/install-pytorch.html
- 6: https://rocm.docs.amd.com/projects/radeon-ryzen/en/latest/docs/install/installrad/native_linux/install-pytorch.html
- 7: https://rocm.docs.amd.com/projects/radeon-ryzen/en/docs-7.2.1/docs/install/installrad/native_linux/install-pytorch.html
- 8: torch 2.2.0 on ROCm 5.7 wheel failed to pip install with error "No matching distribution found for pytorch-triton-rocm==2.2.0" pytorch/pytorch#118233
- 9: https://github.com/pytorch/pytorch/blob/2fcfae72b4fc478dc7922b1efc637d9cb64dd42f/RELEASE.md
- 10: https://rocm.docs.amd.com/projects/radeon-ryzen/en/latest/docs/install/installrad/native_linux/install-triton.html
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.
1 participant
Fix ROCm/AMD GPU support in Docker build
Problem
The ROCm build path in the Dockerfile silently falls back to a CUDA-only PyTorch build even when
PYTORCH_VARIANT=rocmis set, causingtorch.cuda.is_available()to returnFalseon AMD GPUs and, once partially fixed,ImportError: libcudart.so.13: cannot open shared object filewhen loading models (e.g. Qwen3-TTS) that depend ontorchaudio's compiled extension.Root causes (two separate issues stacked on top of each other):
PyTorch's "wheel variant" system (introduced in 2.9) relies on a provider plugin to detect the local ROCm environment and pick the right backend variant. That detection can't work inside a
docker buildcontext (no ROCm runtime present at build time), so an unpinnedtorchinstall resolves to the default CUDA build even when pointed at the ROCm-specific package index (--index-url https://download.pytorch.org/whl/rocm6.3).git+https://github.com/QwenLM/Qwen3-TTS.gitis installed without--no-deps(unlikechatterbox-ttsandhume-tada, which correctly use--no-deps). This pulls in unconstrained, CUDA-oriented transitive dependencies (triton,nvidia-*,cuda-bindings,cuda-toolkit) and a second, unpinnedtorch/torchaudioinstall. Because these installs happen across separatepip install --prefix=/installinvocations,pip --force-reinstalldoes not reliably remove the stale.dist-infodirectories from the first install — both versions end up coexisting in the final image, and the wrong one gets loaded at runtime.Fix
torch,torchvision, andtorchaudioto explicit versions (2.7.1/0.22.1/2.7.1) in the initial ROCm install step, since this is the last PyTorch generation using classic (non-variant)+rocm6.3-tagged wheels — avoids the wheel-variant auto-detection problem entirely.PIP_EXTRA_INDEX_URL=(empty) to the final pinned reinstall step, since/etc/pip.conf(written by the initial ROCm install) setsextra-index-url = https://pypi.org/simple, which otherwise lets pip silently fall back to PyPI's default CUDA build even with an explicit--index-url.triton,nvidia-*,cuda_bindings,cuda_pathfinder,cuda_toolkit) and the unpinnedtorch/torchaudio/torchvision/functorchdirectories/dist-info from/install, so the subsequent pinned reinstall starts from a clean slate instead of relying onpip's reinstall detection across separate invocations.Verification
Before:
torch-2.13.0.dist-infoandtorch-2.7.1+rocm6.3.dist-info(and same for torchaudio) present simultaneously.After: only
torch-2.7.1+rocm6.3.dist-info,torchvision-0.22.1+rocm6.3.dist-info,torchaudio-2.7.1+rocm6.3.dist-info.Full container startup log now reports:
and
completes without the
libcudart.so.13error.Note for maintainers
scripts/package_rocm.pydeclarestorch_compat=">=2.9.0,<2.10.0"as the default compatibility range, which conflicts with the2.7.1pin needed here to avoid the wheel-variant detection issue. I haven't touched that file — flagging it separately since reconciling it (e.g. testing whether a 2.9.x pin can be made to work with an explicit provider override, or updating the declared range) seems like a separate discussion from this build fix.Summary by CodeRabbit