摘要:本文详细记录了在 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 目录下创建该脚本。

关键配置说明:

  1. --gpu-memory-utilization 0.05:BGE-M3 模型参数量仅 569M,FP16 精度下占用极低。设置 5% 的显存占用足以运行,避免浪费宝贵的显存资源。
  2. **移除 --task embed**:新版 vLLM 会自动识别 BERT/RoBERTa 类架构为 Embedding 任务,显式指定该参数在某些版本中可能会报错 unrecognized arguments
  3. 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%),为服务器上运行其他大型模型留出了充足的空间。

Logo

脑启社区是一个专注类脑智能领域的开发者社区。欢迎加入社区,共建类脑智能生态。社区为开发者提供了丰富的开源类脑工具软件、类脑算法模型及数据集、类脑知识库、类脑技术培训课程以及类脑应用案例等资源。

更多推荐