tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2025-12-06 04:12:13 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Export TRITON_HIP_LLD_PATH in speedrun.sh for AMD ROCm

Browse Source 
When running on AMD ROCm using `uv`-installed packages (`rocm-sdk-core`), the `ld.lld` linker is not in the default `/opt/rocm/llvm/bin/` location expected by `pytorch-triton-rocm`. This causes `InductorError` during `torch.compile`.

This change updates `speedrun.sh` to dynamically find the `ld.lld` binary within the active python environment's site-packages (`_rocm_sdk_core`) and export the `TRITON_HIP_LLD_PATH` environment variable, allowing Triton to locate the linker correctly.
This commit is contained in:
google-labs-jules[bot] 2025-11-23 08:19:07 +00:00
parent 1f9b734358
commit 68148b1bf3
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
speedrun.sh
View File

@ -49,6 +49,13 @@ if [ "$EXTRAS" == "amd" ]; then
# Uninstalling triton may have deleted the shared 'triton' directory, breaking pytorch-triton-rocm. # Uninstalling triton may have deleted the shared 'triton' directory, breaking pytorch-triton-rocm.
# Reinstall pytorch-triton-rocm to ensure it's intact. # 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 uv pip install --force-reinstall --index-url https://repo.amd.com/rocm/whl/gfx1151 pytorch-triton-rocm
# Find and export the path to ld.lld from rocm-sdk-core if available, as torch.compile/triton needs it
ROCM_LLD_PATH=$(python -c "import sysconfig; import os; p = f\"{sysconfig.get_paths()['purelib']}/_rocm_sdk_core/lib/llvm/bin/ld.lld\"; print(p) if os.path.exists(p) else print('')")
if [ -n "$ROCM_LLD_PATH" ]; then
export TRITON_HIP_LLD_PATH=$ROCM_LLD_PATH
echo "Exported TRITON_HIP_LLD_PATH=$TRITON_HIP_LLD_PATH"
fi
fi fi
# ----------------------------------------------------------------------------- # -----------------------------------------------------------------------------