如何根据LLM的参数估算内存（显存）使用量

简介

基于Transformer架构的大型语言模型（LLMs）已经变得越来越普遍，例如Mistral AI团队推出的Mistral 7B模型。理解其推理、微调和训练时的内存需求对于有效部署和使用这些模型至关重要。

推理内存估算：，对于一个70亿参数（7B）的模型，预计需要的内存需求如下：

• 浮点精度：28GB
• BF16精度：14GB
• int8精度：7GB

这个估算逻辑可以等比例应用于其他版本。

训练所需内存：，保守的估计是，训练所需的内存是具有相同参数数量和类型的推理内存的四倍。例如，使用浮点精度训练一个7B模型大约需要112GB（28GB * 4）。

训练大语言模型的内存要求：

• 梯度：所需的梯度内存等于参数的数量。
• 优化器状态：优化器状态所需的内存取决于使用的优化器类型。
  • 使用AdamW优化器，需要的内存是参数数量的两倍。
  • 使用SGD优化器，需要的内存相当于参数数量。

使用LoRA/QloRA技术的内存使用情况，以LoRA为例：

• LoRA涉及在原始模型上运行推理并训练一个较小的模型，以实现与训练原始参数几乎相同的效果。

• 例如，如果需要微调大小为1024x512的参数，使用LoRA并选择Rank为8，只需要微调以下数量的参数：1024x8 + 512x8。

• 这个过程需要使用原始参数量运行一次推理（不需要梯度和优化器状态），但在计算过程中仍需要一些内存来存储数据。总内存使用量是这些需求的总和。

背景知识

在计算机内存/磁盘存储中，基本单位是字节，系统基于1024。单位包括KB、MB、GB和TB。重要的是不要将其与十进制系统混淆：1 KB = 1024 字节；1 MB = 1024 KB；1 GB = 1024 MB。

参数计数

以Mistral-7B为例，“7B”表示LLM有70亿个参数。

数据类型

对于训练好的模型，参数类型可能包括：

• 浮点（float）：32位浮点数，4字节
• 半精度（half/BF16）：16位浮点数，2字节
• int8：8位整数，1字节
• int4：4位整数，0.5字节

参数类型所需的存储量越小，性能通常越低。

推理使用内存估算

虽然其他因素也在占用内存，但推理期间使用内存的主要是参数。

• 比如，Mistral-7B-BF16模型需要的内存等于参数数量乘以类型大小：70亿参数 * 2字节 = 140亿字节。因此，140亿字节 = 14 * 1,000 * 1,000 * 1,000 / 1024 / 1024 / 1024 ≈ 13 GB（考虑1000/1024）³ ≈ 0.93。

• 注1：(1000/1024)³ ≈ 0.93

• 注2：为了估算目的，简单地将此比率视为1。因此，对于7B-BF16模型，内存需求大约是7 * 2 = 14 GB。这个估算略高于精确计算，但更实用，因为推理需要超出参数之外的额外内存。

• 假设要估算llama2-13B模型的内存需求，对应各种类型的分别是：float：13 * 4 = 52 GB，half/BF16：13 * 2 = 26 GB，int8：13 GB，int4：13 * 0.5 = 6.5 GB

训练内存使用估算

为了确保训练期间模型收敛，参数类型不能是int8或int4。通常使用float，对于稍低的性能，BF16也是一种选择。

由于反向传播、Adam优化和Transformer架构等因素，训练所需的内存通常是相同大小LLM推理所需的3到4倍。一个保守的估计是使用4作为因子进行计算。

假设要训练Qwen-7B模型，所需的内存为：

• 对于float类型：7（10亿参数）* 4（float的字节数）* 4 = 112 GB
• 对于half/BF16类型参数：7（10亿参数）* 2（每个BF16参数字节数）* 4 = 56 GB

结论

理解像Mistral-7B这样的模型的内存需求量对于优化其部署和使用至关重要。对于考虑使用云计算服务进行模型训练和推理的人来说更是如此，因为它会影响到硬件的选择和整体成本。