大语言模型参数规模与显存大小的关系估算方法
开源大语言模型的发展非常迅速,其强大的能力也吸引了很多人的尝试与体验。尽管预训练大语言模型的使用并不复杂,但是,因为其对GPU资源的消耗很大,导致很多人并不能很好地运行加载模型,也做了很多浪费时间的工作。其中一个比较的的问题就是很多人并不知道自己的显卡支持多大参数规模的模型运行。本文将针对这个问题做一个非常简单的介绍和估算。
大模型消耗的显存简介
在详细说明大模型需要消耗的显存大小之前我们需要先明确几个概念。
一个就是大模型在不同阶段对显存的消耗是不同的。但是大致可以分为三个阶段或者说三个场景。即大模型预训练阶段、大模型微调阶段和大模型推理阶段。
在预训练阶段,大模型通常选择较大规模的数据集获取泛化能力,因此需要较大的批次等来保证模型的训练强大。而模型的权重也是从头开始计算,因此通常也会选择高精度(如32位浮点数)进行训练。需要消耗大量的GPU显存资源。
在微调阶段,通常会冻结大部分参数,只训练小部分参数。同时,也会选择非常多的优化技术和较少的高质量数据集来提高微调效果,此时,由于模型已经在预训练阶段进行了大量的训练,微调时的数值误差对模型的影响通常较小。也常常选择16位精度训练。因此通常比预训练阶段消耗更低的显存资源。
在推理阶段,通常只是将一个输入数据经过模型的前向计算得到结果即可,因此需要最少的显存即可运行。
下面是针对大模型在不同阶段消耗的显存的详细估算。
大模型预训练阶段的显存消耗
大模型在预训练阶段需要加载完整的模型结构和大量预训练数据,所以显存需求很大,通常需要几十GB到上百GB。常见的预训练数据集如ImageNet图片数据集、维基百科文本corpus等都比较大。
预训练大模型的显存消耗估算涉及多个因素。以下是一个简化的步骤来估算显存消耗:
模型参数
每个参数通常是一个32位浮点数(float32),需要4字节的存储空间。
如果模型有P个参数,那么参数的总大小为:P×4字节
中间激活
对于每一层,都会有一些中间的激活值。这些激活值的数量通常与模型的宽度和输入数据的大小有关。
假设每层的激活需要A个浮点数,那么激活的总大小为:A×4字节
梯度
每个模型参数在反向传播时都会有一个对应的梯度。
梯度的总大小与模型参数的总大小相同:P×4字节
优化器状态
一些优化器(如Adam)为每个参数存储额外的状态信息。例如,Adam存储每个参数的动量和梯度的平方值。
对于Adam,状态的总大小为:2×P×4字节
批次大小
批次大小会影响中间激活的数量。假设批量大小为B,那么激活的总大小为:B×A×4字节
其他因素
正则化、Dropout、Batch Normalization等可能需要额外的显存。
一些特殊的操作或层(如特殊的卷积)可能有额外的显存需求。
综上所述,预训练大模型的显存消耗估算可以通过以下公式得到:
总显存=(P+B×A+P+2×P)×4字节
这只是一个简化的估算。实际的显存消耗可能会受到其他因素的影响,如模型的具体结构、使用的库和框架、显卡的架构等。为了更准确地估算显存消耗,建议使用专门的工具或在实际环境中进行测试。
进一步简化一下只考虑模型的参数和批次大小,那么大约是4倍批次大小乘以参数的关系。
总结一下就是,大模型预训练阶段需要高精度数值来保证训练过程的准确性,需要大规模数据让模型获得泛化和涌现能力,也会一次使用较大批次加速训练,因此需要大量显存消耗。
大模型微调阶段的显存消耗估算
第二个场景就是大模型微调阶段需要的显存。大模型微调是指在一个预训练大模型的基础上,对模型进行细微的调整,使其适应特定的任务。这通常涉及在新的、特定任务的数据集上进行额外的训练。
在深度学习中,预训练和微调是两个常见的阶段。预训练通常涉及在大型数据集上从头开始训练模型,而微调则是在预训练模型的基础上,使用特定的任务数据进行进一步的训练。通常情况下,微调阶段的显存消耗通常远低于预训练阶段:
- 数据集大小:预训练通常使用的是大型的数据集,如Wikipedia、Common Crawl等,这些数据集包含数十亿的单词。而微调时使用的数据集通常针对特定任务,相对较小。处理更小的数据集时,显存的需求通常会减少。
- 批次大小:由于预训练数据集的大小,通常会使用更大的批次来加速训练。大的批次会消耗更多的显存。而在微调时,由于数据集较小,可以使用较小的批次,从而减少显存使用。
- 模型结构:虽然预训练和微调通常使用相同的模型结构,但在微调时,有时会冻结模型的某些部分(例如,不训练预训练模型的某些层),这可以减少显存的使用。
- 累积梯度:在预训练阶段,由于使用大型数据集和大批次,可能需要累积多个批次的梯度来进行一次参数更新,这会增加显存的使用。而在微调阶段,由于数据集和批次大小都较小,这种需求会减少。
- 正则化和技巧:预训练阶段可能会使用更多的正则化技巧,如Dropout、Layer Normalization等,这些技巧可能会增加显存的使用。而在微调阶段,这些技巧可能会被减少或调整。
- 优化器:某些优化器,如Adam,会为每个参数保存额外的状态信息,从而增加显存使用。在预训练阶段,可能会使用这些优化器,而在微调阶段,可能会选择不同的优化策略。
总之,虽然预训练和微调都是深度学习训练的重要阶段,但由于它们的目标、数据集和策略的差异,微调阶段的显存消耗通常会远低于预训练阶段。这里也没有一个统一的公式估算,主要还是依赖于微调选择的技术和参数。尤其是如果我们冻结了很多参数的话,微调所需要的显存资源也会很低。
简单来说,大模型微调阶段不需要大幅度更新模型权重,因此可以以较低精度来实现,也只需要在较小规模的高质量数据集上进行小批次的输入,因此比大模型预训练阶段消耗更低的GPU显存资源。
大模型推理阶段的显存消耗
大模型在推理阶段的显存估计相对简单,主要考虑以下几点:
- 模型参数量:这个与训练阶段相同,决定了模型本身的显存占用量。
- 单次输入数据:推理只需要加载单次输入,而不是像训练时的 batch 数据。所以输入数据量很小。
- 激活复用:推理阶段可以重用激活数据,不需要像训练时存储多个步骤的激活,大幅减少激活显存。
- 无需保存梯度:推理阶段不需要计算和保存梯度信息,减少消耗。
- 优化器状态:推理阶段不需要优化器状态,减少消耗。
所以一个粗略的估计方式是:参数数量×2字节数。例如,一个7B大小的模型,推理阶段通常需要的显存约7*2=14GB。原因是推理阶段通常使用16位精度,一个参数占用2个字节,7b的参数规模是70亿,二者相乘再换算到GB,也就是14GB左右了。
类似的,13B通常需要26GB显存,65B通常需要130GB显存左右才能推理。
不过也有一个需要考虑的是量化,在某些时候,大模型可以继续将16位精度量化成INT8或者INT4的精度,那么量化后的大模型在推理阶段所需要的显存大小可以进一步降低。例如,7B模型的INT4量化可能只需要3.5GB就行。如果模型是混合精度量化,则介于全精度和完全量化之间了。
相比训练阶段,推理阶段显存需求可以降低10倍以上。这使得很大的模型也可以部署到资源受限的环境中进行推理。
总之,推理阶段显存需求远低于训练阶段,原因在于推理只需要前向传播而不需要计算图的反向传播过程。这大大减少了激活和梯度的存储需求。