4GB 显存跑 70B、12GB 跑 671B：AirLLM 的逐层流式推理

📝

AirLLM 不把整模型塞进显存，而是让权重按层短暂经过 GPU——70B 约 4GB、Llama 3.1 405B 约 8GB、DeepSeek-V3 671B 约 12GB 显存即可运行。省的是权重显存峰值，代价是磁盘 I/O 与 KV cache。

本文转载自微信公众号，如有侵权请联系删除。

AirLLM 不把整模型塞进显存：让权重按层短暂经过 GPU。目标是 70B 约 4GB、Llama 3.1 405B 约 8GB、DeepSeek-V3 671B 约 12GB 显存。

边界：省的是权重显存，不是一切内存

AirLLM 是 Python 推理包，入口为 AutoModel.from_pretrained()，可传 Hugging Face ID 或本地目录。Llama 2/3/3.1/3.3/4、Qwen 1/2/2.5/3（MoE、FP8）、DeepSeek V2/V3/R1、Mistral/Mixtral、Phi、Gemma、ChatGLM、Baichuan、InternLM、Yi 等。当前 setup.py 版本为 3.0.1，运行时依赖 PyTorch、Transformers、Accelerate、Safetensors、Hub。

保留 Transformers 的 tokenizer、generate() 与生成参数；没有定义 CLI 服务、推理 API 或 OpenAI 兼容服务。默认 CUDA；macOS 要求 Apple Silicon、mlx、torch，且源码走 Llama 专用 MLX 实现。

权重不常驻显存，也不等于没有上限：生成仍保留 KV cache，长 prompt、batch、激活和缓存都会增长。AirLLM 改变的是模型权重的峰值，不是注意力和上下文的一切成本。

功能：不只是”逐层加载”

自动分派。 AutoModel 读取 AutoConfig.architectures；ChatGLM、旧 QWen、Baichuan、InternLM 走专用子类，其余走通用 AirLLMBaseModel。通用路径让 Transformers 执行真实 forward/generate，避免为每个新架构复制推理循环。

拆层、落盘、复用。 初次加载把 embedding、每个 decoder layer、norm、lm_head 分成独立层文件；非 macOS 用 Safetensors，MLX 路径写 .mlx.npz。完整文件和 .done 标记存在便复用。源码兼容带索引分片、单文件 model.safetensors、pytorch_model.bin 与 tied embedding。

流式运行。 模型先由 Accelerate 在 meta device 建空壳；embedding、每个 decoder layer、norm、lm_head 都挂前/后钩子。前钩子从磁盘取当前层并移到设备，后钩子将其移回 meta、清缓存，因此 GPU 不需要整模型权重。

预取和剖析。 默认 prefetching=True，单线程在当前层计算时读取下一层，CUDA 下固定 CPU 页；profiling_mode=True 累积读取/解压时间至 model.profiler。压缩会关闭预取。

4/8-bit 块压缩、FP8。 compression='4bit'/'8bit' 需要 bitsandbytes。4-bit 使用 NF4、块 64；8-bit 使用 blockwise、块 2048，并保存量化状态。README 将其定位为压缩磁盘读取量、最高约 3 倍推理加速、精度损失很小，不是量化激活。当前代码还保留 FP8 权重与 weight_scale_inv，接入 Transformers 预量化配置。

下载与空间管理。 远程仓库先取配置/索引，再按需拿原始权重分片；hf_token 用于 gated 模型，layer_shards_saving_path 指定层文件盘，delete_original=True 会删除拆完的原 checkpoint。

安装、源码安装与第一段可运行代码

python -m pip install -U airllm
## 仅使用压缩时需要
python -m pip install -U bitsandbytes

git clone https://github.com/lyogavin/airllm.git
cd airllm/air_llm       # setup.py 在这里，不在仓库根目录
python -m pip install -e .

源码包要求 torch>=2.4、transformers>=4.49,<5.13、accelerate>=1.0 等。根目录 requirements.txt 是 Anima 训练实验的旧依赖集，不是 AirLLM 运行时安装单。macOS 另装 mlx、torch，且只支持 Apple Silicon。

from airllm import AutoModel

model = AutoModel.from_pretrained(
    "Qwen/Qwen3-32B",
    # layer_shards_saving_path="/fast-ssd/airllm-layers",
    # compression="4bit",       # 或 "8bit"
    # hf_token="HF_TOKEN",      # gated 模型
)

tokens = model.tokenizer(
    ["What is the capital of the United States?"],
    return_tensors="pt", return_attention_mask=False,
    truncation=True, max_length=128, padding=False,
)
out = model.generate(
    tokens["input_ids"].cuda(), max_new_tokens=20,
    use_cache=True, return_dict_in_generate=True
)
print(model.tokenizer.decode(out.sequences[0]))

首次构造应视为预处理：下载和拆层都很吃磁盘，空间不足常报 MetadataIncompleteBuffer。层文件放快 SSD；慢盘会拖慢每个 token。先以”未压缩+预取”建基线，再比较 4/8-bit——文件小不必然更快。没有 padding token 时设 padding=False 或补 pad。

源码如何做到：一层一层经过 GPU

一次 generate 不是加载一层只算一次。自回归每产出一个 token，都要过 embedding、全部 decoder layer、norm、lm_head；AirLLM 让它们每轮循环进出 GPU。代价是重复磁盘 I/O、CPU→GPU 传输和 Python 钩子调度：小显存、快磁盘、输出不长最符合设计；追求极限 tokens/s、超长上下文或大批量服务时，换入换出会主导成本。

初始化时 init_empty_weights() 先置空参数，AutoModelForCausalLM.from_config() 优先 SDPA、失败回退 eager attention；默认 dtype 取模型配置，源码避免深模型一律强制 FP16。执行时 set_module_tensor_to_device() 注入当前层；bitsandbytes 权重重建量化参数，FP8/scale tensor 原样搬运。普通层执行后 module.to('meta')，量化层逐参数回 meta。于是”无需量化也能小显存运行”与”压缩可再加速”可以同时成立。

同仓库的三组实验：完整，但不属于 AirLLM 运行时

training/：Anima 33B 的 QLoRA SFT。支持 Alpaca、CHIP2、Self-Instruct、HH-RLHF、OASST1、中文 Vicuna、本地 JSON/JSONL/CSV/TSV，含 4/8-bit、LoRA、梯度检查点、断点恢复、PEFT 保存、样例生成、MMLU 回调；给出的复现脚本是 Guanaco 33B、batch 1、累积 16、10000 steps，文档记录测试于 80GB H100 单卡或两张 40GB A100。
rlhf/：QLoRA + DPO。比较 policy/reference 对 chosen/rejected 的 token log-prob，以 -logsigmoid(beta × log-ratio 差) 训练；示例是 100 steps、LoRA r=256、beta=0.1，不是 AirLLM 推理功能。
anima_100k/：Anima-7B-100K 长上下文训练，包含自定义 Flash Llama、FlashAttention2 分块注意力、XEntropy inplace logits、paged 8-bit AdamW、LoRA；脚本设 source/target 为 92000/1024。不增加 AirLLM 的上下文长度，也不会随 pip install airllm 自动启用。

AirLLM 的工程交换很明确：容量交给磁盘分层，正确性回到 Transformers，显存峰值收敛至单层；代价是首轮拆分、I/O 与 KV cache。先用小模型核对输出与层文件路径，再放大模型。