核能挑战：8GB显存本地跑Llama 405B！

今天是整活时间，我要挑战用8GB显存GPU跑Llama 3.1 405B模型！

Llama 405B模型有820GB! 是8GB显存的103倍！

完全放不下，怎么搞呢？太难了！

01

4bit量化

首先，我们用4-bit量化技术，把16-bit浮点数变成4-bit，节省4倍内存。

量化之后，所有的浮点数都会被分配到4个bit，也就是16个桶之一。深度神经网络的浮点数可以表达的范围从：-3.40282347E+38 to 3.40282347E+38，这么多的浮点数，都用这16个桶之一来表达，这靠谱吗？

特别靠谱。

首先最重要的是保证这些参数的数值均匀的分布到这16个桶中。

一般情况下这是绝对做不到的。不均匀就会有很大的精度损失。

好在深度神经网络的参数一般都是服从正态分布的。因此只需要做一个简单的变换就可以保证理论上绝对的均匀分布。

当然，服从统计分布，不意味着没有一些特殊情况异常值，所谓的outlier。

比如我的存款就叫正态分布，马斯克的存款就叫outlier。

那我们只需要用一些专门的存储空间，把这些outlier异常值特殊记录下来就好了。这就是所谓的outlier-dependent量化。

经过大量的实验验证，4bit量化对于大语言模型的精度几乎没有任何损失。（甚至有的情况下精度更高了！）

经过了一轮丧心病狂的4bit量化，LLama 405B的模型大小已经缩小到了230GB，距离加载到我的8GB显卡已经《不远了》。

02

分层推理

达成挑战第二个魔法是分层推理。

其实transformer的结果推理过程只需要一层一层的加载模型就可以。不需要整体将模型全部加载到内存。

Llama 405B的模型在层数上只增加了50%，到126层。

但是向量维数增加了一倍，多头注意力头数也增加了一倍，因此每层的参数量大概来到了原来的4倍。

逐层加载并推理，最多的显存占用大概只需要5GB。

挑战达成！

03

开源项目AirLLM

和某些大佬的看法不太一样，我的观察是，AI行业各个大模型的差距在快速接近。用什么模型差异已经不大。

越来越多的公司会愿意采用开源模型，并且自己部署大模型，保证可以根据业务需要灵活的调整配置。

本人也是开源坚定的拥护者，相信AI的未来一定也是属于开源的。

这个方法已经也更新到我的开源项目AirLLM， (https://github.com/lyogavin/airllm)。大家都可以使用。

pip安装：

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)