折腾 Zhēteng

GPU order helper script for Ollama

Created: 5 min read

使用 ollama 这种基于 llama.cpp 的推理框架，多 gpu 同时使用的时候，如果是在消费级主板上面，那么这些显卡的调用顺序对于性能是会有影响的。

llm 里面靠前那些层使用频率会更高（好像是这么回事，或者说它们的 KV cache 被使用的频率会更高）。所以频繁使用的这些靠前的十几层最好一定要放在速度最快的 GPU 上面。尤其是能够使用 pcie x16 通道的插槽。

但是有时候，如果全都插满，会发现 cuda 里面 device 的顺序实际上并不理想，这时候就需要我们调整一下顺序。让最慢的卡跑在最后面。如果模型小，可能根本就用不上，如果模型大，就算用上，也最小化了影响。

当然了，像我这种连 pcie x1 的插槽都不放过，而非要用个 riser card 插着显卡用的情况也确实比较极端，看不得手上有闲置的硬件，然后来老是守着一堆过时的垃圾舍不得扔。。。。

gpu_info.sh:

#!/bin/bash
 
# ==============================================================================
# GPU Inspector & Sorter for Ollama/AI Rigs
# 功能：显示显卡 UUID、型号、PCI 插槽及实际带宽，并生成按带宽排序的 CUDA 环境变量
# ==============================================================================
 
# 1. 依赖检查
if ! command -v nvidia-smi &> /dev/null; then
    echo "Error: nvidia-smi not found."
    exit 1
fi
 
if ! command -v lspci &> /dev/null; then
    echo "Error: lspci not found. Please install pciutils (sudo apt install pciutils)."
    exit 1
fi
 
# 创建临时文件用于排序
TMP_FILE=$(mktemp)
 
echo "================================================================================================================================================"
printf "%-5s | %-14s | %-10s | %-30s | %-22s | %-38s\n" "CUDA" "PCI Bus ID" "Link(BW)" "Vendor" "Model" "GPU UUID"
echo "------------------------------------------------------------------------------------------------------------------------------------------------"
 
# 2. 遍历显卡信息
# 使用 nvidia-smi 获取基础信息，然后用 lspci 补充带宽和厂商信息
nvidia-smi --query-gpu=index,pci.bus_id,name,uuid --format=csv,noheader | while IFS=, read -r idx pci_id name uuid; do
    
    # 清理字符串空白
    idx=$(echo "$idx" | xargs)
    pci_id=$(echo "$pci_id" | xargs)
    name=$(echo "$name" | xargs)
    uuid=$(echo "$uuid" | xargs)
 
    # A. 获取厂商信息 (Subsystem)
    brand_info=$(lspci -v -s "$pci_id" 2>/dev/null | grep "Subsystem" | cut -d: -f2- | xargs)
    [ -z "$brand_info" ] && brand_info="Unknown"
 
    # B. 获取 PCIe 链路宽度 (需要 sudo 权限才能读取准确的 LnkSta)
    # 提取类似 "Width x16" 中的数字
    width_raw=$(sudo lspci -vv -s "$pci_id" 2>/dev/null | grep "LnkSta:" | grep -o "Width x[0-9]*" | head -n 1 | cut -d'x' -f2)
    
    if [ -z "$width_raw" ]; then 
        width_raw="0"      # 无法读取时视为最慢
        width_disp="?"
    else
        width_disp="x${width_raw}"
    fi
 
    # C. 打印可视化表格
    printf "%-5s | %-14s | %-10s | %-30s | %-22s | %-38s\n" "$idx" "$pci_id" "$width_disp" "${brand_info:0:30}" "$name" "$uuid"
 
    # D. 写入临时文件: [带宽数值] [UUID]
    echo "${width_raw} ${uuid}" >> "$TMP_FILE"
done
 
echo "================================================================================================================================================"
echo ""
echo ">>> Optimized Environment Variable (Sorted by Bandwidth: High -> Low) <<<"
echo "Paste the following line into your .bashrc or docker-compose environment:"
echo "-----------------------------------------------------------------------"
 
# 3. 排序逻辑
# sort -k1,1nr : 按第一列(带宽) 数值(n) 逆序(r) 排序
# awk '{print $2}' : 只取 UUID 列
# paste -sd "," : 用逗号合并为一行
SORTED_UUIDS=$(sort -k1,1nr "$TMP_FILE" | awk '{print $2}' | paste -sd ",")
 
echo "export CUDA_VISIBLE_DEVICES=${SORTED_UUIDS}"
echo "-----------------------------------------------------------------------"
 
# 清理
rm "$TMP_FILE"

最后产出的 string 可以直接给 docker compose 用，总之就是 env var 了，如果是 systemd 管理的 ollama, 自然就要添加到 systemd config 里面了

Example:

services:
  ollama:
    image: ollama/ollama:latest
    # ... 其他配置 (ports, volumes 等) ...
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    environment:
      # 强制指定 CUDA 顺序：x16 -> x4 -> x4 -> x1 (最慢)
      - CUDA_VISIBLE_DEVICES=GPU-xxxxx,GPU-yyyyy,GPU-zzzz,GPU-wwwww

Graph View