vLLM与Torchrun结合的大模型分布式推理实践

1. 为什么需要Torchrun与vLLM结合

在大模型推理场景中，单卡GPU的内存容量往往成为瓶颈。以7B参数的模型为例，仅加载FP16精度的模型就需要约14GB显存，这还未计算推理过程中的KV Cache等开销。vLLM（Vectorized Large Language Model）作为新一代推理引擎，通过PagedAttention等创新技术优化了显存利用率，而PyTorch的torchrun工具则为分布式训练/推理提供了基础设施支持。

实际部署中常见两种并行策略：

• 张量并行（Tensor Parallelism）：将模型参数拆分到多个设备
• 流水线并行（Pipeline Parallelism）：按层拆分模型

vLLM当前主要支持张量并行方式，通过设置tensor_parallel_size参数即可启用。而torchrun的作用是管理多进程环境，使开发者无需手动处理进程间通信（IPC）等底层细节。二者结合后，可以实现以下典型场景：

• 单机多卡推理（如8卡A100服务器）
• 多机多卡推理（需配合NCCL等通信库）
• 弹性推理（动态增减计算资源）

关键配置经验：当tensor_parallel_size大于1时，必须使用torchrun或类似工具启动，否则会导致进程 hanging。

2. 基础环境搭建与验证

2.1 硬件与驱动要求

推荐配置：

• GPU：NVIDIA A100/A10G（Ampere架构最佳）
• 驱动版本：>=525.60.13
• CUDA：11.8或12.1
• 内存：每卡至少预留5GB系统内存

验证环境完整性的命令：

bash复制nvidia-smi  # 确认GPU识别正常
nvcc --version  # 检查CUDA版本
python -c "import torch; print(torch.cuda.is_available())"  # 验证PyTorch CUDA支持

2.2 依赖安装

使用conda创建隔离环境：

bash复制conda create -n vllm python=3.9 -y
conda activate vllm
pip install vllm torch>=2.1.2

常见安装问题处理：

• 遇到Could not build wheels for vllm：升级setuptools到最新版
• CUDA版本不匹配：通过conda install cuda -c nvidia指定版本
• 内存不足：添加--no-cache-dir参数

3. 最小可运行示例解析

3.1 单卡基础启动

创建simple.py：

python复制from vllm import LLM

llm = LLM(model="meta-llama/Llama-2-7b-chat-hf")
output = llm.generate("Explain quantum computing")
print(output)

启动命令：

bash复制python simple.py

3.2 多卡Torchrun启动

创建torchrun_example.py：

python复制import os
from vllm import LLM, SamplingParams

def main():
    # 必须从环境变量获取rank信息
    local_rank = int(os.getenv("LOCAL_RANK", "0"))
    
    # 模型配置
    model = "Qwen/Qwen1.5-7B-Chat"
    tp_size = int(os.getenv("TP_SIZE", "2"))  # 张量并行度
    
    # 初始化LLM实例
    llm = LLM(
        model=model,
        tensor_parallel_size=tp_size,
        trust_remote_code=True
    )
    
    # 只有rank0进程执行生成
    if local_rank == 0:
        sampling_params = SamplingParams(temperature=0.7, top_p=0.9)
        outputs = llm.generate(
            ["The future of AI is"],
            sampling_params=sampling_params
        )
        print(outputs[0].text)

if __name__ == "__main__":
    main()

启动脚本start.sh：

bash复制#!/bin/bash
TP_SIZE=2  # 等于GPU数量
torchrun --nproc_per_node=$TP_SIZE torchrun_example.py

关键参数说明：

• --nproc_per_node：每台机器的进程数（通常等于GPU数）
• LOCAL_RANK：自动分配的进程标识符（0到N-1）
• TP_SIZE：必须与nproc_per_node一致

4. 生产级部署实践

4.1 性能调优参数

在LLM初始化时添加优化参数：

python复制llm = LLM(
    model=model,
    tensor_parallel_size=4,
    max_model_len=4096,  # 控制最大序列长度
    gpu_memory_utilization=0.9,  # 显存利用率
    enforce_eager=True,  # 禁用CUDA Graph（调试用）
    kv_cache_dtype="fp8"  # KV Cache量化
)

4.2 多模型负载均衡

通过--master_port参数支持多组服务：

bash复制# 第一组服务
torchrun --nproc_per_node=2 --master_port=29500 serve_model_a.py

# 第二组服务（不同端口）
torchrun --nproc_per_node=2 --master_port=29501 serve_model_b.py

4.3 常见错误排查

1. CUDA out of memory

   • 降低gpu_memory_utilization（默认0.9）
   • 减小max_model_len
   • 启用swap_space参数使用磁盘交换

2. NCCL通信失败

   bash复制export NCCL_DEBUG=INFO
   export NCCL_SOCKET_IFNAME=eth0  # 指定网卡
   

3. 进程hang住

   • 确认所有rank的TP_SIZE一致
   • 检查防火墙设置
   • 添加--standalone参数测试单机模式

5. 进阶技巧与监控

5.1 动态批处理配置

在服务端脚本中添加：

python复制from vllm.engine.arg_utils import AsyncEngineArgs

engine_args = AsyncEngineArgs(
    model=model,
    tensor_parallel_size=tp_size,
    max_num_seqs=256,  # 最大批处理量
    max_paddings=1024  # 最大padding长度
)
llm = LLM.from_engine_args(engine_args)

5.2 Prometheus监控集成

添加监控端点：

python复制from vllm import AsyncLLMEngine
from vllm.entrypoints.openai import api_server

engine = AsyncLLMEngine.from_engine_args(engine_args)
app = api_server.app(engine)

# 添加/metrics端点
@app.route("/metrics")
def metrics():
    from prometheus_client import generate_latest
    return generate_latest()

5.3 自定义分布式策略

继承DistributedExecutorBackend：

python复制from vllm.executor.distributed_executor import DistributedExecutorBackend

class CustomBackend(DistributedExecutorBackend):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 自定义初始化逻辑
        
    def execute(self, *args, **kwargs):
        # 自定义执行逻辑
        return super().execute(*args, **kwargs)

启动时指定后端：

bash复制torchrun --nproc_per_node=2 --executor_backend=custom serve.py

6. 性能对比实测

测试环境：2x A100 80GB，Llama2-13B模型

关键发现：

1. 张量并行在TP=2时达到近线性加速
2. vLLM的PagedAttention显著降低显存占用
3. 当TP超过物理GPU数时性能下降明显

7. 容器化部署方案

7.1 Dockerfile示例

dockerfile复制FROM nvidia/cuda:12.1.1-base

RUN apt-get update && \
    apt-get install -y python3.9 python3-pip && \
    ln -s /usr/bin/python3.9 /usr/bin/python

WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .
ENTRYPOINT ["torchrun", "--nproc_per_node=2", "serve.py"]

7.2 Kubernetes部署要点

1. 资源请求配置：

   yaml复制resources:
     limits:
       nvidia.com/gpu: "2"
     requests:
       cpu: "4"
       memory: 16Gi
   

2. 使用StatefulSet保证rank稳定性

3. 通过Headless Service实现Pod发现

8. 模型热加载实践

实现不中断服务的模型切换：

python复制import signal
from vllm import LLM

current_llm = LLM(model="model_v1")

def reload_model(signum, frame):
    global current_llm
    new_llm = LLM(model="model_v2")
    current_llm = new_llm

signal.signal(signal.SIGHUP, reload_model)

触发方式：

bash复制kill -HUP <pid>