难以置信！使用这项新技术可以在单个 4GB GPU 上运行 70B LLM 推理

大型语言模型需要大量的 GPU 内存。是否可以在单个 GPU 上运行推理？如果是这样，所需的最小 GPU 内存是多少？

70B大语言模型参数大小为130GB。仅将模型加载到 GPU 中就需要 2 个 A100 GPU，每个 GPU 具有 100GB 内存。
在推理过程中，整个输入序列还需要加载到内存中以进行复杂的“注意力”计算。这种注意力机制的内存需求与输入长度呈二次方关系。除了 130GB 型号大小之外，还需要更多内存。
那么什么技术可以节省如此多的内存并能够在单个 4GB GPU 上进行推理呢？
请注意，这里的内存优化技术不需要任何模型压缩，例如量化、蒸馏、修剪，否则会牺牲模型性能。
今天我们将讲解大型模型极限内存优化的关键技术。
文章最后我们还分享了通过几行代码实现这一点的开源库！

01分层推理

最关键的技术是分层推理。这本质上是计算机科学中基本的分而治之的方法。
我们首先看一下大语言模型的架构。如今的大型语言模型都采用了Google论文《Attention is all you need》中提出的Multi-head self-attention结构。这就是后来人们所说的Transformer结构。

大语言模型首先有一个嵌入投影层。之后是80个完全相同的变压器层。最后有一个标准化和全连接层来预测 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 注意力也许是当今大型语言模型开发中最重要、最关键的优化之一。
所有各种大型语言模型本质上都使用相同的底层代码，其中 Flash Attention 是最大的改进。
Flash注意力优化的想法并不完全新颖，我们不得不提到另一篇论文“Self-attentionDoes Not Need O(n²) Memory”。
最初，自注意力需要 O(n²) 内存（n 是序列长度）。
本文提出我们实际上不需要保留 O(n²) 中间结果。我们可以按顺序计算它们，不断更新一个中间结果并丢弃其他所有结果。这将内存复杂度降低到 O(logn)。
Flash Attention 本质上类似，内存复杂度 O(n) 稍高，但Flash Attention 深度优化了 cuda 内存访问，实现推理和训练的数倍加速。

如图所示，最初的自注意力计算并存储 O(n²) 中间结果。闪存注意力将计算分成许多小块，逐块计算并将内存减少到一个块的大小。

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 安装airllm
然后可以像普通 Transformer 模型一样进行分层推理：

from airllm import AirLLMLlama2 

MAX_LENGTH = 128 
# 可以使用拥抱脸部模型 repo 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 = [ 
        '美国的首都是什么？' , 
    ] 

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 )

输出 = model.tokenizer.decode( Generation_output.sequences[ 0 ])

打印（输出）

我们已在 16GB Nvidia T4 GPU 上测试了此代码。整个推理过程使用的GPU内存不到4GB。
请注意，像 T4 这样的低端 GPU 的推理速度会相当慢。不太适合聊天机器人这样的交互场景。更适合一些离线数据分析，如RAG、PDF分析等。
目前仅支持基于 Llam2 的模型。如果您需要其他型号的支持，请发表评论！

06 70B 训练可以在单个 GPU 上进行吗？

虽然可以通过分层来优化推理，但训练是否可以在单个 GPU 上进行类似的工作？
Inference在执行下一个transformer层时只需要上一层的输出，因此可以在有限的数据下分层执行。
训练需要更多数据。训练过程首先计算前向传播以获得每层和张量的输出。然后进行反向传播来计算每个张量的梯度。
梯度计算需要保存前面前向层的结果，因此分层执行并不会减少内存。
还有一些其他技术（例如梯度检查点）可以实现类似的效果。
如果您对梯度检查点如何显着降低训练内存需求感兴趣，请发表评论！

文章搬运地址：https://medium.com/ai-advances/unbelievable-run-70b-llm-inference-on-a-single-4gb-gpu-with-this-new-technique-93e2057c7eeb