【AI大模型】Transformers大模型库（七）：单机多卡推理之device_map

【AI大模型】Transformers大模型库（七）：单机多卡推理之device_map-LMLPHP

一、引言

二、单机多卡推理之device_map

2.1 概述

2.2 自动配置，如device_map="auto"

2.3 手动配置，如device_map="cuda:1"

三、总结

一、引言

这里的Transformers指的是huggingface开发的大模型库，为huggingface上数以万计的预训练大模型提供预测、训练等服务。

本文重点介绍如何使用device_map进行单机多卡推理。

二、单机多卡推理之device_map

2.1 概述

device_map是一个在使用Hugging Face的transformers库，特别是处理大型模型和多GPU环境时非常重要的参数。它用于指定模型的各个部件应加载到哪个具体的计算设备上，以实现资源的有效分配和利用。这个参数在进行模型并行或分布式训练时特别有用。

2.2 自动配置，如device_map="auto"

使用device_map="auto"时，Hugging Face的transformers库会尝试自动分配模型的各个部件到可用的GPU设备上，以实现最佳的并行计算和资源利用。这个选项特别适用于拥有多个GPU的环境，它会根据每个GPU的内存大小和其他运行中的进程智能地分配模型的层，尽量做到负载均衡。

安装transformers和accelerate两个python包

pip install transformers -i https://mirrors.cloud.tencent.com/pypi/simple
pip install accelerate -i https://mirrors.cloud.tencent.com/pypi/simple

采用CUDA_VISIBLE_DEVICES=1,2,3指定推理代码可见的GPU设备。

CUDA_VISIBLE_DEVICES=1,2,3 python trans_glm4.py

在采用AutoModelForCausalLM.from_pretrained模型加载时，加入device_map="auto"，模型会自动分配至CUDA_VISIBLE_DEVICES指定的GPU显卡，编号从0开始

model = AutoModelForCausalLM.from_pretrained(model_dir,device_map="auto",trust_remote_code=True,torch_dtype=torch.float16)

除了device_map="auto"，还包括balanced、balanced_low_0、sequential等参数

2.3 手动配置，如device_map="cuda:1"

假设想要模型的某些部分在第一张显卡，另一部分在第二张显卡，需要知道模型的层名或者按照模型的组件大小进行合理分配。不过，具体层名需要根据实际模型来确定，这里提供一个概念性的示例：

device_map = {
    "transformer.h.0": "cuda:0",  # 第一部分放在GPU 0
    "transformer.h.1": "cuda:1",  # 第二部分放在GPU 1
    # ... 根据模型结构继续分配
}
model = AutoModelForCausalLM.from_pretrained(model_dir, device_map=device_map)