疯狂挑战：在8GB显存GPU上运行Llama 405B

我正在挑战在只有8GB显存的GPU上运行Llama 3.1 405B模型。

Llama 405B模型的大小是820GB！这相当于8GB显存容量的103倍！

它显然无法完全载入8GB显存。那么我们如何才能使其运行呢？

4比特量化

首先，我们使用4比特量化技术将16位浮点数转换为4位，从而节省四倍的内存。

量化后，所有浮点数将被分配到4比特的16个桶中的一个。深度神经网络中浮点数的范围从-3.40282347E+38到3.40282347E+38。这么大的浮点数范围能用仅仅16个桶来表示吗？

是的，可以。

最重要的是要确保这些参数均匀分布在16个桶中。

通常，这几乎是不可能实现的。不均匀的分布会导致显著的精度损失。

幸运的是，深度神经网络的参数通常遵循正态分布。因此，一个简单的转换可以确保理论上的均匀分布。

当然，遵循统计分布并不意味着没有异常值。

我们只需使用一些专门的存储空间来特殊记录这些异常值。这被称为依赖于异常值的量化。

大量实验表明，4比特量化几乎不影响大型语言模型的准确性。（在某些情况下，准确性甚至更高！）

经过一轮广泛的4比特量化后，Llama 405B模型的大小已减小到230GB，使我们“更接近”在我的8GB GPU上加载它。

逐层推理

实现这一挑战的第二个魔法是逐层推理。

事实上，Transformer的推理过程只需要逐层加载模型。没有必要一次性将整个模型加载到内存中。

Llama 405B模型有126层，层数增加了50%。

然而，向量维度翻了一番，多头注意力头的数量也翻了一番，因此每层的参数数量大约是原来的四倍。

通过逐层加载和推理，最大显存使用量约为5GB。

挑战完成！

现在我可以在我的8GB GPU上成功运行Llama 405B了！

开源项目AirLLM

AI行业中各种大型模型之间的差距正在迅速缩小。模型之间的差异变得不那么显著。

越来越多的公司愿意采用开源模型并自行部署大型模型，以确保他们可以根据业务需求灵活控制和调整其模型。

我也坚信开源，并认为AI的未来属于开源。

此方法已在我的开源项目AirLLM（https://github.com/lyogavin/airllm）中分享。

pip install airllm

你只需几行代码即可实现

from airllm import AutoModel

model = AutoModel.from_pretrained（
    "unsloth/Meta-Llama-3.1-405B-Instruct-bnb-4bit")

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=128, 
      padding=False)

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

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

print(output)

更多详情请见此处。

我们将继续关注最新最酷的AI技术，并持续分享开源工作。欢迎关注我们，敬请期待！