Ubuntu 下 AMD AI MAX 395+ 部署 Qwen 模型与 ROCm 加速实战

Ubuntu 环境下 AMD AI MAX 395+ 本地部署 Qwen 模型指南

本文基于 Ubuntu 22.04 系统，演示如何在 AMD AI MAX 395+ 硬件上配置 ROCm 驱动，并利用 Docker 和 vLLM 本地部署千问（Qwen）系列模型。内容涵盖环境搭建、镜像离线迁移及多场景服务启动。

一、ROCm 7.0 驱动安装

参考官方文档选择对应系统版本的 ROCm 进行安装。若需更换版本，可查阅 AMD GPU Install 仓库。

# 更新 apt 缓存并安装基础工具
sudo apt update && sudo apt install wget -y

# 可选：卸载旧驱动
sudo apt autoremove amdgpu-dkms
sudo rm /etc/apt/sources.list.d/amdgpu.list
sudo rm -rf /var/cache/apt/*
sudo apt clean all
sudo apt update

# 下载并安装 ROCm 7.0.3 版本
wget https://repo.radeon.com/amdgpu-install/7.0.3/ubuntu/jammy/amdgpu-install_7.0.3.70003-1_all.deb
sudo apt install ./amdgpu-install_7.0.3.70003-1_all.deb

# 安装依赖并配置用户组
sudo apt install python3-setuptools python3-wheel
sudo usermod -aG render,video $LOGNAME

# 安装核心驱动包
sudo apt install rocm
sudo apt update
sudo apt install "linux-headers-$(uname -r)" "linux-modules-extra-$(uname -r)"
sudo apt install amdgpu-dkms

# 确保当前用户加入 render 和 video 组
sudo usermod -aG render $USER
sudo usermod -aG video $USER

# 重启系统
reboot

# 验证 GPU 识别情况
rocminfo | grep gfx

二、Docker 环境准备（vLLM）

1. 安装并配置 Docker

# 更新软件源
sudo apt update -y

# 安装依赖包
sudo apt-get install apt-transport-https ca-certificates curl software-properties-common lrzsz -y

# 添加阿里云 Docker GPG 密钥
sudo curl -fsSL https://mirrors.aliyun.com/docker-ce/linux/ubuntu/gpg | sudo apt-key add -

# 添加阿里云 Docker 软件源
sudo add-apt-repository "deb [arch=amd64] https://mirrors.aliyun.com/docker-ce/linux/ubuntu $(lsb_release -cs) stable"

# 再次更新并安装 Docker CE
sudo apt update -y
sudo apt-get install docker-ce -y

# 验证版本
docker version

# 配置镜像加速器
sudo tee /etc/docker/daemon.json <<-'EOF'
{
  "registry-mirrors": [
    "https://docker.1panel.live",
    "https://hub.rat.dev"
  ]
}
EOF

# 重载配置并重启服务
sudo systemctl daemon-reload
sudo systemctl restart docker

2. 拉取 vLLM 镜像

针对无法联网的主机，建议先在联网机器上打包镜像文件至 U 盘，再导入目标环境。

2.1 镜像打包

# 在联网主机上拉取镜像（版本可根据需求调整）
docker pull rocm/vllm:rocm7.0.0_vllm_0.11.2

# 导出为 tar 文件
docker save -o vllm_rocm7.tar rocm/vllm:rocm7.0.0_vllm_0.11.2

# 将 vllm_rocm7.tar 拷贝至 U 盘（确保格式支持大文件，如 exFAT 或 NTFS）

2.2 导入镜像

# 找到 U 盘挂载路径，例如 /media/用户名/U 盘名称
# 将文件拷贝到本地目录
cp /media/用户名/U 盘名称/vllm_rocm7.tar ~/

# 导入镜像
docker load -i vllm_rocm7.tar

# 验证镜像存在
docker images

三、部署千问模型

1. Qwen3-32B 对话模型

1.1 下载模型

# 创建模型存储目录
mkdir -p /opt/qwen_models/Qwen-32B-AWQ
export MODEL_DIR=/opt/qwen_models/Qwen-32B-AWQ

# 安装 modelscope 依赖
pip3 install modelscope

# 执行下载脚本
python3 -c "
import os
from modelscope.hub.snapshot_download import snapshot_download
model_id = 'qwen/Qwen-32B-AWQ'
snapshot_download(
    model_id=model_id,
    cache_dir=os.environ.get('MODEL_DIR'),
    revision='master'
)
"

1.2 启动模型

使用 HSA_OVERRIDE_GFX_VERSION 环境变量强制指定架构以兼容 RX 7900 XTX 等显卡。

docker run -it \
--network=host \
--group-add=video \
--ipc=host \
--cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined \
--device /dev/kfd \
--device /dev/dri \
-v /opt/qwen_models/Qwen-32B-AWQ/qwen/Qwen3-32B-AWQ:/model \
-e HSA_OVERRIDE_GFX_VERSION=11.0.0 \
rocm/vllm:rocm7.0.0_vllm_0.11.2_20251210 \
vllm serve /model \
--quantization awq \
--dtype float16 \
--served-model-name Qwen3-32B-AWQ \
--trust-remote-code \
--max-model-len 8192

1.3 验证模型

curl http://localhost:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{ "model": "Qwen3-32B-AWQ", "prompt": "你是谁？", "max_tokens": 2000, "temperature": 0.7 }'

2. Qwen3-Embedding 向量化模型

2.1 下载模型

python3 -c "
from modelscope import snapshot_download
snapshot_download('Qwen/Qwen3-Embedding-8B', cache_dir='/opt/qwen_models')
"

2.2 启动模型

注意修改端口号和服务名称。

docker run -d \
--name vllm-embedding \
--restart=always \
--network=host \
--group-add=video \
--ipc=host \
--cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined \
--device /dev/kfd \
--device /dev/dri \
-v /opt/qwen_models/Qwen3-Embedding-8B:/model \
-e HSA_OVERRIDE_GFX_VERSION=11.0.0 \
rocm/vllm:rocm7.0.0_vllm_0.11.2_20251210 \
vllm serve /model \
--port 8001 \
--task embed \
--dtype float16 \
--max-model-len 8192 \
--gpu-memory-utilization 0.4 \
--trust-remote-code \
--served-model-name qwen-embedding

2.3 验证模型

# 检查容器状态
docker ps

# 查看日志
docker logs vllm-embedding

# 发送测试请求
curl http://localhost:8001/v1/embeddings \
-H "Content-Type: application/json" \
-d '{ "model": "qwen-embedding", "input": "你好，测试一下向量化服务" }'

3. Qwen3-Reranker 重排序模型

由于 vLLM 原生暂不支持 Rerank 任务，需通过 FastAPI 自定义服务实现。

3.1 下载模型

python3 -c "
from modelscope import snapshot_download
snapshot_download('Qwen/Qwen3-Reranker-8B', cache_dir='/opt/qwen_models')
"

3.2 配置启动脚本与 uv 管理

在项目目录下创建 Python 服务脚本 rerank_service.py 及依赖配置文件 pyproject.toml。

rerank_service.py

import torch
import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
from transformers import AutoModelForCausalLM, AutoTokenizer

# === 配置区域 ===
MODEL_PATH = "/model"
PORT = 8002

app = FastAPI()
print(f"Loading model from {MODEL_PATH} ...")

# 1. 加载 Tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, padding_side='left', trust_remote_code=True)

# 2. 加载模型
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    trust_remote_code=True,
    device_map="auto",
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2"
).eval()

# 3. 准备 Token IDs
token_false_id = tokenizer.convert_tokens_to_ids("no")
token_true_id = tokenizer.convert_tokens_to_ids("yes")

# 4. 准备前后缀
prefix = "<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\".<|im_end|>\n<|im_start|>user\n"
suffix = "<|im_end|>\n<|im_start|>assistant\n\n\n"
prefix_tokens = tokenizer.encode(prefix, add_special_tokens=False)
suffix_tokens = tokenizer.encode(suffix, add_special_tokens=False)
max_length = 8192

print("Model loaded successfully!")

# === 核心处理逻辑 ===
def format_instruction(instruction, query, doc):
    if instruction is None or instruction == "":
        instruction = 'Given a web search query, retrieve relevant passages that answer the query'
    return "<Instruct>: {instruction}\n<Query>: {query}\n<Document>: {doc}".format(instruction=instruction, query=query, doc=doc)

def process_inputs(pairs):
    inputs = tokenizer(
        pairs,
        padding=False,
        truncation='longest_first',
        return_attention_mask=False,
        max_length=max_length - len(prefix_tokens) - len(suffix_tokens)
    )
    for i, ele in enumerate(inputs['input_ids']):
        inputs['input_ids'][i] = prefix_tokens + ele + suffix_tokens
    inputs = tokenizer.pad(inputs, padding=True, return_tensors="pt", max_length=max_length)
    for key in inputs:
        inputs[key] = inputs[key].to(model.device)
    return inputs

@torch.no_grad()
def compute_scores(inputs):
    batch_scores = model(**inputs).logits[:, -1, :]
    true_vector = batch_scores[:, token_true_id]
    false_vector = batch_scores[:, token_false_id]
    batch_scores = torch.stack([false_vector, true_vector], dim=1)
    batch_scores = torch.nn.functional.log_softmax(batch_scores, dim=1)
    scores = batch_scores[:, 1].exp().tolist()
    return scores

# === API 定义 ===
class RerankRequest(BaseModel):
    model: str = "qwen-reranker"
    query: str
    documents: List[str]
    top_n: Optional[int] = None
    instruction: Optional[str] = None

@app.post("/v1/rerank")
async def rerank(request: RerankRequest):
    try:
        query = request.query
        documents = request.documents
        instruction = request.instruction
        if not documents:
            return {"results": []}
        pairs = [format_instruction(instruction, query, doc) for doc in documents]
        inputs = process_inputs(pairs)
        scores = compute_scores(inputs)
        results = []
        for i, score in enumerate(scores):
            results.append({
                "index": i,
                "relevance_score": float(score),
                "document": documents[i]
            })
        results.sort(key=lambda x: x["relevance_score"], reverse=True)
        if request.top_n:
            results = results[:request.top_n]
        return {
            "model": request.model,
            "results": results,
            "usage": {"total_tokens": inputs.input_ids.numel()}
        }
    except Exception as e:
        print(f"Error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=PORT)

pyproject.toml

[project]
name = "qwen3-reranker-service"
version = "0.1.0"
description = "Rerank service using Qwen3 and ROCm"
readme = "README.md"
requires-python = ">=3.10"
dependencies = [
    "transformers>=4.51.0",
    "fastapi",
    "uvicorn",
    "modelscope",
    "accelerate",
    "pydantic"
]
# 注意：我们故意不把 torch 写在这里，因为我们要用系统自带的 AMD 版本

3.3 构建自定义镜像

为了便于迁移，将依赖环境和业务代码打包进镜像。

# 启动临时构建容器
docker run -it --name builder \
--network=host \
-v /opt/qwen_project:/tmp_build \
rocm/vllm:rocm7.0.0_vllm_0.11.2_20251210 \
bash

# 进入容器后操作
# 1. 安装 uv
pip install uv -i https://pypi.tuna.tsinghua.edu.cn/simple

# 2. 创建应用目录
mkdir -p /app

# 3. 复制代码
cp /tmp_build/rerank_service.py /app/
cp /tmp_build/pyproject.toml /app/

# 4. 进入应用目录
cd /app

# 5. 使用 uv 安装依赖
uv pip install --system -r pyproject.toml -i https://pypi.tuna.tsinghua.edu.cn/simple

# 6. 验证环境
pip list | grep transformers
ls /app

# 7. 退出容器
exit

# 提交为新镜像
docker commit builder qwen-rerank:v1

# 删除临时容器
docker rm builder

# 后续可导出镜像
docker save -o qwen-rerank-v1.tar qwen-rerank:v1

3.4 启动与检验

# 启动最终服务
docker run -d \
--name final_reranker \
--restart=always \
--network=host \
--group-add=video \
--ipc=host \
--cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined \
--device /dev/kfd \
--device /dev/dri \
-v /opt/qwen_models/Qwen3-Reranker-8B:/model \
-e HSA_OVERRIDE_GFX_VERSION=11.0.0 \
qwen-rerank:v1 \
python3 /app/rerank_service.py

# 检查容器运行状态
docker ps

# 查看日志
docker logs final_reranker

# 访问 Swagger UI (http://localhost:8002/docs)

# 发送测试请求
curl http://localhost:8002/v1/rerank \
-H "Content-Type: application/json" \
-d '{ "model": "qwen-reranker", "query": "中国的首都在哪里？", "documents": [ "重力是万有引力。", "中国的首都是北京。", "香蕉很好吃。" ] }'