如何在多GPU加载Hugging Face模型并执行推理（非训练/微调）

当然可以在多GPU环境下加载Hugging Face模型并执行推理，以下分别给出基于Accelerate和DeepSpeed的实现方案，均针对推理场景优化：

1. 安装依赖

pip install accelerate

2. 修改后的推理代码

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from accelerate import Accelerator

# 初始化Accelerator，自动识别多GPU环境
accelerator = Accelerator()

tokenizer = AutoTokenizer.from_pretrained("togethercomputer/LLaMA-2-7B-32K")
# 加载模型时指定device_map="auto"，让Accelerate自动将模型层分配到可用GPU
model = AutoModelForCausalLM.from_pretrained(
    "togethercomputer/LLaMA-2-7B-32K",
    torch_dtype=torch.float16,
    device_map="auto"
)

# 用Accelerator完成模型的设备同步与配置
model = accelerator.prepare(model)

input_context = "Your text here"
input_ids = tokenizer.encode(input_context, return_tensors="pt").to(accelerator.device)

# 启用自动混合精度推理
with accelerator.autocast():
    output = model.generate(input_ids, max_length=256, temperature=0.7)

# 解码并输出结果
output_text = tokenizer.decode(output[0], skip_special_tokens=True)
print(output_text)

• 关键说明：device_map="auto"会根据GPU显存情况自动拆分模型层到不同GPU，无需手动指定；Accelerator会自动处理设备同步、混合精度等细节。

1. 安装依赖

pip install deepspeed

2. 创建DeepSpeed推理配置文件（命名为ds_config.json）

{
  "inference": true,
  "tensor_parallel": {
    "tp_size": 2  // 对应你的GPU数量
  },
  "zero_optimization": {
    "stage": 0  // 推理场景下无需Zero优化，设为0即可
  },
  "fp16": {
    "enabled": true
  }
}

3. 修改后的推理代码

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import deepspeed

tokenizer = AutoTokenizer.from_pretrained("togethercomputer/LLaMA-2-7B-32K")
model = AutoModelForCausalLM.from_pretrained(
    "togethercomputer/LLaMA-2-7B-32K",
    torch_dtype=torch.float16
)

# 初始化DeepSpeed推理引擎，指定GPU数量与配置文件
model = deepspeed.init_inference(
    model,
    mp_size=2,  // GPU数量，需与配置文件中tp_size一致
    dtype=torch.float16,
    config="ds_config.json"
)

input_context = "Your text here"
input_ids = tokenizer.encode(input_context, return_tensors="pt").to(model.device)

# 执行推理
output = model.generate(input_ids, max_length=256, temperature=0.7)

# 解码并输出结果
output_text = tokenizer.decode(output[0], skip_special_tokens=True)
print(output_text)

4. 运行命令

deepspeed --num_gpus 2 your_script.py

• 关键说明：DeepSpeed通过张量并行（Tensor Parallelism）将模型拆分到多GPU，适合大模型推理；mp_size需与配置文件中的tp_size保持一致。

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