硬件教程

4GB 显存为什么能跑 70B？把 AirLLM 的“逐层搬运”拆开讲明白

yiwan-cat2026-08-07 21:08

项目地址：lyogavin/airllm
协议：Apache-2.0
本文核对版本：仓库 main 分支，提交 3a0ad55（2026-08-06）

70B 模型的 BF16 权重，常规情况下要占约 140GB 空间。可 AirLLM 给出的说法是：一张 4GB 显存卡，也能推理原始精度的 Llama 3.x 70B。

第一次看到这句话，最自然的怀疑是：它是不是偷偷量化了？或者只是把模型的某一部分跑起来？

都不是。AirLLM 的关键不是把 140GB 权重“压进”4GB 显存，而是干脆不让它们同时出现：权重绝大部分留在 SSD；GPU 每次只拿到眼前要计算的一层，算完就立刻释放，再换下一层。

它像一条把大型家具搬进窄电梯的流水线。你不需要把整栋楼的家具都塞进电梯；一次搬一件，慢一些，但能到达。

这篇不把它写成一篇“低显存跑大模型”的宣传稿。我们重点把它的真实实现流程、它为什么对 MoE 更有意思，以及最重要的边界讲清楚。

先把结论摆在前面：它解决“能不能加载”，不解决“能不能实时聊天”

AirLLM 把显存需求从“整个模型有多大”，换成了“当前这一个模块有多大，加上 KV Cache、激活和运行余量有多大”。所以，70B 模型可以在小显存 GPU 上被执行；但每生成一个 token，依然要让大量层轮流从磁盘读取、传入 GPU、参与计算、再释放。

这带来两个同时成立的结论：

对离线摘要、文档批处理、评测、学习验证，它提供了一个原本没有的本地入口；
对在线对话、代码补全、客服这类持续交互任务，磁盘 I/O 往往使它体验很差，不能拿“4GB 能跑”推导出“4GB 能用得顺”。

这是理解 AirLLM 的第一把钥匙：它是 out-of-core inference（外存推理） 的工程实现，而不是显存扩容魔法。

常规模型加载为什么会爆显存？

普通 Transformers 推理的默认姿势是：实例化完整网络，再把所有参数放进 GPU。模型的每一层都要随时可被调用，因此显存大致要容下：

$$\
\text{VRAM} \approx \text{全部权重} + \text{KV Cache} + \text{激活/临时张量} + \text{运行余量}.\
$$

对于自回归 Transformer，一次前向传播的层顺序是固定的：词嵌入 $\rightarrow$ 第 0 层 $\rightarrow$ 第 1 层 $\rightarrow \cdots \rightarrow$ 归一化 $\rightarrow$ 输出头。既然同一时刻真正计算的只有一层，AirLLM 就问了一个很工程的问题：其余层为什么要常驻显存？

它给出的答案是：不用常驻，把它们放到磁盘中。

关键不是重写模型，而是在原模型前后“插队”

早期理解 AirLLM 时，很容易以为它自己手写了 Llama、Qwen、DeepSeek 的整套 attention 和生成循环。当前仓库的实现更巧：它仍然让 Hugging Face Transformers 负责正常的 forward() 与 generate()，自己只负责权重在何时出现、何时消失。

实现上有三个关键动作。

第一步：先把检查点改造成“按模块可取”的磁盘文件

第一次加载某个 Hugging Face 模型时，AirLLM 会检查权重索引，把 embedding、每个 decoder layer、norm、lm_head 拆成独立的磁盘 shard。之后就可以按模块名读取，而不必每次把整个检查点读入内存。

这也是为什么第一次启动并不轻：要下载模型、处理分片，并在磁盘上留下派生文件。原始实现还会优先对“本来就一模块一文件”的 safetensors 权重做硬链接，而非再复制一份。对于数百 GB 甚至 TB 级模型，这一点决定了它是否现实。

如果磁盘紧张，可以设置 delete_original=True，在完成转换后删除原始下载权重；但这不是无代价的开关——请先确认保留的分层文件完整、且自己不再需要原目录。

第二步：在 `meta` 设备上搭出一个“空壳模型”

AirLLM 用 init_empty_weights() 在 PyTorch 的 meta 设备上构造真实的 AutoModelForCausalLM。模型的层结构、generate()、attention 逻辑都在，但大参数并没有分配真实显存或内存。

这种设计有两个实际好处：

不必为每一种模型结构复制一套推理代码；只要 Transformers 已支持新结构，适配门槛会低很多。
AirLLM 不接管 token 生成的语义逻辑，主要处理“参数从哪里来”。

当然，“主要”不等于“零适配”。仓库仍然要识别不同架构的模块路径；例如 Kimi K3 的文本 decoder 位于多模态包装器更深的一层，需由专门的 AirLLMKimiK3 类声明路径。

第三步：在每层的前后挂钩子，完成真正的搬运

AirLLM 给 embedding、每一个 decoder layer、最终 norm 和 lm head 注册了两个 hook：

forward pre-hook： 从磁盘加载这一模块的权重到 CPU，再放到 GPU；
forward hook： 模块算完后把刚刚放入的参数迁回 meta，并清理 CUDA 缓存。

整个单 token 的主流程可以概括成：

flowchart LR
    A["SSD: 分层权重"] --> B["CPU: 读取当前层"]
    B --> C["GPU: 当前层计算"]
    C --> D["释放当前层参数"]
    D --> E["下一层重复"]

注意，模型的中间激活和 KV Cache 仍在 GPU；被循环搬运的是大头——各层权重。因此上下文越长、并发越大，KV Cache 越可能成为新的显存瓶颈。4GB 不是一条对所有输入、所有长度都成立的硬上限。

SSD 为什么会直接决定体验？预取只是在遮住一部分等待

层权重搬运的路径是 SSD → CPU RAM → GPU VRAM。当 GPU 正在算第 $i$ 层时，AirLLM 可以在一个后台线程预读第 $i+1$ 层，让部分磁盘读取与当前计算重叠。这就是预取（prefetching）。

可以把理想的单层耗时理解为：

$$\
t\_{\text{layer}} \approx \max(t\_{\text{disk+transfer}},;t\_{\text{compute}}),\
$$

而不是两者严格相加。仓库 README 给出的历史更新中，预取带来约 10% 的提速。

但它不会消灭 I/O：当单层计算很快、模型层又很大时，SSD 读取依然在主导时间。当前代码还限制单个预取层最多约 2GB 的 pinned memory；更大的层退回普通可分页内存，避免为了预取把主机内存锁死。

一个值得注意的版本细节是：README 的配置说明仍写着预取“目前仅 AirLLMLlama2 支持”，但当前 main 的通用 AirLLMBaseModel 已将预取逻辑放入通用 streaming hook。实际使用时，应把你的 AirLLM、Transformers、CUDA 与模型版本固定下来测试，而不是只根据一行 README 估算速度。

为什么 MoE 模型会更有戏？从“按层”进一步变成“按专家”

对于稠密模型，AirLLM 一次至少要搬一整层；对于 MoE（Mixture of Experts）模型，一个 token 在一层中只会被 router 分到少数专家。于是当前仓库加入了更细的策略：不是把整层专家都搬进 GPU，而是给每个 expert 单独挂 hook，只在它真的被路由命中时读取该 expert 的张量。

例如仓库对 Kimi K3 的注释写明：每层有 896 个专家、每个 token 路由到 16 个。全部专家展开后约 55GB，而该 token 实际需要的专家权重约 1GB。这里的节省来自稀疏路由本身，不是 AirLLM 凭空让 55GB 变成 1GB。

这也解释了 README 中两个看似夸张、但需要准确理解的数字：DeepSeek-V3（671B）约 12GB、Kimi K3（2.8T）实测端到端 3.72GB。它们描述的是特定硬件、特定版本与特定推理路径下的峰值显存占用；总模型仍需要相应的本地存储空间，且生成速度高度依赖读取路径。

想要更快：它确实提供压缩，但这时就不再是“原始权重推理”

AirLLM 默认 compression=None，对应不额外压缩其分层权重。如果把初始化参数设为 compression='8bit' 或 compression='4bit'，它会调用 bitsandbytes 对磁盘分层文件做 block-wise quantization，减少读盘和传输的数据量。

图 1 官方仓库给出的同一推理案例：未压缩 449 s，8bit 为 237 s，4bit 为 157 s

图中的数字说明一件事：在 I/O 主导的模式下，少搬数据会明显更快。官方将这一选项描述为“最高约 3 倍”加速；但它不是免费午餐：权重经量化后，精度/鲁棒性是否可接受，仍应在自己的任务上评估。

另外，当前代码会在开启 compression 时关闭预取，因为两者暂不兼容。不要把“压缩 3 倍”和“预取 10%”机械相乘，实际结果取决于模型、SSD、PCIe、CPU RAM、输入长度和采样长度。

快速开始：复制这段代码，跑出第一条结果

这一节的目标不是第一天就去挑战 235B 或 DeepSeek-V3，而是先确认三件事：GPU 和 CUDA 能用、模型能下载、分层流式加载能完整走完。 AirLLM 官方 README 当前也以 Qwen/Qwen3-32B 作为 AutoModel 的快速开始示例。

0. 开跑前，先留出三种空间

AirLLM 的“低显存”只针对 GPU；第一次运行仍会下载模型、生成按层文件，并在 CPU 内存与显存之间搬运。建议先确认：

NVIDIA GPU + 可用 CUDA PyTorch： 以下最小示例会调用 .cuda()；无 CUDA 环境不能原样执行。
本地 SSD： 分层过程会占用大量磁盘；将派生层文件放在 NVMe SSD 上，速度差异非常明显。
系统内存： 权重经由 CPU RAM 进入 GPU，内存也不能只按“4GB 显存”来估算。

如果你只是第一次验证流程，先从自己机器能较快下载、磁盘也容得下的开放模型开始；不要在缓存空间没确认前直接拉取超大模型。

1. 安装 AirLLM

先在已经能使用 CUDA PyTorch 的 Python 环境中安装：

pip install -U airllm

先不要装压缩依赖。 最小路径默认 compression=None，对应本文前面讨论的原始权重流式路径。跑通后，如果你的目标是缩短读盘与传输时间，再额外安装：

pip install -U bitsandbytes

2. 保存一个可直接运行的脚本

新建 run_airllm.py，粘贴下面完整代码。这里保留官方示例使用的 padding=False，可避开部分 tokenizer 没有 padding token 的报错；同时指定 layer_shards_saving_path，避免把大体积的分层文件不知不觉堆进系统盘。

from airllm import AutoModel

MAX_LENGTH = 128

# 官方 README 的当前示例模型。把分层文件明确存到大容量 SSD。
model = AutoModel.from_pretrained(
    "Qwen/Qwen3-32B",
    layer_shards_saving_path="/data/airllm_layers",
)

input_text = ["请用三句话解释什么是外存推理。"]
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=80,
    use_cache=True,
    return_dict_in_generate=True,
)

output = model.tokenizer.decode(
    generation_output.sequences[0],
    skip_special_tokens=True,
)
print(output)

把 "/data/airllm_layers" 改成你自己确实存在且空间充足的 SSD 目录。Windows 用户可改成例如 "D:/airllm_layers"；路径不需要与 Hugging Face 缓存目录相同。

然后执行：

python run_airllm.py

第一次运行的典型过程是：下载检查点 → 按模块拆分/建立分层文件 → 对每层依次加载、计算和释放 → 打印生成结果。第一次的下载与拆分耗时不代表后续单次生成速度；第二次运行会复用已经生成的分层文件。

3. 跑通后，如何换成更大的模型？

代码逻辑不变，只替换 from_pretrained() 中的模型 ID。官方 README 给出的两个当前示例如下：

# Qwen3-235B-A22B（MoE）：官方标注约 3GB 峰值显存
model = AutoModel.from_pretrained("Qwen/Qwen3-235B-A22B")

# DeepSeek-V3（MoE）：官方标注约 12GB 峰值显存
model = AutoModel.from_pretrained("deepseek-ai/DeepSeek-V3")

这里的“约 3GB / 12GB”是仓库披露的特定运行路径下显存峰值，不是完整部署成本：模型下载、分层文件和 CPU 内存需求仍会随模型规模增加。第一次换成超大模型前，先把 layer_shards_saving_path 指向有足够余量的 SSD。

如果改用 Hugging Face 上的受限模型（例如部分 Llama），在确认已经获得访问权限后传入 token：

import os

model = AutoModel.from_pretrained(
    "meta-llama/Llama-3.1-70B",
    hf_token=os.environ["HF_TOKEN"],
)

先在终端设置 HF_TOKEN，或由密钥管理工具注入；不要把真实 token 写进代码、文章截图或 Git 仓库。

4. 可选：用 4bit / 8bit 压缩换取更短等待

确认无压缩路径可跑后，安装过 bitsandbytes 的读者可以这样打开 AirLLM 的 block-wise 权重量化：

model = AutoModel.from_pretrained(
    "Qwen/Qwen3-32B",
    layer_shards_saving_path="/data/airllm_layers",
    compression="8bit",  # 或 "4bit"
)

它的直接收益是少读、少传输权重；代价是这已不再是原始精度推理，且当前实现会关闭预取。建议先比较同一 prompt 的输出质量和耗时，再决定是否将它用于你的任务。

5. 三个最常见的“明明照着做却跑不起来”原因

现象	优先检查什么
`MetadataIncompleteBuffer` 或分层中断	多半是磁盘空间不足；清出空间后重新下载/分层。
`CUDA out of memory`	缩短输入与生成长度；KV Cache、激活和运行余量仍需要显存。
`.cuda()` 报错或找不到 GPU	当前环境没有可用 CUDA PyTorch；先验证 PyTorch 与驱动，而不是反复重装 AirLLM。

至此，你已经不是“看懂了 AirLLM”，而是实际走完了它的核心路径：模型落盘、按层拆分、从磁盘流向 GPU、生成结果。接下来再回看前面的 hook 与预取设计，会更容易理解它为什么能跑、又为什么不会快。

第一次使用最容易踩到的四个坑

1. 看的是显存，忘了看磁盘

AirLLM 降低的是 VRAM，不是模型文件体积。第一次拆分期间可能同时存在原始权重和派生分层文件；模型越大，SSD 空间和写入时间越重要。把 layer_shards_saving_path 指到容量充足的本地 SSD，是最实在的配置。

2. 机械硬盘不是“能跑就行”的小差别

对 AirLLM 来说，存储介质是推理链路的一部分。机械硬盘会把大量随机/连续读取的等待放大，甚至让可运行变成不可用。NVMe SSD、足够的系统内存和稳定的 PCIe 通道，比在生成参数上微调几个 token 更值得先投入。

3. “显存只够一层”不等于任何上下文都够

长上下文下 KV Cache 会持续增长；多轮对话、多并发或很长的 max_new_tokens 都可能触发新的 OOM。先用短 prompt、较小的生成长度确认流程，再逐步扩大负载。

4. 大模型能加载，不代表依赖组合能自动成功

前沿模型经常有特殊依赖。仓库对 Kimi K3 明确列出了 compressed-tensors、flash-attn、CUDA 12 构建的 PyTorch 与 Transformers 4.56.x 等要求；README 同时提示其远程代码不适配 Transformers 5.x。运行前应按目标模型的仓库说明锁定依赖版本。

它最适合谁？

我会把 AirLLM 放在下面这三类任务里：

场景	是否适合	原因
本地批量摘要、离线标注、夜间评测	很适合	吞吐不必实时；小显存机器获得运行大模型的能力。
模型效果摸底、课程学习、私有数据试验	适合	可在本机验证较大模型的原始权重表现。
在线聊天、代码补全、语音助手	不适合	每个 token 都反复经过存储读取路径，首 token 和持续输出延迟通常不可接受。
生产高并发服务	通常不适合	权重流式搬运与 KV Cache、并发需求天然冲突。

如果你的目标只是“在 8GB/12GB 显存上顺畅聊天”，优先考虑 llama.cpp、Ollama、vLLM 等工具配合合适的量化模型。AirLLM 的独特价值不在于取代它们，而在于：当你不愿量化、或只是想验证超大模型确实能在本地跑通时，它把显存这道硬门槛拆成了可排队通过的层。

最后：AirLLM 真正厉害的地方，是把“不可能同时装下”改成“可以按顺序完成”

AirLLM 的技术主线其实很朴素：分层落盘 → meta 空壳模型 → 前置 hook 加载 → 后置 hook 释放 → 可选预取 → MoE 按专家加载。

它没有改变模型需要计算多少层，也没有让 SSD 变成 HBM；它改变的是权重驻留的位置和时刻。于是，“我的 GPU 放不下这个模型”不再必然等于“我无法在本地运行这个模型”。

但请记住这条边界：低峰值显存，只证明模型可以被执行；它不自动承诺低延迟、低成本或好的交互体验。 这恰好也是 AirLLM 最诚实、最值得研究的地方。