难以置信！使用这项新技术在单张 4GB GPU 上运行 70B LLM 推理
大型语言模型需要大量的 GPU 内存。是否可以在单个 GPU 上运行推理？如果可以，所需的最小 GPU 内存是多少？
70B 大型语言模型的参数大小为 130GB。仅将模型加载到 GPU 中就需要两张各 100GB 内存的 A100 GPU。
在推理过程中，整个输入序列也需要加载到内存中进行复杂的“注意力”计算。这种注意力机制的内存需求与输入长度呈二次方增长。除了 130GB 的模型大小之外，还需要更多的内存。
那么，有哪些技术可以节省如此多的内存，并实现在单个 4GB GPU 上进行推理呢？
请注意，这里的内存优化技术不需要任何模型压缩，如量化、蒸馏、剪枝等会牺牲模型性能的方法。
今天，我们将解释大型模型极端内存优化的关键技术。
在文章的最后，我们还分享了只需几行代码即可实现此功能的开源库！
01
逐层推理
最关键的技术是逐层推理。这本质上是计算机科学中基本的分治法。
我们首先看看大型语言模型的架构。今天的语言模型都采用了 Google 论文《Attention is all you need》中提出的多头自注意力结构。这就是人们后来所说的 Transformer 结构。
大型语言模型首先有一个嵌入投影层。之后有 80 个完全相同的 Transformer 层。最后是一个归一化和全连接层，用于预测 token ID 概率。
在推理过程中，层是顺序执行的。前一层的输出是下一层的输入。一次只有一个层执行。
因此，完全没有必要将所有层都保留在 GPU 内存中。我们可以在执行某个层时，从磁盘加载该层所需的任何数据，进行所有计算，然后完全释放内存。
这样，每个层所需的 GPU 内存仅约为一个 Transformer 层的参数大小，即整个模型的 1/80，大约 1.6GB。
此外，一些输出缓存也存储在 GPU 内存中，其中最大的是 KV 缓存，以避免重复计算。
简单计算一下，对于 70B 模型，KV 缓存大小约为
2 * input_length * num_layers * num_heads * vector_dim * 4
输入长度为 100 时，此缓存 = 2 * 100 * 80 * 8 * 128 * 4 = 30MB GPU 内存。
根据我们的监控，整个推理过程使用的 GPU 内存少于 4GB！
02
单层优化 — Flash Attention
Flash Attention 可能是当今大型语言模型开发中最重要的关键优化之一。
所有各种大型语言模型基本上都使用相同的底层代码，而 Flash Attention 是最大的改进。
Flash Attention 优化的思想并非完全新颖，我们不得不提及另一篇论文“自注意力不需要 O(n²) 内存”。
最初的自注意力需要 O(n²) 内存（n 是序列长度）。
这篇论文提出我们实际上不需要保留 O(n²) 的中间结果。我们可以顺序计算它们，持续更新一个中间结果并丢弃所有其他结果。这将内存复杂度降低到 O(logn)。
Flash Attention 本质上相似，内存复杂度略高，为 O(n)，但 Flash Attention 深度优化了 CUDA 内存访问，以实现推理和训练的数倍加速。
如图所示，最初的自注意力计算并存储 O(n²) 的中间结果。Flash Attention 将计算分成许多小块，逐块计算并将内存减少到一块的大小。
03
模型文件分片
原始模型文件通常被分片为多个块，通常每个块 10GB。
我们的执行是逐层进行的。每层只有 1.6GB。如果根据原始的 10GB 分片加载，每次层执行都需要重新加载整个 10GB 文件，但只使用了 1.6GB。
这个过程浪费了大量的内存用于加载和磁盘读取。磁盘读取速度实际上是整个推理过程中最慢的瓶颈，所以我们希望尽可能地减少它。
因此，我们首先预处理原始的 HuggingFace 模型文件，并按层进行分片。
对于存储，我们使用 safetensor 技术（https://github.com/huggingface/safetensors）。
Safetensor 确保存储格式和内存中格式紧密匹配，并使用内存映射进行加载以最大限度地提高速度。
04
元设备
在实现中，我们使用了 HuggingFace 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(...)
05
开源库
我们开源了所有代码 — AirLLM。它让你只需几行代码即可实现此功能。
它可以在 Anima github 上找到：**https://github.com/lyogavin/Anima/tree/main/air_llm.**
使用方法非常简单。首先安装软件包
pip install airllm
然后可以像普通的 Transformer 模型一样执行分层推理
from airllm import AirLLMLlama2
MAX_LENGTH = 128
# could use hugging face model repo id:
model = AirLLMLlama2("garage-bAInd/Platypus2-70B-instruct")
# or use model's local path...
#model = AirLLMLlama2("/home/ubuntu/.cache/huggingface/hub/models--garage-bAInd--Platypus2-70B-instruct/snapshots/b585e74bcaae02e52665d9ac6d23f4d0dbc81a0f")
input_text = [
'What is the capital of 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 分析等。
目前仅支持基于 Llam2 的模型。如果您需要支持其他模型，请留言！
06
70B 模型训练可以在单个 GPU 上完成吗？
虽然推理可以通过分层进行优化，但训练是否也可以在单个 GPU 上进行类似的操作？
推理在执行下一个 Transformer 层时只需要前一层的输出，因此可以使用有限数据进行分层执行。
训练需要更多的数据。训练过程首先计算前向传播以获得每一层和张量的输出。然后进行反向传播以计算每个张量的梯度。
梯度计算需要保存之前前向层的结果，因此分层执行并不能减少内存。
还有一些其他技术，如梯度检查点（gradient checkpointing），可以实现类似的效果。
如果您对梯度检查点如何显著降低训练内存需求感兴趣，请留言！
07
我们的代码大量参考了 SIMJEG 在 Kaggle 上的实现：https://www.kaggle.com/code/simjeg/platypus2-70b-with-wikipedia-rag/notebook。感谢 Kaggle 社区的杰出贡献！
我们将继续开源人工智能领域最新最有效的新方法和进展，为开源社区做出贡献。请关注我们。
