在单个 4GB GPU 上运行 Meta-Llama-3-8B-Instruct – 使用 airllm 和分层推理

0. 原理

分层推理本质上是“分而治之”的方法

这没有使用量化、蒸馏、修剪或其他模型压缩技术

大型语言模型之所以体积大，占用大量内存，主要是因为它们的结构包含许多“层”。 从LLM嵌入投影层开始，然后是许多相同的变压器层。 8B 型号有32 层。但是在推理过程中，每一层都是独立的，只依赖于前一层的输出。 因此，在运行层后，可以释放其内存，仅保留层的输出。基于这一概念，AirLLM实现了分层推理。 如何在基于LLMTransformer的推理过程中，各层是按顺序执行的。上一层的输出是下一层的输入。一次只执行一个图层。 因此，完全没有必要将所有层都保留在 GPU 内存中。我们可以在执行该层时从磁盘加载所需的任何层，执行所有计算，然后完全释放内存。 这样一来，每层所需的 GPU 内存仅为一个 transformer 层的参数大小，即完整型号的 1/32，约为 417MB。

然后使用闪存深度优化cuda内存访问，实现多倍加速按层分片模型文件。

使用 HuggingFace Accelerate 提供的元设备功能。当您通过元设备加载模型时，实际上不会读取模型数据，只会加载代码。内存使用率为 0。

提供使用“压缩”参数进行量化的选项 “压缩”：支持的选项：4 位、8 位用于 4 位或 8 位块级量化

引用连接：https://github.com/lyogavin/Anima

1. 推理 Meta-Llama-3-8B-Instruct

1.1 查看 config.json

1. 模型类型 (model_type): llama
2. 隐藏层激活函数 (hidden_act): silu
3. 隐藏层大小 (hidden_size): 4096
4. 中间层大小 (intermediate_size): 14336
5. 最大位置嵌入数 (max_position_embeddings): 8192
6. 注意力头数 (num_attention_heads): 64
7. 隐藏层数 (num_hidden_layers): 32
8. 键值头数 (num_key_value_heads): 8
9. 词汇表大小 (vocab_size): 128256
10. 初始化范围 (initializer_range): 0.02
11. 注意力丢失 (attention_dropout): 0.0
12. 标准化余弦阈值 (rms_norm_eps): 1e-05
13. 绑定词嵌入 (tie_word_embeddings): false
14. 张量数据类型 (torch_dtype): bfloat16
15. Transformers版本 (transformers_version): 4.40.0.dev0
16. 是否使用缓存 (use_cache): true
17. 起始标记ID (bos_token_id): 128000
18. 结束标记ID (eos_token_id): 128001
19. Rope theta值 (rope_theta): 500000.0

从配置文件里的参数 num_hidden_layers 为 80，标识模型权重有80层

1.2 尝试加载

参照 airllm 的例子，加载 Meta-Llama-3-8B-Instruct

运行结果

从处理过程来看，其实是把每一层保存到 splitted_mode 目录下了。

我们显示一下文件

第一次要保存这些文件，所以非常慢，大概需要2个多小时,每层大小约为417MB。

那第二次则直接判断这些文件存在不存在，自然就快了。

和前面一样，输入token数为8， 输总token 为28， 所以28-8 = 20， 总共20个回合，每个回合2分多钟

1.3 尝试4bit加载

只需要增加 compression=’4bit’就可以，8bit 标识8位压缩

运行后会重新产生一个4bit 的目录，文件产生大约时间为10分钟左右

每层的文件大小大概为118MB，是前面没压缩的面的1/4。

总结消耗时间

1.4 4bit加载后，保留在内存

原来的代码是每次运行都行从文件里面加载，现在改为，加载后保留内存，先克隆原来的项目（版本号是2.8.3)

1.4.1 修改版本号

修改 setup.py 文件里面的版本号

1.4.2 修改 load_layer_to_cpu 函数

修改 airllm/airllm_base.py 文件两处，首先添加一个全局变量 loaded_layers，用来保存加载后的层

修改后的 load_layer_to_cpu 函数：

1.4.3 安装到你的环境

1.4.5 进行再次测试

可以看到第一次推理的时间没有变：1分59秒，但后面19次的推理时间都是7秒左右，整体时间缩小了10倍。