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

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

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

它显然无法适应8GB的显存。那么我们该如何让它运行起来呢？

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

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

是的，它可以。

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

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

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

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

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

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

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

第二个魔术技巧是逐层推理。

实际上，变压器的推理过程只需要一层层地加载模型。不需要一次性将整个模型加载到内存中。

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

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

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

挑战完成！

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

人工智能行业的各种大型模型之间的差距正在迅速缩小。模型之间的差异变得不那么显著了。

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

我也非常支持开源，并认为 人工智能的未来属于开源。

这种方法已经在我的开源项目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)

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