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

像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位整数）可以在保持最小准确性损失的同时显著减少内存使用量，适用于文本生成等推理任务。

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

</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！这使得可以在像RTX 3090（24GB VRAM）这样的消费者GPU上运行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上的峰值内存需求。
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上的超过十亿参数的LLM时，LLMem实现了令人印象深刻的平均误差率3.0%。