From d4280e1378ccf6a013fb882a6178850828404ec9 Mon Sep 17 00:00:00 2001 From: kingbri Date: Thu, 28 Mar 2024 11:02:56 -0400 Subject: [PATCH] Dependencies: Add pytorch-triton-rocm Required for AMD installs. Signed-off-by: kingbri --- pyproject.toml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 00cda21..0ea59bb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,6 +78,10 @@ cu118 = [ "flash_attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.5.2/flash_attn-2.5.2+cu118torch2.2cxx11abiFALSE-cp310-cp310-linux_x86_64.whl ; platform_system == 'Linux' and platform_machine == 'x86_64' and python_version == '3.10'", ] amd = [ + # Torch triton for ROCm + "pytorch_triton_rocm @ https://download.pytorch.org/whl/pytorch_triton_rocm-2.2.0-cp311-cp311-linux_x86_64.whl ; python_version == '3.11'", + "pytorch_triton_rocm @ https://download.pytorch.org/whl/pytorch_triton_rocm-2.2.0-cp310-cp310-linux_x86_64.whl ; python_version == '3.10'", + # Torch "torch @ https://download.pytorch.org/whl/rocm5.6/torch-2.2.1%2Brocm5.6-cp311-cp311-linux_x86_64.whl ; python_version == '3.11'", "torch @ https://download.pytorch.org/whl/rocm5.6/torch-2.2.1%2Brocm5.6-cp310-cp310-linux_x86_64.whl ; python_version == '3.10'",