4GB 显存跑 70B，12GB 跑 671B：拆解 AirLLM 的"分层流式"魔法

大模型圈有一类话题永远有热度：”我家显卡只有 X GB，能跑多大模型？”

以前这个问题的答案很残酷——8B 要 16GB，70B 要 80GB，405B 得上 4 卡 H100，Kimi K3 这种 2.8T 的怪物根本别想本地玩。但 GitHub 上有个叫 AirLLM 的项目，把这套心智模型掀了个底朝天：

模型	参数量	AirLLM 实测 VRAM
Qwen3-8B / Mistral-7B	8B	~1–2 GB
Qwen3-30B / Mixtral	30–47B	~1–3 GB
Qwen3.8-27B (含视觉)	27B	3.33 GB
Qwen3.8-Flash-Next	~180B (MoE + PLE)	5.95 GB
Qwen3-235B	235B	~3 GB
Llama 3.x 70B	70B	~4 GB
Llama 3.1 405B	405B	~8 GB
DeepSeek-V3	671B	~12 GB
Kimi K3	2.8T	3.72 GB

你没看错——671B 全量化训练挂一层在 12GB 卡上能跑，2.8T 的 Kimi K3 居然在 3.72GB 上能跑完一整轮生成。

它没有用蒸馏、剪枝，也没有用通常意义上的量化。它做的是一件更朴素、但工程上更狠的事：把模型拆成一层一层，按需从硬盘搬到显存，算完就扔。

这一篇，我就把它的源码扒一遍，告诉你这件事到底是怎么做到的，以及为什么”它能在 4GB 上跑 70B”这件事，没有看上去那么神奇，但绝对比你想象的更精巧。

一、先说朴素想法：你不需要同时把模型放显存里

我们直觉里，”跑 70B 模型” 等于 “70B 参数都得驻在显存里”。这是错的。

Transformer 是顺序执行的——第 N 个 token 的 KV 缓存只依赖前 N-1 个，而每一层 decoder 的计算，都只依赖当前这一层的权重和上一层输出的 hidden state。

也就是说：任意时刻，显卡上真正”必须存在”的，其实只有当前那一层 decoder 的权重 + KV 缓存。 其它几十层、几百层，全都可以在硬盘上躺着。

理论上，一个 70B 模型、单层大概十几个 GB。等等，这也不对啊，4GB 怎么装得下？——因为作者把所有权重都用bf16 加载到 meta 上”假装存在”，然后按 layer 切到磁盘，每次只把单层的 safetensors shard 读到内存再拷到 GPU。

这就是 AirLLM 整个设计的灵魂：Layer-wise Streaming（分层流式加载）。

二、核心原理：meta 占位 + 前向 hook 调度

整个 AirLLM 的代码量其实不大（核心 4k 行左右），关键技术秘密就在 airllm_base.py 里。我把它拆成三步：

2.1 步骤一：把检查点先在硬盘上按层切好

模型从 HuggingFace 下载下来通常是按”分片”切的（比如 model-00001-of-00030.safetensors），AirLLM 拿到的是一组散落的分片，它做的第一件事是把它们重新组织成”按层”：

splitted_model/
  model.embed_tokens.safetensors
  model.layers.0.safetensors
  model.layers.1.safetensors
  ...
  model.layers.79.safetensors
  model.norm.safetensors
  lm_head.safetensors

每个文件恰好对应模型 forward 路径上的一个 module：embedding、每一层 decoder、最终 norm、输出头。

这一步做完之后，每一层的”独立打包”就齐了，后面的流式加载才有意义。

2.2 步骤二：模型建在 meta 设备上，零显存占用

接下来用 transformers 的 AutoModelForCausalLM（或 AutoModelForImageTextToText）把整个模型建在 meta 设备上。meta 设备上的 tensor 不分配真实内存，只占”形状”——整个 70B 模型在你机器里此刻0 字节。

这时候整个模型的骨架（每一层的 Linear、RMSNorm、Attention、MLP）已经搭好，forward 路径也通了，KV 缓存的逻辑也都就绪。唯一缺的，是每层 Linear 的真实 weight——这才是 AirLLM 介入的地方。

2.3 步骤三：用 forward pre/post hook 把权重”按需请进 GPU”

最关键的魔法来了。AirLLM 给每一个需要流式加载的 module（embed / 每个 decoder layer / norm / lm_head）都注册了一对 hook：

# 伪代码，简化自 airllm_base.py
def _pre_hook(module, args):
    # 进这一层之前：把这一层的权重从硬盘搬到 GPU
    state_dict = self._load_streamed_layer(idx)
    module._airllm_moved = self.move_layer_to_device(state_dict)

    # 顺手把下一层异步预取到内存（关键加速技巧）
    if self.prefetching:
        submit_next_layer_to_threadpool(idx + 1)

def _post_hook(module, args, output):
    # 出这一层之后：把这一层的权重扔回 meta
    set_module_tensor_to_device(module, 'meta')
    clean_memory()
    return output

也就是说，每次 forward 只在 GPU 上同时存在一层 decoder 的权重 + 上一层留下的 hidden state + KV 缓存。这就是为什么 4GB 显存能跑 70B——它的有效内存占用 = 单层权重 + KV 缓存，而不是整个模型。

三、那些让它真正能用的工程细节

朴素想法谁都能想到，但 AirLLM 能跑 Kimi K3，靠的是一堆硬核工程细节。我挑几个最重要的讲。

3.1 预取：把 IO 和计算重叠起来

朴素实现的瓶颈是：算一层 → 加载下一层 → 算下一层 → 加载下下一层……加载时间是浪费的。

AirLLM 的解法是双缓冲预取。它开了一个单线程的 ThreadPoolExecutor：

def _pre_hook(module, args):
    # 如果"下一层"已经被后台线程加载好了，直接拿结果
    if self.prefetch_future and self.prefetched_idx == target_idx:
        state_dict = self.prefetch_future.result()
    else:
        state_dict = self._load_streamed_layer(idx)

    module._airllm_moved = self.move_layer_to_device(state_dict)

    # 把"下一层"的加载扔进后台线程
    next_future = executor.submit(self._load_streamed_layer, idx + 1)

效果：当前层 GPU 在算的时候，后台线程在磁盘上读下一层的 safetensors。IO 和计算并行，整体延迟被加载时间掩盖，实测提速 10%–30%。

3.2 4bit / 8bit 块级量化：再省 3 倍带宽

朴素版本还有一个痛点：硬盘到 GPU 的 PCIe 带宽。一张 RTX 4090 的 PCIe 是 ~32 GB/s，一个 70B 模型单层 bf16 大概 4 GB，读完一层要 100+ ms，而 GPU 算一层只要 50 ms——加载比计算还慢。

AirLLM 提供了一个参数：

model = AutoModel.from_pretrained(
    "garage-bAInd/Platypus2-70B-instruct",
    compression='4bit',  # 或 '8bit'
)

它会在切层的时候用 bitsandbytes 的块级量化把每一层的权重压成 4bit/8bit，同时保留 scale。读 4bit 流所需的 PCIe 带宽只有 bf16 的 1/4，整体推理速度能再快 3 倍。

注意它和普通量化推理的区别——普通量化（GPTQ/AWQ）需要同时量化权重和激活才能加速，因为它们的瓶颈在矩阵乘；AirLLM 的瓶颈是加载，所以只需要压权重，激活保持 bf16 不动。精度几乎不掉，但省下来的 PCIe 带宽是实打实的。

唯一限制：量化后没法做预取（两者不兼容）。

3.3 块级量化 1 层优化，量化 1 层后必须留在 GPU

开启压缩之后还有一个细节：bitsandbytes 的量化是整个 Linear 层一次性算完的，所以层内的反向 hook 不能再在算完后就 meta 化——必须带着整个量化模块一起跑完。这也是为什么 _post_hook 里有个特殊分支：

if (self.hf_quantizer is not None
    or getattr(self, '_expert_streaming', False)
    or getattr(self, '_cpu_resident_params', None)):
    # 只 meta 掉这一次"搬到 GPU 的张量"，不 module.to('meta')
    for param_name in getattr(module, '_airllm_moved', []):
        set_module_tensor_to_device(self.model, param_name, 'meta')
else:
    module.to('meta')  # 整层 meta 化，最干净

这是从代码里抠出来的工程妥协，量化模型和流式模型不能”一刀切”地释放。

3.4 MoE 按专家粒度流式：Kimi K3 的关键

前面那张表里 Kimi K3（2.8T）居然能跑在 3.72GB，这其实是 AirLLM 最戏剧性的部分——单纯按”层”流式还不够。

Kimi K3 是 MoE，每个 token 只激活一小撮专家。如果把整层 MoE（几百个专家）一次性搬上 GPU，单层就 55GB 起步，根本放不下。

AirLLM 的解法是按 expert 流式——它识别到 layer_names_dict['expert_prefix']，然后给每一个专家模块单独挂一对 hook：

def _expert_pre_hook(self, module, args):
    layer_idx, expert_idx = module._airllm_expert
    keys = self._expert_keys[layer_idx][expert_idx]
    # 只加载这一个专家的 tensor（一般就几十 MB）
    state_dict = load_layer_subset(self.checkpoint_path, ..., keys)
    module._airllm_moved = self.move_layer_to_device(state_dict)

效果：真正只有被这个 token 路由到的那些专家才会被搬到 GPU 上。一个 token 只命中 4–8 个专家，每个专家几十 MB，加起来 1–2 GB——3.72GB 就够 Kimi K3 推理了。

代码里还有个细节很关键：

1
2
3

# 启用 expert streaming 之后，layer 自己的 hook 里会跳过这些 keys
keys = self._non_expert_keys.get(idx) if getattr(self, '_expert_streaming', False) else None
# 否则整层 keys，包含非 expert 和 expert

——避免 layer 的 hook 把 expert 的权重也加载一遍，造成双重占用。

3.5 Flash-Next 的 n-gram 表：mmap 留在内存

Qwen3.8-Flash-Next 有一个 51B 参数的 n-gram embedding 表（PLE），整整 102GB 的 bf16，没有任何一层 GPU 能装得下。

AirLLM 的解法：用 mmap 把这张表文件映射到虚拟内存，永远不拷贝到物理 RAM。当 transformer 调用 ngram_embedding(...) 时，它直接从硬盘页面里读：

1 2	table = open_ngram_mmap_table(self.checkpoint_path, name) self._install_mmap_embedding(name, table) # 替换原 module

需要哪一行就读哪一页，操作系统帮你做 page cache。一台 64GB 内存的机器就能跑 125B 的 Flash-Next。

3.6 自动识别模型架构：`AutoModel`

AirLLM 最贴心的一点是 AutoModel.from_pretrained(...) 一行就能跑几乎所有主流模型：

1 2	from airllm import AutoModel model = AutoModel.from_pretrained("Qwen/Qwen3-32B") # 任意 HuggingFace 模型

背后是一段很精妙的”试错链”：

def _auto_model_classes(self):
    archs = self.config.architectures or []
    arch = archs[0] if archs else ""
    names = []
    # 视觉语言架构优先试 text + image factory
    if any(tag in arch for tag in ("ConditionalGeneration", "ImageTextToText", "Multimodal")):
        names.extend(["AutoModelForImageTextToText", "AutoModelForMultimodalLM"])
    # 然后才是常规路径
    names.extend(["AutoModelForCausalLM", "AutoModelForImageTextToText",
                  "AutoModelForMultimodalLM", "AutoModel"])
    ...

然后逐个尝试这些 Auto 类，能实例化出”模块名匹配 layer_names_dict 的那个”就用，否则换下一个。VL 模型如果只走 AutoModelForCausalLM 会拿到一个文本版本（Qwen3.5 之前就踩过这个坑），所以视觉架构必须优先试 image-text factory。

作者甚至还做了个”老 transformers 兼容性补丁”：

_RELOCATED_TRANSFORMERS_SYMBOLS = {
    'OutputRecorder': 'transformers.utils.output_capturing',
    'check_model_inputs': 'transformers.utils.output_capturing',
}

——transformers 5.0 把这两个函数搬了家，但很多模型的 remote code 还在 transformers.utils.generic 里 import 老路径。新版本 transformers 找不到就崩。AirLLM 主动 re-export，让旧 remote code 在新 transformers 下也能加载。非常贴心但也很 hack。

四、更狠的：流式 LoRA 训练

2026 年 9 月，AirLLM 加了一个比推理更炸裂的能力：流式 LoRA 训练。

思路完全一样的对偶：

推理时：base 权重是只读的，可以流式加载。
训练时：base 权重仍然是只读的（冻结），也完全可以流式加载；只有 LoRA 参数需要梯度。

也就是说：训练 125B 模型时，base 权重继续从硬盘流过，只有 LoRA adapter（几十 MB）留在显存里。结果：

模型	训练 VRAM	说明
Qwen3.8-27B	~2 GB	seq 512
Qwen3.8-Flash-Next (125B)	~6 GB	RTX 3060 Ti 即可

用法也很简洁：

from airllm import AirLLMLoRAQwen4Exp

trainer = AirLLMLoRAQwen4Exp(
    "Qwen/Qwen3.8-Flash-Next",
    max_seq_len=512,
    lora_r=16,
    delete_original=True,  # 训练完删掉原检查点，省硬盘
)

loss = trainer.train_step(input_ids.cuda(), attention_mask=...)
trainer.save_adapter("qwen38-flash-next-lora.pt")

或者跑现成脚本：

python air_llm/examples/train_qwen38_flash_next_lora.py \
  --data my_data.jsonl \
  --seq-len 512 \
  --epochs 1 \
  --save-adapter out.pt

底层思路依然是单层 GPU + KV 缓存 + LoRA 参数，只是把”流式加载”和”反向传播 + 优化器 step”调度起来。每一层在 GPU 上时，正向算一次、反向算一次，然后扔掉。没有 HuggingFace Trainer，没有 bitsandbytes QLoRA，纯自己撸——但就是这么朴素的做法，把”125B 模型微调”这件事压到了 6GB。

五、MacOS 支持：MLX 后端

AirLLM 不只是 NVIDIA 的事，它还有一套 MLX 后端（airllm_llama_mlx.py），可以在 Apple Silicon 上跑：

1 2	pip install mlx # 其它用法和 Linux 完全一致

MLX 是 Apple 出的统一内存框架——CPU 和 GPU 共享内存。所以”分层流式加载”在 Mac 上其实是把权重从 SSD 搬到 unified memory，GPU 访问零拷贝。70B 模型在 M2/M3 上也能跑。

六、它不是银弹：必须说清楚的代价

AirLLM 看着像魔法，但它是有明确代价的：

1. 速度。它本质上是”用 IO 换显存”。再叠加预取，单 token 生成速度相比原生 GPU 推理还是要慢几倍。生成 100 token 的 70B 文本可能要几十秒。如果你需要 ms 级响应，请上 H100。

2. 硬盘和带宽。它把 PCIe/SSD 推到了极限。NVMe 比 SATA SSD 强一截；机械硬盘基本不可用。

3. 第一次跑要复制 hash。模型刚下载下来是按分片切的，AirLLM 要先把它按层重切，存成 splitted_model/。这一份和原检查点等大——700GB 的 Kimi K3 就要再花 700GB 临时磁盘。delete_original=True 可以让它切完就删原文件，但风险自负。

4. 顺序执行，没有并行。它无法做 tensor parallel / pipeline parallel 之类的多卡加速（因为每层只在一张卡上）。如果你有 4 张 4090，AirLLM 也只能让你用一张。

5. fp16 要小心。源码里作者特意写了一段注释——强制 fp16 在深层模型（Qwen3-235B 的 94 层）会溢出到 inf/NaN，悄悄毁掉输出。默认用 bf16。

七、为什么这种”反向思考”重要

AirLLM 给我的最大启发不是”4GB 能跑 70B”，而是它演示了一种和主流叙事完全相反的优化思路。

主流叙事：模型越来越大 → 必须有更猛的显卡 → NV 赚麻了。
AirLLM 的回答：模型再大，任意时刻只需要一层的权重 + KV 缓存。剩下的，全都可以在硬盘上睡大觉。

这不只是一个 trick，它实际上指向了 LLM 推理的结构性事实——Transformer 的 sequential 性质意味着”模型大小”和”推理所需显存”之间根本没有线性关系，前者除以层数才是真正的显存下界。这件事，AirLLM 把它工程化做到了极致。

同样的思路可以延展到很多地方：本地 RAG 系统的轻量化、边缘设备部署、教育/科研的低成本入门——你不需要顶配显卡，也能在自己的笔记本上跑起一个”会思考”的大模型。

八、参考资料

AirLLM GitHub
AirLLM 引用的块级量化论文：LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale
同作者的灵感来源：SimJeg 的 Kaggle 70B with Wikipedia RAG