8G显卡挑战Llama3 405B

我正在接受挑战，在只有 8GB VRAM 的 GPU 上运行 Llama 3.1 405B 模型。

Llama 405B 模型有 820GB！这是 8GB VRAM 容量的 103 倍！

显然，8GB VRAM 无法容纳它。那么我们如何让它工作呢？

1、4 位量化

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

量化后，所有浮点数将分配到 4 位的 16 个存储桶中的一个。深度神经网络中浮点数的范围从 -3.40282347E+38 到 3.40282347E+38。仅使用 16 个存储桶可以表示如此广泛的浮点数范围吗？

是的，可以。

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

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

幸运的是，深度神经网络的参数一般都服从正态分布。因此，简单的变换就可以确保理论上的均匀分布。

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

我们只需要使用一些专用的存储空间来专门记录这些异常值。这被称为异常值相关量化。

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

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

2、逐层推理

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

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

Llama 405B型号共有126层，层数增加了50％。

但是向量维度增加了一倍，多头注意力头的数量也增加了一倍，所以每层的参数数量大概是原来的四倍。

通过逐层加载和推断，最大 VRAM 使用量约为 5GB。

挑战完成！

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

3、开源项目 AirLLM

AI 行业中各种大型模型之间的差距正在迅速缩小。模型之间的差异越来越不明显。

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

我也是开源的坚定信徒，相信 AI 的未来属于开源。

本文介绍的方法已在我的开源项目 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)

原文链接：Crazy Challenge: Run Llama 405B on a 8GB VRAM GPU

BimAnt翻译整理，转载请标明出处