如何在多GPU上分布式部署PyTorch大模型进行推理？

方法1：Hugging Face Accelerate + 自动设备映射 #

这是最省心的方式，让Accelerate自动将模型层分配到多块GPU，无需手动拆分：

1. 首次运行初始化Accelerate环境（按提示选择多GPU模式，生成配置文件）：

accelerate config

2. 加载模型时指定自动设备映射，同时开启磁盘卸载应对显存溢出：

from accelerate import Accelerator
from transformers import InstructBlipProcessor, InstructBlipForConditionalGeneration

accelerator = Accelerator()
processor = InstructBlipProcessor.from_pretrained("Salesforce/instructblip-vicuna-7b")
model = InstructBlipForConditionalGeneration.from_pretrained(
    "Salesforce/instructblip-vicuna-7b",
    device_map="auto",
    offload_folder="./offload",
    offload_state_dict=True
)

# 准备输入并统一设备
image = processor(images="your_image.jpg", return_tensors="pt").pixel_values
prompt = "Describe this image in detail."
inputs = processor(text=prompt, return_tensors="pt").to(accelerator.device)
image = image.to(accelerator.device)

# 执行推理
outputs = model.generate(
    image=image,
    **inputs,
    max_length=512
)
print(processor.decode(outputs[0], skip_special_tokens=True))

若之前遇到张量设备不匹配错误，检查所有输入是否都通过accelerator.device指定，确保和模型分配的设备一致。

方法2：结合BitsAndBytes量化降低显存压力 #

如果模型体量仍超出单GPU承载，叠加4/8位量化进一步压缩显存占用：

from transformers import InstructBlipForConditionalGeneration, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

model = InstructBlipForConditionalGeneration.from_pretrained(
    "Salesforce/instructblip-vicuna-7b",
    device_map="auto",
    quantization_config=bnb_config
)

量化配合自动设备映射，基本能在4块A10上顺利加载模型。

方法3：DeepSpeed推理专用配置 #

之前用DeepSpeed失败大概率是误用了训练配置，改用推理模式即可实现张量并行：

1. 创建ds_config.json配置文件：

{
  "inference": {
    "tensor_parallel": {
      "enable": true,
      "tp_size": 4
    },
    "enable": true
  }
}

2. 加载并初始化推理模型：

from transformers import InstructBlipForConditionalGeneration
import deepspeed
import torch

model = InstructBlipForConditionalGeneration.from_pretrained("Salesforce/instructblip-vicuna-7b")
model = deepspeed.init_inference(
    model,
    config="./ds_config.json",
    dtype=torch.float16,
    replace_method="auto"
)

# 输入需匹配模型设备，推理步骤同常规流程

该配置会自动将模型张量拆分到4块GPU，避免单GPU加载全量模型。

核心注意事项 #

• 所有输入张量（图像、文本）不要手动指定单GPU（如cuda:0），统一用model.device或accelerator.device获取设备，避免设备不匹配。
• 优先使用transformers最新版本，自动设备映射逻辑在新版本中更稳定。
• 使用Accelerate时，不要手动调用.to(device)移动模型，device_map="auto"会自动完成设备分配。

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