在单独一个4GB GPU上运行70B LLM推理的创新技术

在单独一个4GB GPU上运行70B LLM推理的创新技术

大型语言模型需要大量的GPU内存。70B大型语言模型拥有130GB的参数大小，加载此模型到GPU需要2个A100型号的GPU，各自拥有100GB内存。在推理过程中，整个输入序列也需要加载到内存中进行复杂的“注意力”计算，这使得内存需求随输入长度的增加而成平方增长。因此，我们需要新的技术来节省大量的内存并使4GB的GPU能够运行推理。

我们今天将解释节省大型模型内存的关键技术，并于文章最后分享用于实现此目标的开源库。

首要的技术就是分层推理，它是计算机科学中的基本分治法。大型语言模型首先有一个嵌入投影层，然后有80个完全相同的transformer层，最后有归一化和全连接层以预测令牌ID概率。在推理过程中，层按顺序执行，前一层的输出是下一层的输入，一次只执行一个层。因此，完全没必要把所有的层都保持在GPU内存里，我们可以在执行某一个层的时候从磁盘中加载这个层，进行所有的计算，然后完全释放内存。这样，每个层所需的GPU内存只需要一个transformer层的参数大小，也就是全模型的1/80，大约1.6GB。

第二个关键技术是单层优化——Flash Attention。它是大型语言模型发展中非常重要且关键的优化。所有的各种大型语言模型基本上都用的是同一种基础代码，Flash Attention就是这些模型中最大的提升。

另一个优化技术是模型文件Sharding，模型文件通常被分为多个部分，每个部分大约10GB。每次执行一个层，我们只加载1.6GB的数据，而不是重新加载整个10GB的文件。

我们还使用HuggingFace Accelerate提供的meta device特性，在实现时，这是一个专门为运行超大型模型设计的虚拟设备。当你通过meta device加载模型时，模型数据实际上并没有被读入，只有代码被加载，内存占用为0。你可以在执行过程中，将模型的部分从meta device动态转移到真实的设备如CPU或GPU，只有在这时它才会被实际加载到内存。

我们将所有的代码开源——名为AirLLM的库，你只需要几行代码就可以使用了。这个库可以在Anima github中找到。使用非常简单，先安装包，然后像运行一般的Transformer模型一样运行分层推理。

请注意，低端的GPU例如T4对于推理来说可能会比较慢。目前只支持基于Llam2的模型。如果你需要其他模型的支持，请留言。

推理可以通过分层进行优化，那训练是否可以在单个GPU上进行呢？推理只需要在执行下一个transformer层时，获取前一层的输出，因此，可以进行数据有限的分层执行。然而，训练需要更多的数据，它首先计算向前传播，以得到每个层和张量的输出，然后进行反向传播，计算每个张量的梯度。梯度计算需要保存前向层的结果，因此分层执行并不减少内存。但是还有其他一些技术，比如梯度检查点，可以达到类似的效果。