本篇博文我们来详细讲解一下 CUDA_VISIBLE_DEVICES 这个环境变量，以及如何在 Ollama 中使用它来控制使用哪块 GPU。

1. CUDA_VISIBLE_DEVICES 是什么？

CUDA_VISIBLE_DEVICES 是一个环境变量，用于控制 CUDA 应用程序（比如 Ollama、PyTorch、TensorFlow 等）可以看到和使用哪些 GPU 设备。

  • 作用：它像一个“过滤器”，在程序启动前设定好，程序就会只认为这些指定的 GPU 是存在的。
  • 好处：
  • 资源分配：在多用户或多任务服务器上，可以为不同任务分配不同的 GPU，避免争抢。
  • 性能隔离：确保计算密集型任务不会相互干扰。
  • 指定计算设备：强制让程序在你想要的特定 GPU 上运行。

2. 如何查看可用的 GPU

在设置之前，你需要知道系统中有哪些 GPU以及它们的编号。

使用 nvidia-smi 命令：

nvidia-smi

输出结果类似如下：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  On   | 00000000:01:00.0  On |                  N/A |
|  0%   45C    P8    10W / 250W |    682MiB / 11264MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA GeForce ...  On   | 00000000:02:00.0 Off |                  N/A |
|  0%   35C    P8     9W / 250W |      2MiB / 11264MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

这里可以看到有两块 GPU：

  • GPU 0: 正在使用中（Memory-Usage 为 682MiB）。
  • GPU 1: 空闲状态。

3. 在 Ollama 中使用 CUDA_VISIBLE_DEVICES

有几种方法可以设置这个环境变量。

方法一：在启动 Ollama 服务前设置（推荐）

这是最常用和有效的方法。你需要在启动 ollama serve 命令的同一个终端会话中设置环境变量。

Linux / macOS (Bash/Zsh)

# 只使用 GPU 1
CUDA_VISIBLE_DEVICES=1 ollama serve

# 使用 GPU 0 和 GPU 1
CUDA_VISIBLE_DEVICES=0,1 ollama serve

# 不使用任何 GPU (强制使用CPU)
CUDA_VISIBLE_DEVICES= ollama serve
# 或者
CUDA_VISIBLE_DEVICES=-1 ollama serve

Windows (Command Prompt)

# 只使用 GPU 1
set CUDA_VISIBLE_DEVICES=1 && ollama serve

# 使用 GPU 0 和 GPU 1
set CUDA_VISIBLE_DEVICES=0,1 && ollama serve

Windows (PowerShell)

# 只使用 GPU 1
$env:CUDA_VISIBLE_DEVICES=1 ; ollama serve

# 使用 GPU 0 和 GPU 1
$env:CUDA_VISIBLE_DEVICES="0,1" ; ollama serve
方法二：导出环境变量（针对当前终端会话）

你也可以先导出环境变量，然后再运行 Ollama。这样，在这个终端里运行的所有后续命令都会继承这个设置。

Linux / macOS

# 导出，指定使用 GPU 1
export CUDA_VISIBLE_DEVICES=1

# 然后运行 Ollama
ollama serve

# 之后在这个终端里运行的任何 CUDA 程序都只会看到 GPU 1

Windows (Command Prompt)

set CUDA_VISIBLE_DEVICES=1
ollama serve

Windows (PowerShell)

$env:CUDA_VISIBLE_DEVICES=1
ollama serve
方法三：在创建系统服务时设置（适用于 Linux 系统服务）

如果你是通过 systemd 服务来运行 Ollama，可以修改其服务文件。

  1. 编辑 Ollama 服务文件：
sudo systemctl edit ollama
  1. 在打开的编辑器中，添加以下内容：
[Service]
Environment="CUDA_VISIBLE_DEVICES=1"
  1. 保存并退出，然后重新加载配置并重启服务：
sudo systemctl daemon-reload
sudo systemctl restart ollama

4. 验证是否生效

设置好后，如何确认 Ollama 真的在使用你指定的 GPU 呢？

  1. 首先，运行一个模型：在另一个终端，使用 ollama run 命令拉取或运行一个模型。
ollama run llama2
  1. 再次运行 nvidia-smi：观察 GPU 的内存使用情况和计算负载。
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    1   N/A  N/A      12345    C+G   ...ollama.exe                 5340MiB   |
+-----------------------------------------------------------------------------+

如果你指定了 CUDA_VISIBLE_DEVICES=1，你应该会看到有一个 ollama 进程正在 GPU 1 上运行并占用显存。

总结与注意事项

  • 编号规则：CUDA_VISIBLE_DEVICES 使用的编号是 nvidia-smi 列出的原始编号。
  • 重新映射：在程序内部，CUDA_VISIBLE_DEVICES=1 会使 GPU 1 被当作“GPU 0”来识别。程序只知道它能看到的那部分 GPU。
  • Ollama 多 GPU：当指定多个 GPU（如 0,1）时，Ollama 会尝试利用所有可见的 GPU 来分摊模型计算，这通常能提高推理速度，尤其是在运行大型模型时。
  • 常见问题：如果设置后 Ollama 仍然在使用错误的 GPU 或者报错，请检查：
  • 确保 Ollama 服务是在设置环境变量 之后 启动的。如果 Ollama 已经作为后台服务在运行，你需要先停止它，再按照上述方法设置环境变量并重新启动。
  • 检查 GPU 编号是否正确。
  • 确保你的 CUDA 驱动和 Ollama 版本正常支持 GPU。