如何准确的估计llm推理和微调的内存消耗

本文涉及的产品
实时数仓Hologres,5000CU*H 100GB 3个月
大数据开发治理平台 DataWorks,不限时长
检索分析服务 Elasticsearch 版,2核4GB开发者规格 1个月
简介: 最近发布的三个大型语言模型——Command-R+ (104B参数), Mixtral-8x22b (141B参数的MoE模型), 和 Llama 3 70b (70.6B参数)——需要巨大的内存资源。推理时,Command-R+需193.72GB GPU RAM,Mixtral-8x22B需262.63GB,Llama 370b需131.5GB。激活的内存消耗根据序列长度、批大小等因素变化。文章详细介绍了计算这些模型内存需求的方法,并探讨了如何通过量化、优化器优化和梯度检查点减少内存使用,以适应微调和推理。

Command-R+, Mixtral-8x22b和Llama 3 70b都在最近的几周内发布了,这些模型是巨大的。它们都有超过700亿个参数:

Command-R+: 104B参数

Mixtral-8x22b:具有141B参数的混合专家(MoE)模型

Llama 370b: 70.6B参数

你能在电脑上微调和运行这些模型吗?

在本文中,我将介绍如何计算这些模型用于推理和微调的最小内存。这种方法适用于任何的llm,并且精确的计算内存总消耗。

推理所需的内存

这三个模型都以16位权重发布:Command-R+为float16, Mixtral和Llama 370b为bfloat16。也就是说一个参数消耗16位或2字节的内存。

10亿个参数则将占用20亿个字节,或者说10亿个字节等于1GB,那么1B个参数占用2GB的内存。100B参数就需要占用200GB内存。这是一个近似值,因为1 KB不等于1,000字节,而是1,024字节。我们通过这种简单的方法可以大概评估内存的占用,后面我们还会有详细的计算过程。

要知道一个模型有多少个参数而不需要下载,可以查看模型卡:

如果在GPU上做快速推理,需要将模型完全加载到GPU RAM上。

对于Command-R+: 193.72 GB的GPU RAM

对于Mixtral-8x22B: 262.63 GB的GPU RAM

对于Llama 370b: 131.5 GB的GPU RAM

或者说你将需要2x80 GB的GPU,例如两个h100。

激活的内存消耗

一旦模型被加载,我们需要更多的内存来存储模型的激活,即在推理过程中创建的张量。这些张量从一层传递到下一层。它们在内存中的大小并不容易估计。为了估计这些激活的内存消耗,我使用了《Reducing Activation Recomputation in Large Transformer Models》提出的方法。

我们需要知道以下内容来估计内存消耗:

 s:最大序列长度(输入中的令牌数量)
 b:批大小
 h:模型的隐藏维度
 a:注意头的数量

标准transformer 层由自注意力块和MLP块组成,每个块由两个layer-norms连接。每个组件的内存消耗估计如下。

1、注意力块

注意块由自注意力机制、线性投射和dropout 组成。内存要求包括:

 线性投影保留其输入激活,其大小为2sbh,而dropout 需要占用sbh的掩码。
 线性投影和自注意力输入激活,各需要2sbh。
 查询(Q)和键(K)矩阵需要4sbh。
 Softmax和它的dropout分别需要2as²b和as²b。
 值(V)上应用的注意存储(V)加起来等于2as²b + 2sbh。

所以注意块所需的总内存为11sbh + 5as²b。

这里的sbh等于 s b h,我们下面也会这样简写。

2、MLP块

MLP块包括两个线性层和一个dropout:

线性层存储花费2sbh和8sbh的输入,GeLU非线性也需要8sbh。我们不需要为推理存储GeLU的激活,但我仍然会对它们进行计数,以防某些推理的框架存储这些激活。Dropout存储的掩码大小为sbh。

所以MLP块需要19sbh的空间。

3、layer-norms

每个LN存储它的输入需要2sbh,两层就为4sbh。

4、对于整个层

总内存需求总计为11sbh + 5as²b(来自注意力块)+ 19sbh(来自MLP块)+ 4sbh(来自LN)。

每层激活内存消耗= 34 sbh + 5as²b

如果我们使用16位数据类型,那么需要将这个数字乘以2,因为每个激活参数将需要2字节。

5、总计

这个等式大致近似于实际内存消耗。大多数推理框架都经过优化,通过在张量无用时立即删除它们,所以一般情况下会比这个数值少。但是在推理期间也会创建各种消耗内存的缓冲区。但是通过实验,我发现我们计算的这个数值基本上近似于Hugging Face的Transformers框架。而对于vLLM和TGI等其他优化得更好的框架,内存消耗则会减少。

如果使用FlashAttention、Alibi或RoPE等高级技术,处理长序列的内存消耗也将大大减少。

6、估算用于推理的Command-R+、Mixtral-8x22B和Llama 370b的内存消耗

在模型的参数和激活都是16位的标准场景下,我们还需要设置解码超参数。

S = 512(序列长度),B = 8(批量大小)

与内存中模型的大小相比,激活的大小可以忽略不计。但是它们的大小会随着批大小和序列长度的增加而迅速增加。

减少推理的内存消耗

大部分用于推理的内存消耗来自模型的参数。最近的量化算法可以显著减少这种内存消耗。他们通过减少大多数参数的位宽来压缩模型,同时尽量保持模型的准确性。

8位量化几乎是无损的,而4位量化只会略微降低性能。4位量化将模型的内存消耗除以4,因为大多数参数都是4位,即0.5字节而不是2字节。我推荐使用AWQ进行4位量化,它运行简单,生成快速模型。

对于超过100B个参数的非常大的模型,精度较低的量化,例如2.5位或3位,仍然可以得到准确的结果。例如,AQLM对Mixtral-8x7B的2位量化表现出良好的性能。但是AQLM的问题在于量化模型的成本非常高。对于非常大的模型,可能需要几个星期的时间。

另一种选择是将模型移动到另一个存储设备,例如,CPU RAM。但是他的缺点和明显,太慢了,特别是在批处理解码时。针对CPU推理进行优化的框架有助于保持推理的合理速度。例如,Neural Speed是使用量化模型在CPU上进行推理的最快框架之一。

如果使用CPU,那么仍然需要大量的CPU RAM来加载模型和存储激活,计算的方法是相同的。

微调所需的内存

对于微调llm,估计内存消耗稍微复杂一些。除了存储模型权重和激活之外,对于所有层,我们还需要存储优化器状态。

优化器状态的内存消耗

AdamW优化器是最流行的微调llm,它为模型的每个参数创建并存储2个新参数。如果我们有一个100B的模型,优化器将创建200B的新参数!为了更好的训练稳定性,优化器的参数为float32,即每个参数占用4字节的内存。

这就是微调比推理消耗更多内存的主要原因。

例如,对于Mixtral-8x22B,优化器创建2*141B = 282B float32参数。它消耗了1053.53 GB的内存,我们必须加上模型本身占用的内存,即262.63 GB。所以总共需要1315.63 GB的GPU内存。这大约是17个80gb的A100 !

并且这还不足以对模型进行微调。我们还需要内存来存储激活。

计算梯度所需的内存

与推理相比,我们只需要在传递给下一层之前存储单个层的激活,微调需要存储前向传播过程中创建的所有激活。这对于计算梯度是必要的,然后用于反向传播误差并更新模型的权重。

为了估计计算梯度所需的内存,我们可以使用用于推理的相同公式,然后将结果乘以层数。

如果L是层数,那么计算梯度所消耗的内存为

 L(34sbh + 5as²b)

估算Llama 3 70b、Mixtral-8x22B和Command R+微调的内存消耗

我们需要估计模型的大小,并添加所有层的激活大小和优化器状态的大小。

我设置了以下超参数进行微调:

S = 512(序列长度)

B = 8(批量大小)

对于优化器状态,我假设它们是float32。

所以得到

这是最坏情况下的内存消耗,也就是说没有使用任何优化来减少内存消耗。幸运的是,我们可以应用许多优化来减少内存需求。

减少微调内存消耗

由于优化器状态消耗大量内存,因此已经进行了大量的研究来减少它们的内存占用,例如:

LoRA:冻结整个模型,并添加一个具有数百万个参数的可训练适配器。使用LoRA,我们只存储适配器参数的优化器状态。

QLoRA: LoRA,但模型量化为4位或更低精度。

AdaFactor和AdamW-8bit:更高效的内存优化器,提供接近AdamW的性能。但AdaFactor在训练期间可能不稳定。

GaLore:将梯度投影到低秩子空间中,这可以将优化器状态的大小减少80%。

内存的另一个重要部分被激活所消耗。为了减少它通常采用梯度检查点。当需要计算梯度时,它会重新计算一些激活。它减少了内存消耗但也减慢了微调速度。

最后,还有一些框架,如Unsloth,在使用LoRA和QLoRA进行微调方面进行了极大的优化。

总结

在本文中,我们介绍了如何估计transformer 模型的内存消耗。这个方法不适用于transformer以外的其他体系结构的模型。例如,曼巴和RWKV的激活记忆消耗明显较低,因为它们不使用注意力机制。

https://avoid.overfit.cn/post/0046a7ef3a47406e9ed98d4ce947a14d

相关实践学习
部署Stable Diffusion玩转AI绘画(GPU云服务器)
本实验通过在ECS上从零开始部署Stable Diffusion来进行AI绘画创作,开启AIGC盲盒。
目录
相关文章
|
1月前
|
机器学习/深度学习 缓存
Block Transformer:通过全局到局部的语言建模加速LLM推理
Block Transformer是一种优化自回归语言模型推理效率的新架构,通过块级自注意力来平衡全局和局部依赖,提高吞吐量。模型包含嵌入器、块解码器和令牌解码器,其中块解码器处理全局依赖,令牌解码器处理局部细节。这种方法减轻了KV缓存的延迟和内存开销,尤其是在长序列处理中。实验显示,尽管Block Transformer参数量增加,但推理速度显著提升,尤其是在大块长度和优化的组件比例下,实现了性能与速度的平衡。
247 7
|
2月前
|
物联网 测试技术 API
LLM 大模型学习必知必会系列(九):Agent微调最佳实践,用消费级显卡训练属于自己的Agent!
LLM 大模型学习必知必会系列(九):Agent微调最佳实践,用消费级显卡训练属于自己的Agent!
LLM 大模型学习必知必会系列(九):Agent微调最佳实践,用消费级显卡训练属于自己的Agent!
|
12天前
|
并行计算 PyTorch 算法框架/工具
LLM推理引擎怎么选?TensorRT vs vLLM vs LMDeploy vs MLC-LLM
有很多个框架和包可以优化LLM推理和服务,所以在本文中我将整理一些常用的推理引擎并进行比较。
43 2
|
6天前
|
人工智能 算法
等不来OpenAI的Q*,华为诺亚探索LLM推理的秘密武器MindStar先来了
【7月更文挑战第13天】华为诺亚方舟实验室推出MindStar,一种增强LLM推理能力的搜索框架。MindStar通过PRM奖励模型和Beam/Levin Search策略选择最佳推理路径,提升开源模型如LLaMA-2-13B、Mistral-7B的性能,与GPT-3.5等闭源模型媲美,但成本更低。尽管推理成本高和需预训练PRM,MindStar为LLM推理研究开辟新途径。[论文链接](https://arxiv.org/pdf/2405.16265v4)
19 9
|
2天前
|
算法 API 数据中心
魔搭社区利用 NVIDIA TensorRT-LLM 加速开源大语言模型推理
魔搭社区于 2022 年 11 月初创建,首次在业界提出了 “模型即服务”( MaaS, Model as a Service)的理念。
|
4天前
LLM用于时序预测真的不行,连推理能力都没用到
【7月更文挑战第15天】LLM在时序预测上的应用遇挫:研究显示,大型语言模型在多个实验中未显优势,甚至被简单注意力层替代时效果不变或更好。预训练知识未能有效利用,处理时序依赖性不足,且在小样本学习中未见提升。[链接:](https://arxiv.org/pdf/2406.16964)**
10 2
|
5天前
|
机器学习/深度学习 自然语言处理 算法
ICML 2024:零阶优化器微调大模型,大幅降低内存
【7月更文挑战第14天】ICML 2024研究表明,零阶优化用于大模型微调能大幅降低内存需求。该论文通过避免反向传播,减少LLM(大型语言模型)微调的内存开销,提出新方法,适用于资源受限环境。虽然性能可能不及一阶优化器,但为高效NLP计算开辟了新途径。论文链接:[arxiv.org/abs/2402.11592](https://arxiv.org/abs/2402.11592)**
14 3
|
9天前
|
测试技术
谷歌DeepMind全新ToT基准:全面评估LLM时间推理能力
【7月更文挑战第10天】DeepMind的ToT基准测试了大型语言模型的时间推理能力,分为ToT-Semantic(合成数据,评估时间逻辑理解)和ToT-Arithmetic(真实数据,检查时间计算)。研究使用Claude-3-Sonnet、GPT-4和Gemini 1.5 Pro进行评估,发现模型在时间逻辑理解上表现各异,而时间计算上均较强。 Gemini 1.5 Pro在复杂问题上表现出色,而GPT-4在数学相关问题上较弱。[[1](https://arxiv.org/pdf/2406.09170)]
8 1
|
1月前
|
存储 机器学习/深度学习 自然语言处理
LLM微调方法(Efficient-Tuning)六大主流方法:思路讲解&优缺点对比[P-tuning、Lora、Prefix tuing等]
LLM微调方法(Efficient-Tuning)六大主流方法:思路讲解&优缺点对比[P-tuning、Lora、Prefix tuing等]
LLM微调方法(Efficient-Tuning)六大主流方法:思路讲解&优缺点对比[P-tuning、Lora、Prefix tuing等]
|
14天前
|
人工智能 安全 物联网
2024年6月后2周重要的大语言模型论文总结:LLM进展、微调、推理和对齐
本文总结了2024年6月后两周发表的一些最重要的大语言模型论文。这些论文涵盖了塑造下一代语言模型的各种主题,从模型优化和缩放到推理、基准测试和增强性能。
29 0