在 Linux 使用 AMD ROCm 運行 Whisper

2024-01-09

為了無痛使用 wayland,最近向公司申請了一部 AMD 顯示卡的電腦,然後便安裝 openai/whisper 測試一下 AI 性能,雖然官方文件說得好像只支援 cuda 一樣,但稍為爬下文就不難發現 AMD 也能使用 ROCm 實現 GPU 運算,表現也不錯,就是安裝過程有點 tricky

安裝 ROCm 驅動

以下是我在 Arch Linux 的安裝過程,其他熱門的 Linux distro 應該不難找到相關的驅動

先安裝 ROCm SDK

> yay -S rocm-hip-sdk rocm-opencl-sdk

這邊要注意一下版本,在之後安裝 pytorch 的時候最好找到對應的版本

> yay -Qi rocm-hip-sdk

Name            : rocm-hip-sdk
Version         : 5.7.1-2
...

把自己加進 videorender group

> sudo usermod -a -G video,render auphone

重登或重啟電腦後檢查一下是否在組別內

> groups

vboxusers video render wheel

然後用 clinfo 檢查一下是否能讀取顯示卡,我的 RX 7700 XT 顯示為 gfx1101,這個號碼之後會用到

> clinfo --list

Platform #0: AMD Accelerated Parallel Processing
 `-- Device #0: gfx1101

安裝 pytorch

先建立一個 venv 並啟動它

> python -m venv ~/whisper
> source ~/whisper/bin/activate

然後安裝 pytorch,因為我的 ROCm 版本是 5.7.1,所以這裡也安裝 5.7,其他版本的指令都可以在官網找到

install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm5.7

檢查 torch 是否能讀取 GPU

> python
>>> import torch
>>> torch.cuda.is_available()
True

如果顯示為 False 可以嘗試一下安裝其他 torch 版本

安裝 whisper dependencies

> git clone https://github.com/openai/whisper.git
> cd whisper
> pip install setuptools-rust
> pip install -r requirements.txt

安裝 whisper

> pip install -U openai-whisper

檢查一下 --device 是否顯示 (default: cuda)

> whisper --help | grep device 

usage: whisper [-h] [--model MODEL] [--model_dir MODEL_DIR] [--device DEVICE]
  --device DEVICE       device to use for PyTorch inference (default: cuda)

最後使用自己顯示卡的 GFX 版本執行 whisper

> HSA_OVERRIDE_GFX_VERSION=1101  whisper test.mp4 --model tiny