大语言模型（LLM）的参数规模不断扩大，但随之而来的是对计算资源的巨大需求。想要运行一个 70B 参数的模型，通常需要数百 GB 的显存。这无疑提高了使用门槛。今天介绍一个推理加速的库——AirLLM，它让我们可以在仅有 4GB 显存的 GPU 上运行 70B 级别的Qwen,甚至可以在 8GB 显存上运行 405B 的 Llama3.1 。这是如何实现的？让我们一起来了解一下。
AirLLM 的核心原理
AirLLM 的核心思想是基于"分而治之"的策略，通过层次推理（layered inference）来优化内存使用。
- 层次独立性：大语言模型通常由多个相同的 Transformer 层组成。在推理过程中，每一层只依赖于前一层的输出，这意味着我们可以逐层处理。
- 动态加载与释放：AirLLM 只在需要时从磁盘加载特定层，完成计算后立即释放内存。这大大减少了同时占用的 GPU 内存。
- 元设备加载：利用 Hugging Face Accelerate 提供的 meta device 功能，初始加载时只读取模型结构，不加载实际参数，从而实现零内存占用的模型加载。
- Flash Attention 优化：使用 Flash Attention 技术深度优化 CUDA 内存访问，实现多倍的速度提升。
AirLLM 使用方法
使用 AirLLM 非常简单，以下是一个基本的使用示例：
from airllm import AutoModel
MAX_LENGTH = 128
# could use hugging face model repo id:
model = AutoModel.from_pretrained("garage-bAInd/Platypus2-70B-instruct")
# or use model's local path...
#model = AutoModel.from\_pretrained("/home/ubuntu/.cache/huggingface/hub/models--garage-bAInd--Platypus2-70B-instruct/snapshots/b585e74bcaae02e52665d9ac6d23f4d0dbc81a0f")
input_text = [
'What is the capital of United States?',
#'I like',
]
input_tokens = model.tokenizer(input_text,
return_tensors="pt",
return_attention_mask=False,
truncation=True,
max_length=MAX_LENGTH,
padding=False)
generation_output = model.generate(
input_tokens['input\_ids'].cuda(),
max_new_tokens=20,
use_cache=True,
return_dict_in_generate=True)
output = model.tokenizer.decode(generation_output.sequences[0])
print(output)
这段代码演示了如何使用 AirLLM 加载一个 70B 参数的 Platypus2 模型，并进行简单的文本生成。
进一步优化
AirLLM 不仅解决了内存问题，还提供了模型压缩功能，进一步提升了推理速度：
- 块级量化模型压缩：支持 4 位和 8 位块级量化（block-wise quantization），与全局量化不同，块级量化是对模型的权重矩阵进行分块处理。每个块可以根据其重要性进行不同的量化，从而提高压缩效率和模型性能。，这一方法可将推理速度提升至 3 倍，同时保持模型精度。
- 预取机制：通过预取技术，实现模型加载与计算的重叠，提升约 10%的速度。
- 多模型支持：除了 Llama 系列，AirLLM 还支持 ChatGLM 、QWen 、Baichuan 、Mistral 和 InternLM 等多种模型。
局限性
尽管 AirLLM 在低内存环境下运行大模型方面取得了突破，但它也存在一些局限性：
- 延迟增加：由于需要频繁从磁盘读取数据，整体推理延迟会显著增加。
- I/O 瓶颈：推理速度受限于 SSD 的读取速度，这可能成为性能瓶颈。
- RAM 利用：目前 AirLLM 还不支持将模型层卸载到系统 RAM，而是直接使用 SSD 。
针对这些问题，AirLLM可以在以下方面进行优化：
- 优化数据读取策略，减少 I/O 瓶颈
- 支持将模型层卸载到系统 RAM
- 进一步提升量化技术，在保持精度的同时减小模型体积
结语
AirLLM 获得社区的追捧，很大程度上反映了推理成本仍然是广大开发者最为关注的问题。随着 LLM 逐步走向工业界，从数据中心走向端侧设备，这样的项目热度就不会降低。对于该主题感兴趣的朋友可以关注本公众号了解更多最新动态。
后台回复“进群”入群讨论。