突破极限：AirLLM实现70B大模型在4GB GPU上无损推理

Introduction

大型语言模型（LLM）需要大量的GPU内存，这引出了一个关键问题：是否可以在单个GPU上运行推理？如果可以，所需的最小GPU内存是多少？一个700亿参数的LLM大小约为130GB。仅加载这个模型通常就需要两个内存为100GB的A100 GPU。在推理过程中，整个输入序列也必须加载到内存中进行复杂的"注意力"计算。这种注意力机制的内存需求与输入长度呈二次方关系，这意味着除了基础的130GB模型大小之外，还需要更多的内存。那么，什么技术可以节省如此巨大的内存，从而实现在单个4GB GPU上进行推理呢？关键的是，这些内存优化技术不涉及任何会牺牲模型性能的模型压缩方法，如量化、蒸馏或剪枝。今天，我们将深入探讨大型模型极限内存优化的关键技术。在文章最后，我们还将分享一个只需几行代码即可实现此功能的开源库！

Key Technologies for Memory Optimization

1. Layered Inference

最关键的技术是分层推理。这本质上是将计算机科学中基本的"分而治之"方法应用于模型执行。我们首先看一下大语言模型的架构。现代LLM建立在谷歌开创性论文《Attention Is All You Need》中提出的多头自注意力结构之上，即通常所说的Transformer架构。

一个典型的大型模型始于一个嵌入投影层，之后是堆叠的（例如80个）相同的Transformer层，最后是一个归一化层和一个全连接层来预测词元概率。在推理过程中，这些层按顺序执行。一层的输出作为下一层的输入。在任何给定时刻，只有一层是活跃的。因此，完全不需要同时将所有层保留在GPU内存中。

我们可以在某一层即将执行时从磁盘加载它，执行其所有计算，然后完全释放其内存。采用这种方法，任何时刻所需的GPU内存大约仅为单个Transformer层的参数大小——约为完整模型的1/80，对于一个700亿参数的模型来说大约是1.6GB。

此外，一些输出缓存也存储在GPU内存中，其中最大的是键值（KV）缓存，用于避免注意力机制中的重复计算。对一个700亿参数模型的简单计算显示，KV缓存大小约为：
2 * 输入长度 * 层数 * 头数 * 向量维度 * 4 字节
对于长度为100的输入，该缓存等于 2 * 100 * 80 * 8 * 128 * 4 ≈ 30MB 的GPU内存。根据我们的监测，整个推理过程使用的GPU显存不到4GB！

2. Single-Layer Optimization — Flash Attention

Flash Attention 可以说是现代LLM开发中最重要、最关键的优化之一。几乎所有大型语言模型都使用相同的底层计算内核，而Flash Attention代表了一项重大改进。

Flash Attention优化的核心思想并非完全新颖；我们必须提及更早的论文《Self-attention Does Not Need O(n²) Memory》。最初，自注意力需要O(n²)内存（其中n是序列长度）。该论文证明，我们不需要保留所有O(n²)的中间结果。我们可以顺序计算它们，不断更新一个运行的中间结果并丢弃其他结果，从而将内存复杂度降低到O(log n)。

Flash Attention遵循类似的原理，尽管其内存复杂度稍高，为O(n)。然而，其真正的威力在于深度优化了CUDA内存访问模式，与原始实现相比，在推理和训练上都实现了显著的加速（数倍）。正如相关图示所示，原始的自注意力计算并存储O(n²)的中间矩阵。Flash Attention将计算分解为许多小块，逐块处理，并将内存占用减少到一个块的大小。

3. Model File Sharding

原始模型文件通常被分片成多个块，每块通常约为10GB。我们的执行过程是逐层的。每层只有大约1.6GB。如果我们基于原始的10GB分片加载，那么每一层的执行都需要重新加载整个10GB文件，而只使用其中的1.6GB。这个过程浪费了大量的内存带宽，更关键的是，导致了过多的磁盘I/O。

磁盘读取速度通常是整个推理流程中最慢的瓶颈，因此我们力求尽可能地减少它。为此，我们对原始的Hugging Face模型文件进行预处理，执行分层分片。对于存储，我们利用safetensors技术（https://github.com/huggingface/safetensors）。Safetensors确保存储格式与内存布局紧密匹配，并采用内存映射加载以实现最大速度。

4. Meta Device

from accelerate import init_empty_weights
with init_empty_weights():
    my_model = ModelClass(...)

在我们的实现中，我们利用了Hugging Face Accelerate提供的元设备功能（https://huggingface.co/docs/accelerate/usage_guides/big_modeling）。元设备是一种虚拟设备，专门为运行超大型模型而设计。当你通过元设备加载模型时，实际的模型数据*不会*被读入内存；只有代码结构（"骨架"）被加载。内存使用量实际上为**0**。

你可以在执行期间动态地将模型的部分内容从元设备传输到真实设备（例如CPU或GPU）。只有在那时，数据才真正加载到内存中。使用 init_empty_weights() 允许通过元设备加载模型：

from accelerate import init_empty_weights
with init_empty_weights():
    my_model = ModelClass(...)

Open-Source Library: AirLLM

pip install airllm

from airllm import AirLLMLlama2

MAX_LENGTH = 128
# Use a Hugging Face model repo ID:
model = AirLLMLlama2("garage-bAInd/Platypus2-70B-instruct")

# Or use a local path to the model...
# model = AirLLMLlama2("/path/to/your/model")

input_text = [
    'What is the capital of the United States?',
]

input_tokens = model.tokenizer(input_text,
    return_tensors="pt",
    return_attention_mask=False,
    truncation=True,
    max_length=MAX_LENGTH,
    padding=True)

generation_output = model.generate(
    input_tokens['input_ids'].cuda(),
    max_new_tokens=20,
    use_cache=True,
    return_dict_in_generate=True)

output = model.tokenizer.decode(generation_output.sequences[0])
print(output)

我们已经将所有代码开源为 AirLLM，使您能够用几行代码实现这一功能。它可以在Anima GitHub仓库中找到：https://github.com/lyogavin/Anima/tree/main/air_llm。

使用方法非常简单。首先，安装包：

pip install airllm

然后，可以像使用标准Transformer模型一样进行分层推理：

from airllm import AirLLMLlama2
>
MAX_LENGTH = 128
# 使用Hugging Face模型仓库ID：
model = AirLLMLlama2("garage-bAInd/Platypus2-70B-instruct")
>
# 或者使用模型的本地路径...
# model = AirLLMLlama2("/path/to/your/model")
>
input_text = [
    'What is the capital of the United States?',
]
>
input_tokens = model.tokenizer(input_text,
    return_tensors="pt",
    return_attention_mask=False,
    truncation=True,
    max_length=MAX_LENGTH,
    padding=True)
>
generation_output = model.generate(
    input_tokens['input_ids'].cuda(),
    max_new_tokens=20,
    use_cache=True,
    return_dict_in_generate=True)
>
output = model.tokenizer.decode(generation_output.sequences[0])
print(output)

我们已在16GB的Nvidia T4 GPU上测试了此代码。整个推理过程使用的GPU内存不到4GB。请注意，在像T4这样的低端GPU上，推理速度可能相当慢。这种方法不太适合聊天机器人这样的交互式场景，但对于RAG（检索增强生成）或PDF分析等离线数据分析任务来说是理想的选择。目前，支持仅限于基于Llama 2架构的模型。如果您需要其他模型系列的支持，请发表评论！

Can 70B Model Training Be Done on a Single GPU?

虽然推理可以通过分层进行优化，但训练能否类似地适配到单个GPU上运行呢？关键区别在于数据需求。

推理只需要上一层的输出来执行下一个Transformer层，允许在有限的数据持久化下进行顺序的、逐层的执行。然而，训练需要保留更多的数据。训练过程涉及前向传播以计算每层和张量的输出，然后是反向传播以计算每个参数的梯度。

梯度计算需要保存所有前层前向传播的结果。因此，简单地顺序执行层并不能减少内存，因为反向传播所需的中间激活值必须保持可访问状态。需要其他技术，例如梯度检查点，来实现训练中类似的内存减少。如果您对梯度检查点如何显著降低训练内存需求感兴趣，请发表评论！