容器直接访问AMD GPU

作为开源友好的GPU厂商(也是为了和NVIDIA竞争)，AMD ROCm的设计高度依赖Linux内核的原生机制(例如 amdgpu 驱动提供的设备文件)，甚至可以不依赖 AMD Container Toolkit 就内为容器内部提供GPU的直接访问:

  • 从Linux角度来看，GPU就是 /dev/kfd (计算接口) 和 /dev/dri/renderD128 (渲染接口)

  • Docker 只需要通过标准的 --device 参数将上述两个文件映射进容器，容器内的 ROCm 运行时就能通过这两个接口直接和内核驱动对话

准备工作

Ubuntu 24.04 安装 amdgpu 驱动
wget https://repo.radeon.com/amdgpu-install/6.4.3/ubuntu/noble/amdgpu-install_6.4.60403-1_all.deb
sudo apt install ./amdgpu-install_6.4.60403-1_all.deb
sudo apt update
sudo apt install "linux-headers-$(uname -r)" "linux-modules-extra-$(uname -r)"
sudo apt install amdgpu-dkms

  • 安装完成后检查AMD GPU对应的模块设备:

检查 amdgpu 内核驱动敌营设备
ls -l /dev/kfd /dev/dri/render*

这里显示输出:

看到4个dri设备是因为我的主机同时安装了2块AMD MI50和2块NVIDIA Tesla A2
crw-rw---- 1 root render 226, 128 Feb 17 10:24 /dev/dri/renderD128
crw-rw---- 1 root render 226, 129 Feb 17 10:24 /dev/dri/renderD129
crw-rw---- 1 root render 226, 130 Feb 17 10:24 /dev/dri/renderD130
crw-rw---- 1 root render 226, 131 Feb 17 10:24 /dev/dri/renderD131
crw-rw---- 1 root render 237,   0 Feb 17 10:24 /dev/kfd

备注

我的服务器安装了 2块 NVIDIA Tesla A2 GPU运算卡 和 2块 AMD Radeon Instinct MI50 ，所以在上面的设备检查会看到4个 renderD... 设备:

在Linux的Direct Rendering Manager(DRM)子系统中，所有具备渲染能力的GPU都会占用一个 renderD 编号。要区分设备厂商，可以通过 renderD*/device/vendor 来区分:

  • AMD的厂商ID是 0x1002

  • NVIDIA的厂商ID是 0x10de

通过vender ID来区分设备厂商
# grep -H "" /sys/class/drm/renderD*/device/vendor | grep "0x1002"
/sys/class/drm/renderD128/device/vendor:0x1002
/sys/class/drm/renderD129/device/vendor:0x1002

# grep -H "" /sys/class/drm/renderD*/device/vendor | grep "0x10de"
/sys/class/drm/renderD130/device/vendor:0x10de
/sys/class/drm/renderD131/device/vendor:0x10de

上述输出可以看到 128 和 129 是AMD的设备 MI50，而 130 和 131 是NVIDIA的设备 A2

手工检测的方法比较麻烦，所以安装 AMD Container Toolkit 和 NVIDIA Container Toolkit 就带来了极大的便利，可以直接对各自厂商的GPU设备进行操作而不用编写复杂的检测脚本

需要注意，设备的属主组是 render ，所以需要将使用者的uid加入到这个组:

将自己的用户帐号加入到 render 组
sudo usermod -aG video $USER
sudo usermod -aG render $USER
# 注意：重新登录才能生效，如果是ssh multiplexing需要kill掉本地ssh客户端进程重新登录

  • 安装docker:

在Ubuntu中安装发行版提供的docker
sudo apt install docker.io

为方便管理，将使用者(自己)的uid加入( 无需sudo运行docker ):

将当前用户添加到 docker 用户组
sudo usermod -aG docker $USER

运行

由于Docker容器和host OS共享内核，所以ROCm内核模式驱动( amdgrpu-dkms )在物理主机上安装完成后，只需要授权Docker容器访问主机AMD GPU就可以在容器中于行ROCm:

授权docker容器访问host主机的AMD GPU
docker run -d \
    --device /dev/kfd \
    --device /dev/dri/renderD128 \
    --device /dev/dri/renderD129 \
    --security-opt seccomp=unconfined \
    -v ollama-amd:/root/.ollama \
    -p 11435:11434 \
    --name ollama-amd \
    ollama/ollama:rocm

备注

现在从官方下载运行的Ollama镜像是精简没有包含 rocminfo 等维护工具

验证

rocminfo 能够提供HSA系统属性以及agents的信息，此外 amd-smi 则提供命令行接口来操纵和监控amdgpu内核:

  • docker inpsect ollama-amd 可以看到运行容器将host主机的GPU设备bind进入容器:

通过 inspect 指令检查容器bind的设备
...
            "Devices": [
                {
                    "PathOnHost": "/dev/kfd",
                    "PathInContainer": "/dev/kfd",
                    "CgroupPermissions": "rwm"
                },
                {
                    "PathOnHost": "/dev/dri/renderD128",
                    "PathInContainer": "/dev/dri/renderD128",
                    "CgroupPermissions": "rwm"
                },
                {
                    "PathOnHost": "/dev/dri/renderD129",
                    "PathInContainer": "/dev/dri/renderD129",
                    "CgroupPermissions": "rwm"
                }
            ],
...

  • 通过 docker exec -it ollama-amd /bin/bash 进入运行容器内部检查: 需要单独安装 rocminfo 工具

由于Ollama精简镜像不包含维护工具，所以需要单独安装 rocminfo
apt update
apt install rocminfo

安装了 rocminfo 之后，在容器中就可以观察驱动信息: 直接执行 rocminfo 可以看到如下输出，验证了容器内部已经绑定了host主机的GPU设备:

rocminfo 输出信息
=====================    
HSA System Attributes    
=====================    
Runtime Version:         1.1
System Timestamp Freq.:  1000.000000MHz
Sig. Max Wait Duration:  18446744073709551615 (0xFFFFFFFFFFFFFFFF) (timestamp count)
Machine Model:           LARGE                              
System Endianness:       LITTLE                             
Mwaitx:                  DISABLED
DMAbuf Support:          YES

==========               
HSA Agents               
==========
...
*******                  
Agent 2                  
*******                  
  Name:                    gfx906                             
  Uuid:                    GPU-6a92216172f0fcb5               
  Marketing Name:          AMD Radeon Graphics                
  Vendor Name:             AMD                                
  Feature:                 KERNEL_DISPATCH                    
  Profile:                 BASE_PROFILE                       
  Float Round Mode:        NEAR                               
  Max Queue Number:        128(0x80)                          
  Queue Min Size:          64(0x40)                           
  Queue Max Size:          131072(0x20000)                    
  Queue Type:              MULTI                              
  Node:                    1                                  
  Device Type:             GPU                                
  Cache Info:              
    L1:                      16(0x10) KB                        
    L2:                      8192(0x2000) KB                    
...
  Internal Node ID:        1                                  
  Compute Unit:            60                                 
  SIMDs per CU:            4                                  
  Shader Engines:          4                                  
  Shader Arrs. per Eng.:   1                                  
  WatchPts on Addr. Ranges:4                                  
...

运行模型

一切就绪，现在可以 在Docker中Ollama使用AMD GPU运行大模型

参考