4GB GPU 运行 70B 大模型的新技术

4GB GPU 运行 70B 大模型的新技术 | 前端进阶之旅

Full translation

完整中文译文

大型语言模型需要海量的 GPU 显存。能否在单块 GPU 上运行推理？如果可以，最少需要多少 GPU 显存？

70B 大型语言模型的参数规模为 130GB。仅仅加载模型到 GPU 就需要 2 块 A100 GPU，每块都有 100GB 显存。

推理过程中，整个输入序列也需要加载到显存中，以进行复杂的 attention 计算。这个 attention 机制的显存需求与输入长度成二次方关系增长。除了 130GB 的模型大小之外，还需要大量额外的显存。

那么有什么技术能节省这么多显存，使得在单块 4GB GPU 上进行推理成为可能呢？

需要注意的是，这里的显存优化技术不需要任何模型压缩方式，如量化、蒸馏、剪枝等会损害模型性能的方法。

今天我们将讲解大模型显存优化的关键技术。

文章最后我们也分享了开源库，仅需几行代码即可实现！

Layer-wise 推理

最关键的技术是 layer-wise 推理。这本质上是计算机科学中基本的分而治之思想。

首先来看一下大型语言模型的架构。如今的大型语言模型都采用了谷歌论文《Attention is all you need》中提出的 Multi-head self-attention 结构。这就是人们后来称之为 Transformer 结构的东西。

大型语言模型首先有一个 embedding 投影层。之后有 80 个完全相同的 transformer 层。最后是一个归一化层和全连接层来预测 token ID 的概率。

推理过程中，各层是按顺序执行的。前一层的输出是下一层的输入。同一时刻只有一层在执行。

因此，完全没有必要把所有层都保留在 GPU 显存中。我们可以在执行某层时从磁盘加载该层，完成所有计算，然后在执行完后彻底释放显存。

这样，每层所需的 GPU 显存只是一个 transformer 层的参数大小，即完整模型的 1/80，约 1.6GB。

此外，还有一些输出缓存也存储在 GPU 显存中，其中最大的是 KV cache，用于避免重复计算。

简单计算一下，对于 70B 模型，这个 KV cache 的大小约为：

2 * input_length * num_layers * num_heads * vector_dim * 4

输入长度为 100 时，这个 cache = 2 * 100 * 80 * 8 * 128 * 4 = 30MB GPU 显存。

根据我们的监测，整个推理过程所用 GPU 显存少于 4GB！

单层优化 — Flash Attention

Flash Attention 可能是当今大型语言模型发展中最重要、最关键的优化之一。

各种大型语言模型本质上使用的是相同的底层代码，Flash Attention 是最大的改进。

Flash Attention 优化的思想本身并非全新的，我们必须提到另一篇论文《Self-attention Does Not Need O(n²) Memory》。

原始的 self-attention 需要 O(n²) 的显存（n 是序列长度）。

这篇论文提出我们其实不需要保留 O(n²) 的中间结果。我们可以顺序计算它们，不断更新一个中间结果并丢弃其他所有结果。这将显存复杂度降低到 O(logn)。

Flash Attention 本质上是类似的，显存复杂度略高为 O(n)，但 Flash Attention 深度优化了 CUDA 显存访问，在推理和训练中都能实现多倍的速度提升。

如图所示，原始的 self-attention 计算并存储 O(n²) 个中间结果。Flash Attention 将计算分成许多小块，逐块计算，将显存减少到一个块的大小。

模型分片优化

原始的模型文件通常被分片成多个块，一般每个约 10GB。

我们按层执行。每层只有 1.6GB。如果基于原始的 10GB 分片加载，每次执行一层都需要重新加载整个 10GB 文件，但只使用其中 1.6GB。

这个过程浪费了大量显存用于加载，且磁盘读取是整个推理过程中最慢的瓶颈，所以我们希望尽可能地最小化它。

因此，我们首先预处理原始的 HuggingFace 模型文件，按层进行分片。

存储我们使用 safetensor 技术（https://github.com/huggingface/safetensors）。

Safetensor 确保存储格式和内存中的格式紧密匹配，并使用内存映射进行加载以最大化速度。

在实现中，我们使用了 HuggingFace Accelerate 提供的 meta device 特性（https://huggingface.co/docs/accelerate/usage_guides/big_modeling）。

Meta device 是一个为运行超大模型而专门设计的虚拟设备。通过 meta device 加载模型时，模型数据并不会被真正读入，只有代码被加载。显存使用为 0。

在执行过程中，你可以动态地将模型的部分从 meta device 传输到 CPU 或 GPU 等真实设备。只有到这时才会真正加载到显存中。

使用 init_empty_weights() 可以通过 meta device 加载模型。

from accelerate import init_empty_weights

with init_empty_weights():
    my_model = ModelClass(...)

我们开源了所有代码 — AirLLM。仅需几行代码即可实现这一功能。

它可以在 Anima github 找到：https://github.com/lyogavin/Anima/tree/main/air_llm。

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

pip install airllm

然后可以像正常的 Transformer 模型一样进行分层推理：

from airllm import AirLLM

Llama2

MAX_LENGTH = 128

# 可以使用 hugging face 模型仓库 id：
model = AirLLMLlama2("garage-bAInd/Platypus2-70B-instruct")

# 或使用模型的本地路径...
# 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 分析等。

目前仅支持基于 Llama2 的模型。如果你需要支持其他模型，请留言！

70B 训练能否在单块 GPU 上进行？

推理可以通过分层来优化，那么训练是否也能在单块 GPU 上类似地进行呢？

推理时执行下一个 transformer 层时只需要前一层的输出，所以有限数据的分层执行成为可能。

训练需要更多的数据。训练过程首先计算前向传播以获得每一层和每个张量的输出。然后进行反向传播来计算每个张量的梯度。

梯度计算需要保存之前的前向层的结果，所以分层执行不能减少显存。

有一些其他技术，如梯度检查点（gradient checkpointing），可以实现类似的效果。

如果你对梯度检查点如何显著减少训练显存需求感兴趣，请留言！

我们的代码参考了很多 SIMJEG 在 Kaggle 上的实现：https://www.kaggle.com/code/simjeg/platypus2-70b-with-wikipedia-rag/notebook。向 Kaggle 社区为他们的贡献致敬！

我们将继续开源最新、最有效的 AI 新方法和进展，为开源社区做出贡献。请关注我们。