環境変数CUDA_VISIBLE_DEVICESでハマったのでメモ
最近はすっかりLLMにハマってるわけですが、ハマり過ぎたあまり一台のPCに複数のGPUを搭載(といっても、RTX A4000(VRAM 16GB)とQuadro M6000(VRAM 24GB)という妙な組み合わせですが)してたりするのですが、どうにもLLMの出力速度が遅いなーなどと思っていました。
感覚的には主にM6000の方が使われてる感じで、こいつはVRAMこそ多いもののRTX A4000より圧倒的に遅い(数世代前)ので、メインで使ってほしくない。というわけで、CUDA_VISIBLE_DEVICES
という環境変数をいじってやればいいってところまではたどりついたのですが、CUDA_VISIBLE_DEVICES=2
とかしてやってcuda:1
みたいに指定すると妙なエラーが出る。色々調べてわかったのは、
export CUDA_VISIBLE_DEVICES=0,1 python llm.py # この中でcuda:1なりcuda:0なりを指定する
のようにすればいいらしいということ。この環境変数はデバイスの数ではなくて、GPUのリストを指定するものだったわけですね。盲点でした。