【AI大模型】Transformers大模型库(七):单机多卡推理之device_map

目录​​​​​​​

一、引言

二、单机多卡推理之device_map

[2.1 概述](#2.1 概述)

[2.2 自动配置,如device_map="auto"](#2.2 自动配置,如device_map="auto")

[2.3 手动配置,如device_map="cuda:1"](#2.3 手动配置,如device_map="cuda:1")

三、总结


一、引言

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

🤗 Transformers 提供了数以千计的预训练模型,支持 100 多种语言的文本分类、信息抽取、问答、摘要、翻译、文本生成。它的宗旨是让最先进的 NLP 技术人人易用。

🤗 Transformers 提供了便于快速下载和使用的API,让你可以把预训练模型用在给定文本、在你的数据集上微调然后通过 model hub 与社区共享。同时,每个定义的 Python 模块均完全独立,方便修改和快速研究实验。

🤗 Transformers 支持三个最热门的深度学习库: Jax, PyTorch 以及 TensorFlow --- 并与之无缝整合。你可以直接使用一个框架训练你的模型然后用另一个加载和推理。

本文重点介绍如何使用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设备。

python 复制代码
CUDA_VISIBLE_DEVICES=1,2,3 python trans_glm4.py 

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

python 复制代码
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等参数

  • "auto" 和 "balanced" 将会在所有的GPU上平衡切分模型。主要是有可能发现更高效的分配策略。"balanced" 参数的功能则保持稳定。
  • "balanced_low_0" 会在除了第一个GPU上的其它GPU上平衡划分模型,并且在第一个 GPU 上占据较少资源。这个选项符合需要在第一个 GPU 上进行额外操作的需求,例如需要在第一个 GPU 执行 generate 函数
  • "sequential" 按照GPU的顺序分配模型分片,从 GPU 0 开始,直到最后的 GPU(那么最后的 GPU 往往不会被占满,和 "balanced_low_0" 的区别就是第一个还是最后一个,以及非均衡填充)

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

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

python 复制代码
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)

三、总结

本文简要介绍了device_map="auto"等使用方法,多数情况下与CUDA_VISIBLE_DEVICES=1,2,3一起使用,可以简单高效的进行多卡分布式推理及训练计算,至于多机多卡场景,多用torchrun和deepspeed等,后面文章会专门进行系统讲解。

如果您还有时间,可以看看我的其他文章:

《AI---工程篇》

AI智能体研发之路-工程篇(一):Docker助力AI智能体开发提效

AI智能体研发之路-工程篇(二):Dify智能体开发平台一键部署

AI智能体研发之路-工程篇(三):大模型推理服务框架Ollama一键部署

AI智能体研发之路-工程篇(四):大模型推理服务框架Xinference一键部署

AI智能体研发之路-工程篇(五):大模型推理服务框架LocalAI一键部署

《AI---模型篇》

AI智能体研发之路-模型篇(一):大模型训练框架LLaMA-Factory在国内网络环境下的安装、部署及使用

AI智能体研发之路-模型篇(二):DeepSeek-V2-Chat 训练与推理实战

AI智能体研发之路-模型篇(三):中文大模型开、闭源之争

AI智能体研发之路-模型篇(四):一文入门pytorch开发

AI智能体研发之路-模型篇(五):pytorch vs tensorflow框架DNN网络结构源码级对比

AI智能体研发之路-模型篇(六):【机器学习】基于tensorflow实现你的第一个DNN网络

AI智能体研发之路-模型篇(七):【机器学习】基于YOLOv10实现你的第一个视觉AI大模型

AI智能体研发之路-模型篇(八):【机器学习】Qwen1.5-14B-Chat大模型训练与推理实战

AI智能体研发之路-模型篇(九):【机器学习】GLM4-9B-Chat大模型/GLM-4V-9B多模态大模型概述、原理及推理实战

《AI---Transformers应用》

【AI大模型】Transformers大模型库(一):Tokenizer

【AI大模型】Transformers大模型库(二):AutoModelForCausalLM

【AI大模型】Transformers大模型库(三):特殊标记(special tokens)

【AI大模型】Transformers大模型库(四):AutoTokenizer

【AI大模型】Transformers大模型库(五):AutoModel、Model Head及查看模型结构

相关推荐
HPC_fac130520678161 小时前
以科学计算为切入点:剖析英伟达服务器过热难题
服务器·人工智能·深度学习·机器学习·计算机视觉·数据挖掘·gpu算力
网易独家音乐人Mike Zhou3 小时前
【卡尔曼滤波】数据预测Prediction观测器的理论推导及应用 C语言、Python实现(Kalman Filter)
c语言·python·单片机·物联网·算法·嵌入式·iot
安静读书3 小时前
Python解析视频FPS(帧率)、分辨率信息
python·opencv·音视频
小陈phd3 小时前
OpenCV从入门到精通实战(九)——基于dlib的疲劳监测 ear计算
人工智能·opencv·计算机视觉
Guofu_Liao4 小时前
大语言模型---LoRA简介;LoRA的优势;LoRA训练步骤;总结
人工智能·语言模型·自然语言处理·矩阵·llama
小二·5 小时前
java基础面试题笔记(基础篇)
java·笔记·python
小喵要摸鱼6 小时前
Python 神经网络项目常用语法
python
秀儿还能再秀7 小时前
神经网络(系统性学习三):多层感知机(MLP)
神经网络·学习笔记·mlp·多层感知机
一念之坤8 小时前
零基础学Python之数据结构 -- 01篇
数据结构·python
wxl7812278 小时前
如何使用本地大模型做数据分析
python·数据挖掘·数据分析·代码解释器