如何用PyTorch在Nvidia A10的4个GPU上实现并发推理？已试DataParallel无效

针对你遇到的单GPU占用、多GPU闲置问题，以下是经过验证的解决思路和实现方案：

一、DataParallel失效的核心原因 #

DataParallel是单进程多线程模式，本质依赖主GPU调度数据，推理场景下（尤其是小batch size）极易出现主GPU负载高、其他GPU闲置的情况，不适合多GPU推理场景。

二、可行解决方案 #

1. 分布式推理（DistributedDataParallel，DDP） #

DDP采用多进程模式，每个GPU对应独立进程，能最大化利用多GPU资源，是推荐的标准方案。

实现步骤：

• 初始化分布式环境，指定NCCL后端（适配Nvidia GPU）
• 每个进程绑定对应GPU，加载并包装模型
• 拆分推理数据，各进程独立处理后汇总结果
• 通过torchrun启动多进程

代码示例：

import torch
import torch.distributed as dist

def setup_distributed():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    torch.cuda.set_device(rank)

def cleanup_distributed():
    dist.destroy_process_group()

def main():
    setup_distributed()
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    # 加载模型并移至对应GPU
    model = YourInferenceModel()
    model = model.to(rank)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank])
    model.eval()

    # 拆分推理数据（按进程数均分）
    total_data = load_your_inference_data()
    chunk_size = len(total_data) // world_size
    start_idx = rank * chunk_size
    end_idx = start_idx + chunk_size if rank != world_size-1 else len(total_data)
    local_data = total_data[start_idx:end_idx]

    # 推理（禁用梯度计算）
    with torch.no_grad():
        local_outputs = model(local_data.to(rank))

    # 汇总结果（按需选择）
    all_outputs = [torch.zeros_like(local_outputs) for _ in range(world_size)]
    dist.all_gather(all_outputs, local_outputs)

    cleanup_distributed()

if __name__ == "__main__":
    main()

启动命令：

torchrun --nproc_per_node=4 your_inference_script.py

2. 手动多进程独立推理 #

无需分布式框架，直接启动4个独立进程，每个进程绑定一块GPU处理部分数据。

实现方式：

• 通过环境变量指定GPU：

# 后台启动4个进程，分别绑定0-3号GPU
CUDA_VISIBLE_DEVICES=0 python your_script.py --data=data_split_0 &
CUDA_VISIBLE_DEVICES=1 python your_script.py --data=data_split_1 &
CUDA_VISIBLE_DEVICES=2 python your_script.py --data=data_split_2 &
CUDA_VISIBLE_DEVICES=3 python your_script.py --data=data_split_3 &

• 或在脚本内指定GPU：

import os
gpu_id = 0  # 每个进程设置不同值（0-3）
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
model = YourModel().to("cuda")

3. 模型流水线并行（Pipeline Parallelism） #

针对大模型（如LLM），将模型不同层分配到不同GPU，实现流水线式推理，充分利用多GPU内存和计算资源。

示例（基于Hugging Face Accelerate）：

from accelerate import Accelerator
from transformers import AutoModelForCausalLM, AutoTokenizer

accelerator = Accelerator()
tokenizer = AutoTokenizer.from_pretrained("your-model-name")
model = AutoModelForCausalLM.from_pretrained("your-model-name")

# 自动分配模型到多GPU
model, tokenizer = accelerator.prepare(model, tokenizer)

# 推理
inputs = tokenizer("your input text", return_tensors="pt").to(accelerator.device)
with torch.no_grad():
    outputs = model.generate(**inputs)

4. TensorRT多GPU优化 #

使用Nvidia TensorRT对模型进行编译优化，原生支持多GPU并行推理，适合高吞吐量场景。

核心步骤：

• 将PyTorch模型导出为ONNX格式
• 通过TensorRT工具将ONNX模型编译为多GPU兼容的引擎
• 每个GPU创建独立推理上下文，并行处理请求

三、排查要点 #

• 确保推理时添加torch.no_grad()，避免不必要的内存占用
• 验证CUDA（≥11.4）、PyTorch（≥1.10）版本与A10显卡兼容
• 检查数据是否正确分配到对应GPU，避免数据集中在主GPU

内容的提问来源于stack exchange，提问作者Murali Mohan