优化内存以进行大型语言模型推理和微调

GPT-4、Bloom 和 LLaMA 等大型语言模型 (LLM) 通过扩展到数十亿个参数而实现了卓越的功能。然而，由于其巨大的内存需求，部署这些大规模模型进行推理或微调具有挑战性。在本技术博客中，我们将探索在 LLM 推理和跨各种硬件设置进行微调期间估计和优化内存消耗的技术。

了解内存要求

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

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

例如，加载 175B 参数 GPT-3 模型将需要大约 350GB bfloat16 精度的 VRAM。截至目前，最大的商用 GPU（例如 NVIDIA A100 和 H100）仅提供 80GB VRAM，因此需要张量并行和模型并行技术。

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

量化推理内存

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

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: Please write a Python function to convert bytes to gigabytes.\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，这与我们估计的用于加载 bfloat31 格式的模型参数的 16GB 一致。

通过量化优化推理内存

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

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

</div>
# 8-bit quantization
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-bit quantization
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 15 (3090GB VRAM) 等消费类 GPU 上运行 24B 参数 OctoCode 模型。

但请注意，与 4 位或 bfloat8 精度相比，更激进的量化（例如 16 位）有时会导致精度下降。用户应根据其用例评估内存节省和准确性之间的权衡。

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

估计内存以进行微调

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

由于以下方面的额外内存需求，微调期间的峰值内存消耗通常比推理高 3-4 倍：

• 渐变
• 优化器状态
• 存储用于反向传播的前向传递的激活

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

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

分布式微调技术

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

1. 数据并行：经典的数据并行方法在多个 GPU 上复制整个模型，同时分割和分配训练数据批次。这会随着 GPU 数量的增加而线性减少训练时间，但不会减少每个 GPU 的峰值内存需求。
2. ZeRO 第 3 阶段：一种高级形式的数据并行性，可跨 GPU 划分模型参数、梯度和优化器状态。与传统的数据并行相比，它通过在训练的不同阶段仅在每个 GPU 上保留所需的分区数据来减少内存。
3. 张量并行：张量并行不是复制模型，而是将模型参数划分为行或列，并将它们分布在 GPU 上。每个 GPU 都在一组分区的参数、梯度和优化器状态上运行，从而节省大量内存。
4. 流水线并行：该技术将模型层划分到不同的 GPU/worker 上，每个设备执行层的子集。激活在工作人员之间传递，减少了峰值内存，但增加了通信开销。

估计这些分布式方法的内存使用量并非易事，因为参数、梯度、激活和优化器状态的分布因技术而异。此外，不同的组件（例如变压器主体和语言建模头）可能会表现出不同的内存分配行为。

法学硕士解决方案

研究人员最近提出 法学硕士，一种在跨多个 GPU 对 LLM 应用分布式微调方法时准确估计 GPU 内存消耗的解决方案。

估计 GPU 内存使用情况以微调预训练的 LLM

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

实验结果表明，LLMem 可以估计峰值 GPU 内存使用量，以便在单个 GPU 上微调 LLM，错误率高达 1.6%，优于最先进的 DNNMem 的平均错误率 42.6%。当将分布式微调方法应用于多个 GPU 上具有超过 10 亿个参数的 LLM 时，LLMem 实现了令人印象深刻的平均错误率 3.0%.