NVIDIA Grace (ARM64) 上使用 vLLM 部署 BGE-M3 Embedding 模型
摘要:本文详细记录了在 NVIDIA Grace (ARM64) 架构服务器上,不使用 Docker 容器,而是直接通过 Python 环境源码运行 vLLM 来部署 BGE-M3 Embedding 服务的全过程。该方案解决了 ARM 架构下 HuggingFace TEI (Text Embeddings Inference) 官方镜像不兼容的问题,并实现了极低的显存占用。
1. 背景与挑战
-
硬件环境:NVIDIA DGX / Spark (基于 NVIDIA Grace CPU),架构为
aarch64。 -
需求:部署 BGE-M3 提供高性能 Embedding (向量化) 服务。
-
痛点:
-
架构兼容性:HuggingFace 官方的 TEI Docker 镜像主要支持 x86_64,在 ARM64 上无法直接运行。
-
编译困难:在 ARM 上从源码编译 Rust 依赖(如 Flash Attention)极其耗时且易出错。
-
解决方案:利用 vLLM 对 ARM 架构的良好支持,直接在 Python 环境中运行 BGE-M3。
2. 环境准备
确保服务器已安装 Miniconda,并且 CUDA 驱动正常。
2.1 创建 Python 环境
建议使用 Python 3.10 或以上版本:
conda create -n vllm-env python=3.11 -y
conda activate vllm-env
2.2 安装 vLLM
由于是在 ARM64 架构上,建议使用 pip 安装(会自动拉取适配 ARM 的 wheel 包)或源码编译。
pip install vllm
2.3 下载模型
推荐使用 ModelScope 国内镜像加速下载 BGE-M3 模型,速度更快且稳定。
# 创建模型目录
mkdir -p /data/models/bge-m3
# 下载模型
python -c "from modelscope import snapshot_download; snapshot_download('Xorbits/bge-m3', cache_dir='./', revision='master')"
# 整理目录结构 (将下载的哈希文件夹内容移动到标准路径)
# 确保 /data/models/bge-m3 下直接包含 config.json 和 pytorch_model.bin / model.safetensors
mv Xorbits/bge-m3/* /data/models/bge-m3/
3. 部署脚本编写
为了规范管理,建议在项目目录下创建独立的运行时目录,并使用脚本启动服务。
3.1 创建目录结构
# 假设项目根目录为 /data/workspaces/vllm-project
cd /data/workspaces/vllm-project
mkdir -p bge_m3_runtime
cd bge_m3_runtime
3.2 编写启动脚本 start_bge_m3.sh
请在 bge_m3_runtime 目录下创建该脚本。
关键配置说明:
--gpu-memory-utilization 0.05:BGE-M3 模型参数量仅 569M,FP16 精度下占用极低。设置 5% 的显存占用足以运行,避免浪费宝贵的显存资源。- **移除
--task embed**:新版 vLLM 会自动识别 BERT/RoBERTa 类架构为 Embedding 任务,显式指定该参数在某些版本中可能会报错unrecognized arguments。 VLLM_WORKER_MULTIPROC_METHOD="spawn":关键参数。在 ARM/Blackwell 架构下,Python 默认的 fork 模式容易导致 CUDA 初始化死锁,必须强制设置为spawn。
#!/bin/bash
# ================= 配置区 =================
# 1. 项目代码根目录
WORKDIR="/data/workspaces/vllm-project"
# 2. 运行时目录 (存放日志和PID)
RUNTIME_DIR="/data/workspaces/vllm-project/bge_m3_runtime"
# 3. 日志与进程锁
LOG_FILE="$RUNTIME_DIR/server.log"
PID_FILE="$RUNTIME_DIR/server.pid"
# 4. Python 环境路径 (根据实际 Conda 路径修改)
PYTHON_PATH="/root/miniconda3/envs/vllm-env/bin/python"
# 5. 模型绝对路径
MODEL_PATH="/data/models/bge-m3"
# ================= 环境变量 =================
# 显式指定后端,防止自动探测出错
export VLLM_TARGET_DEVICE="cuda"
# 解决多进程死锁问题 (ARM/Blackwell 必备)
export VLLM_WORKER_MULTIPROC_METHOD="spawn"
# 指定 Tokenizer 缓存路径,防止内网下载失败
export TIKTOKEN_CACHE_DIR="/root/.cache/tiktoken"
# ================= 预检逻辑 =================
if [ -f "$PID_FILE" ]; then
OLD_PID=$(cat "$PID_FILE")
if ps -p "$OLD_PID" > /dev/null; then
echo "❌ BGE-M3 Service is already running (PID: $OLD_PID)."
echo " Stop it first: kill $OLD_PID"
exit 1
else
echo "⚠️ Cleaning up stale PID file..."
rm "$PID_FILE"
fi
fi
# ================= 启动逻辑 =================
echo "🚀 Starting BGE-M3 Embedding Service on Port 8023..."
echo "📂 Runtime Dir: $RUNTIME_DIR"
echo "📝 Logs: $LOG_FILE"
cd "$WORKDIR" || exit
# 启动命令
# --max-model-len 8192: BGE-M3 支持的最大长度
# --gpu-memory-utilization 0.05: 极低显存占用模式
nohup "$PYTHON_PATH" -m vllm.entrypoints.openai.api_server \
--model "$MODEL_PATH" \
--served-model-name bge-m3 \
--trust-remote-code \
--dtype bfloat16 \
--gpu-memory-utilization 0.05 \
--max-model-len 8192 \
--host 0.0.0.0 --port 8023 \
> "$LOG_FILE" 2>&1 &
PID=$!
echo "✅ Service started with PID: $PID"
echo "$PID" > "$PID_FILE"
3.3 赋权并启动
chmod +x start_bge_m3.sh
./start_bge_m3.sh
4. 验证与测试
4.1 查看日志
启动需要几秒钟加载模型,可以通过日志观察进度:
tail -f server.log
当看到 Application startup complete 字样,且日志中包含 Model 'bge-m3' is an embedding model 的提示时,说明服务启动成功。
4.2 API 接口测试
vLLM 提供了兼容 OpenAI 格式的 API 接口。
curl http://localhost:8023/v1/embeddings \
-H "Content-Type: application/json" \
-d '{
"model": "bge-m3",
"input": "测试文本向量化服务"
}'
预期输出(返回向量数据):
{
"id": "embd-xxxx",
"object": "list",
"created": 1769152010,
"model": "bge-m3",
"data": [
{
"index": 0,
"object": "embedding",
"embedding": [
-0.02145374,
0.01234567,
... (省略长向量) ...
0.039404828
]
}
],
"usage": {
"prompt_tokens": 6,
"total_tokens": 6,
"completion_tokens": 0
}
}
5. 总结
在 ARM64 架构上部署 AI 模型时,传统的 Docker 镜像方案可能会遇到架构不兼容问题(尤其是涉及 Rust/CUDA 底层优化的库)。直接使用 Python 环境运行 vLLM,并合理配置显存参数和多进程启动方式,是一种灵活且高效的替代方案。
此方案不仅发挥了 NVIDIA Grace 平台的性能,还通过极低的显存占用(仅需 ~5%),为服务器上运行其他大型模型留出了充足的空间。
更多推荐
所有评论(0)