为大型语言模型推理和微调优化内存

像GPT-4、Bloom和LLaMA这样的大型语言模型（LLM）通过扩大到数十亿参数，实现了令人惊叹的能力。然而，部署这些庞大的模型进行推理或微调由于其巨大的内存需求而具有挑战性。在本技术博客中，我们将探索估计和优化LLM推理和微调过程中内存消耗的技术，涵盖各种硬件设置。

理解内存需求

加载LLM所需的内存主要由参数数量和用于存储参数的数字精度决定。一个简单的经验法则是：

• 加载具有X十亿参数的模型大约需要 4X GB 的VRAM，在 32位 浮点精度下
• 加载具有X十亿参数的模型大约需要 2X GB 的VRAM，在 16位 bfloat16/float16精度下

例如，加载175B参数的GPT-3模型大约需要350GB的VRAM，在bfloat16精度下。截至今日，像NVIDIA A100和H100这样的最大商用GPU仅提供80GB的VRAM，需要使用张量并行和模型并行技术。

在推理过程中，内存占用主要由模型参数和产生的临时激活张量支配。推理过程中峰值内存使用的高级估计是加载模型参数所需的内存与激活内存的总和。

量化推理内存

让我们使用OctoCode模型来量化推理内存需求，该模型具有大约15亿参数，采用bfloat16格式（~ 31GB）。我们将使用 Transformers库 加载模型并生成文本：

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder",
torch_dtype=torch.bfloat16,
device_map="auto",
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: 请编写一个Python函数将字节转换为千兆字节。\n\nAnswer:"
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]

def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024

bytes_to_gigabytes(torch.cuda.max_memory_allocated())

输出：

峰值GPU内存使用量约为29GB，与我们对bfloat16格式中加载模型参数的估计31GB一致。

使用量化优化推理内存

虽然bfloat16是训练LLM的常用精度，但研究人员发现，将模型权重量化为较低精度数据类型（如8位整数（int8）或4位整数）可以显著减少内存使用量，同时在文本生成等推理任务中最小化准确率损失。

让我们看看8位和4位量化对OctoCode模型的内存节省：

</div>
# 8位量化
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_8bit=True,
pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())</pre>

输出：

# 4位量化
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=True,
low_cpu_mem_usage=True, pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())

输出：

使用8位量化，内存需求从31GB降低到15GB，而4位量化进一步降低到9.5GB！这允许在具有24GB VRAM的消费者GPU（如RTX 3090）上运行15B参数的OctoCode模型。

但是，请注意，更激进的量化（如4位）有时会导致与8位或bfloat16精度相比的准确率下降。用户应根据自己的使用场景评估内存节省和准确率之间的权衡。

量化是一种强大的技术，可以通过大幅降低内存占用，使LLM能够在资源受限的环境中部署，例如云实例、边缘设备，甚至移动电话。

估计微调内存

虽然量化主要用于高效推理，但像张量并行和模型并行这样的技术对于管理大型语言模型训练或 微调 过程中的内存需求至关重要。

微调过程中的峰值内存消耗通常比推理高3-4倍，这是由于以下额外内存需求：

• 梯度
• 优化器状态
• 用于反向传播的正向传播激活

一个保守的估计是，使用X十亿参数的LLM微调大约需要 4 * (2X) = 8X GB 的VRAM，在bfloat16精度下。

例如，微调7B参数的LLaMA模型大约需要 7 * 8 = 56GB 的VRAM，每个GPU在bfloat16精度下。这超过了当前GPU的内存容量，需要分布式微调技术。

分布式微调技术

已经提出了几种分布式微调方法，以克服大型模型的GPU内存限制：

1. 数据并行: 数据并行的经典方法会在多个GPU上复制整个模型，同时分割和分配训练数据批次。这会将训练时间线性地减少到GPU的数量，但不会减少每个GPU上的峰值内存需求。
2. ZeRO Stage 3: 数据并行的高级形式，它将模型参数、梯度和优化器状态分区到多个GPU上。它通过仅在训练的不同阶段在每个GPU上保留所需的分区数据来减少与经典数据并行相比的内存使用。
3. 张量并行: 与复制模型不同，张量并行将模型参数划分为行或列，并将其分布到多个GPU上。每个GPU操作一组分区的参数、梯度和优化器状态，从而实现了显著的内存节省。
4. 管道并行: 该技术将模型层划分到不同的GPU/工作器上，每个设备执行一组层的子集。激活在工作器之间传递，减少峰值内存，但增加了通信开销。

估计这些分布式方法的内存使用量很棘手，因为参数、梯度、激活和优化器状态的分布在技术之间有所不同。此外，像变压器主体和语言建模头这样的不同组件可能会表现出不同的内存分配行为。

LLMem解决方案

研究人员最近提出了 LLMem，一种可以准确估计在多个GPU上应用分布式微调方法时的大型语言模型的GPU内存消耗的解决方案。

估计预训练LLM的GPU内存使用

LLMem考虑了诸如在计算之前重新组合参数（ZeRO Stage 3）、反向传播中的输出收集（张量并行）以及变压器主体和语言建模头的不同内存分配策略等因素。

实验结果表明，LLMem可以估计单个GPU上微调LLM的峰值GPU内存使用，误差率最高可达1.6%，优于最先进的DNNMem的平均误差率42.6%。当在多个GPU上应用分布式微调方法时，LLMem在超过十亿参数的LLM上实现了平均误差率3.0%的成绩。