Fix AMD Triton runtime error by reinstalling pytorch-triton-rocm

Uninstalling the conflicting `triton` package (upstream) on AMD systems often removes the `triton` directory shared with `pytorch-triton-rocm`, breaking the latter. This caused `ImportError: cannot import name 'Config' from 'triton'`.

This change adds a step to force reinstall `pytorch-triton-rocm` immediately after uninstalling `triton`, ensuring the correct package is present and intact for the runtime.
This commit is contained in:
google-labs-jules[bot] 2025-11-23 05:22:21 +00:00
parent d291a62ad8
commit b92647c580

View File

@ -46,6 +46,9 @@ source .venv/bin/activate
# if the NVIDIA version of triton (e.g. 3.4.0) is accidentally installed.
if [ "$EXTRAS" == "amd" ]; then
uv pip uninstall -q triton || true
# Uninstalling triton may have deleted the shared 'triton' directory, breaking pytorch-triton-rocm.
# Reinstall pytorch-triton-rocm to ensure it's intact.
uv pip install --force-reinstall --index-url https://repo.amd.com/rocm/whl/gfx1151 pytorch-triton-rocm
fi
# -----------------------------------------------------------------------------