如何通过模型并行在多GPU上加载HuggingFace模型进行推理？

针对单卡显存不足的问题，以下是几种实用的多卡分布式加载模型的方法：

1. 自动设备映射（最简方案） #

直接在加载模型时指定device_map="auto"，Transformers库会自动将模型的不同层分配到可用的GPU上，充分利用多卡显存：

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("google/ul2")
# 开启自动设备映射，同时启用低CPU内存占用模式
model = AutoModelForSeq2SeqLM.from_pretrained(
    "google/ul2",
    device_map="auto",
    low_cpu_mem_usage=True
)

# 推理示例
inputs = tokenizer("你的输入文本", return_tensors="pt").to("cuda")
outputs = model.generate(**inputs)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

device_map="auto"会根据各GPU的剩余显存自动分配模型层，无需手动配置。

2. 平衡式设备映射 #

如果需要更均衡地分配模型到多卡，可以使用device_map="balanced"，它会将模型参数平均分配到所有可用GPU：

model = AutoModelForSeq2SeqLM.from_pretrained(
    "google/ul2",
    device_map="balanced",
    low_cpu_mem_usage=True
)

3. 使用Accelerate库实现分布式推理 #

对于更复杂的分布式场景，比如需要控制模型分配细节，可使用HuggingFace Accelerate库：

from accelerate import Accelerator
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# 初始化加速器
accelerator = Accelerator()

tokenizer = AutoTokenizer.from_pretrained("google/ul2")
model = AutoModelForSeq2SeqLM.from_pretrained("google/ul2", low_cpu_mem_usage=True)

# 准备模型（自动适配多卡环境）
model = accelerator.prepare(model)

# 推理示例
inputs = tokenizer("你的输入文本", return_tensors="pt")
# 将输入移到合适的设备
inputs = accelerator.prepare(inputs)
outputs = model.generate(**inputs)
# 解码结果
outputs = accelerator.gather(outputs)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Accelerate会自动处理多卡环境的设备分配、数据同步等逻辑，适合自定义推理流程。

注意事项 #

• 确保你的Transformers库版本足够新（建议>=4.20.0），device_map功能在较新版本中才稳定支持。
• 推理时尽量控制输入批量大小，避免单步推理占用过多显存。
• 如果模型仍然过大，可结合load_in_8bit=True或load_in_4bit=True（需要bitsandbytes库）进一步压缩模型显存占用，配合多卡使用效果更佳。

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