bitsandbytes の NF4 量子化が、ROCm の PyTorch と一緒に使えるかを技術検証しました。PyTorch の ROCm 版の wheel は ROCm のライブラリを中に持っていて、システムに ROCm を入れなくても GPU で動くように見えますが、量子化に使う bitsandbytes は、システムに入れた ROCm を読みに行きます。どのライブラリがどこから ROCm を読むのかを確かめた結果をまとめます。

結論

  • torch 2.13.0+rocm7.2 は、HIP・HSA・rocBLAS・MIOpen・hipBLASLt を wheel の torch/lib から読みます(ldd で確かめました)。
  • bitsandbytes 0.50.1(PyPI の公式 wheel)は ROCm 用の .so を同梱していて、gfx1100 向けのフォークは要りません。ただし RUNPATH が /opt/rocm/lib なので、NF4 の経路にはシステムの ROCm が要ります。
  • ROCm を入れると、apt の依存で Ubuntu 版の ROCm 5.7 のライブラリも入ります。torch には効きません。

試した環境

  • GPU: AMD Radeon RX 7900 XT(gfx1100、VRAM 20 GiB)
  • OS: Ubuntu
  • ROCm 7.2.4(amdgpu-install --usecase=rocm --no-dkms で入れた)
  • ライブラリ: torch 2.13.0+rocm7.2、bitsandbytes 0.50.1、diffusers 0.40.0.dev0、transformers 5.15.0
  • NF4 を使ったモデル: Qwen-Image 系の text_encoder(拡散モデルの transformer は、bitsandbytes の int8 では絵にならず、NF4 では平坦な面に粒が出たので、量子化せず bf16 のまま)

ROCm はどこから読まれるか

次の図は、私たちの環境で確かめたライブラリごとの読み先です。

ライブラリごとの ROCm の読み先

torch は HIP や rocBLAS などを、wheel の中の torch/lib から読みます。システムに一緒に入った ROCm 5.7 のライブラリは、torch には効きません。

一方、bitsandbytes 0.50.1 の ROCm 用の .so は、RUNPATH が /opt/rocm/lib になっていて、システムの ROCm を読みます。torch だけを見て「システムの ROCm は要らない」と考えると、NF4 の経路が読む先がなくなります。torch の動作確認だけでは気づけないので、量子化を使うなら、ここを先に見ておくのが安全です。

読み先の設定(RUNPATH)は readelf で、実際にどこで見つかるかは ldd で確かめられます。

# torch: 読み先が wheel の torch/lib になっているか
ldd <venv>/lib/python3.12/site-packages/torch/lib/<ライブラリ>.so | grep -iE 'hip|hsa|rocblas|miopen|hipblaslt'

# bitsandbytes: ROCm 用の .so の読み先の設定と、実際に見つかるか
readelf -d <venv>/lib/python3.12/site-packages/bitsandbytes/<ROCm 用の .so> | grep -E 'RUNPATH|RPATH'
ldd <venv>/lib/python3.12/site-packages/bitsandbytes/<ROCm 用の .so> | grep -E '/opt/rocm|not found'
荷物を全部背負って進む旅人と、道ばたの小屋で必要な物を受け取ってから進む旅人の絵

bitsandbytes にフォークは要らなかった

ROCm で bitsandbytes を使うには、AMD の GPU 向けのフォークを入れる、という道もあります。私たちの環境では、PyPI の公式 wheel(0.50.1)が ROCm 用の .so を同梱していて、gfx1100 向けのフォークは要りませんでした。要るのは、上で見たシステムの ROCm だけです。

ROCm を入れるとき

ROCm 7.2.4 は、次のコマンドで入れました。

amdgpu-install --usecase=rocm --no-dkms

このとき apt の依存の連鎖で、Ubuntu 版の ROCm 5.7 のライブラリ(libamdhip64-5 など)も一緒に入ります。版の違う ROCm が並ぶので戸惑いますが、torch は wheel の中の物を読むので、5.7 は torch には効きません。

おわりに

ROCm の PyTorch がシステムの ROCm 無しで動いていても、bitsandbytes を使うならシステムの ROCm も入れておいてください。どこから読んでいるかは、readelf と ldd で確かめられます。

※ PyTorch と bitsandbytes の wheel の作りは執筆時点の版(torch 2.13.0+rocm7.2、bitsandbytes 0.50.1)で見た物で、いまは変わっている可能性があります。