AirLLM：让4GB显存跑起70B大模型

2024年的某天，一位名叫Gavin Li的开发者在Kaggle竞赛中遇到了一个头疼的问题： 手里只有一块消费级GPU（显存4GB），却想跑一个70B参数的大语言模型。传统方案需要8块A100才能搞定，要么量化掉精度，要么压缩模型，效果都不理想。他没有妥协，而是花了几个周末写了一个"分层加载"的hack——把模型按层拆开，推理时只把当前计算层加载到显存，其余层留在内存里，按需调度。就这样，70B模型在单卡4GB显存上跑起来了。

这个hack后来演变成了AirLLM，一个开源推理加速库，如今在GitHub上已收获超过18000颗星，成为大模型推理优化领域最受关注的项目之一。

什么是AirLLM？

AirLLM的核心思路非常优雅：不再要求把整个70B模型塞进显存，而是按层（layer）拆分加载。传统推理时，所有模型权重一次性加载到GPU——70B的fp16模型需要约140GB显存，根本不可能。而AirLLM的策略是：把模型分成若干层（Transformer层），每次只把当前正在计算的一层或少数几层放到GPU上，其他层保留在CPU内存中。当一层计算完成后，立即卸载，再加载下一层。

这样做显存占用从O(模型总大小)降到了O(单层大小)。以70B模型为例，每层约400MB，单层推理只需不到4GB显存——普通游戏显卡就能跑。

更重要的是，这个方案不需要量化、不需要蒸馏、不需要剪枝，模型精度完全保留。它本质上是一个推理调度优化，而非模型压缩。

核心技术与实现细节

分层加载（Layer-wise Loading）

AirLLM的核心实现在airllm_base.py中，通过自定义nn.Module重写了transformer的前向传播。以Llama2为例（airllm_llama2.py），每次forward时：

1. 输入token经过embedding层
2. 遍历每一层 TransformerLayer：
   • 将该层加载到GPU（已在CPU则跳过）
   • 执行attention + MLP计算
   • 卸载该层回CPU内存
3. 最终通过lm_head输出token

关键代码模式：

for layer in self.model.layers:
    layer = layer.cuda()
    output = layer(output, ...).cuda()
    layer = layer.cpu()
    torch.cuda.empty_cache()

块级量化压缩（Block-wise Quantization）

AirLLM v2.0引入了块级量化，在分层加载基础上进一步提速3倍。其原理来自QuaIA论文（arXiv:2212.09720）：将模型权重按块（block）为单位进行4bit或8bit量化，只量化权重不量化激活，从而在几乎不损失精度的情况下将模型文件大小减半，加载速度同步提升。

AutoModel自动适配

v2.6版本新增了AutoModel，自动识别模型类型（Llama2/ChatGLM/QWen/Baichuan/Mistral/InternLM等），用户无需手动选择模型类：

from airllm import AutoModel
model = AutoModel.from_pretrained("THUDM/chatglm3-6b-base")

支持的模型覆盖了HuggingFace Open LLM榜单Top 10的主流开源模型，包括Llama3.1 405B（8GB显存即可运行）。

Apple Silicon支持

AirLLM也支持MacOS，通过MLX框架（苹果自研神经网络加速库）实现在M系列芯片上运行70B模型。项目依赖中同时包含PyTorch和MLX，根据平台自动切换后端。

性能表现

图：v2.0块级量化后推理速度提升约3倍

AirLLM的分层加载机制虽然理论上会因频繁的CPU-GPU数据传输拖慢速度（每层计算完都要传输下一层），但由于PCIe带宽远高于计算密度，实际上瓶颈在GPU计算而非传输——层间传输可以与下一层预取并行执行。实测v2.5版本加入prefetching优化后，速度再提升10%。

对于极端场景（4GB显存），AirLLM将原本不可能的推理任务变成了现实——虽然速度比多卡A100慢，但"能跑"本身就是价值。

上手体验

安装极其简单：

pip install airllm

推理代码和标准HuggingFace transformers几乎一样：

from airllm import AutoModel
model = AutoModel.from_pretrained("garage-bAInd/Platypus2-70B-instruct", compression='4bit')
input_tokens = model.tokenizer(["What is AI?"], return_tensors="pt")
output = model.generate(input_tokens['input_ids'].cuda(), max_new_tokens=50)
print(model.tokenizer.decode(output.sequences[0]))

注意事项：首次运行会将模型按层拆分保存到磁盘（默认HuggingFace缓存目录），需要约2倍模型大小的磁盘空间。MacOS仅支持Apple Silicon。

局限与争议

AirLLM并非完美方案，以下几点值得注意：

速度牺牲：分层加载本质上是用时间换空间。单层推理速度虽然和全量GPU推理相同，但频繁的层切换引入额外开销。对于延迟敏感的生产环境（比如在线问答），这可能不适用。

内存带宽瓶颈：虽然prefetching可以重叠传输和计算，但CPU-GPU传输仍然受到PCIe带宽限制。70B模型每层400MB，100层意味着总共40GB的数据传输——这在高吞吐场景下会成为明显瓶颈。

不支持batch推理优化：当前实现主要针对单序列推理，batch并行优化有限。对于需要同时处理多条请求的场景，性能收益不明显。

依赖PyTorch开发版：requirements.txt中明确依赖transformers/peft/accelerate的GitHub开发版，存在版本兼容风险，生产环境部署需要锁定具体commit。

行业意义

AirLLM的出现反映了一个重要趋势：大模型推理正在从"大力出奇迹"走向"精细化工程"。当模型越来越大（GPT-4传闻1.8T参数），硬件增长的速度远远跟不上模型膨胀。分层加载、KV Cache优化、推测解码（Speculative Decoding）等技术，正在把"只有大厂能跑大模型"的大门撬开。

Gavin Li在README中引用了自己的BibTex引用格式，可见他对自己的工作有清晰的学术定位。从Kaggle竞赛的临时hack，到PyPI月下载量稳定在前列的开源项目，AirLLM用实际影响力证明：工程创新有时候比算法创新更有普惠价值——它让每一个有游戏显卡的开发者，都能玩转70B大模型。