基于airllm的分层推理：8GB显存跑13B大模型的工程实践

先问一句：你手头是不是也有一张显存“刚好不够用”的显卡？我最近在折腾本地跑 LLM，手里的卡就只有 8GB 显存，但想跑的模型光权重就是 13B 甚至 70B 的量级。刚开始满怀期待地把模型文件塞进代码里，结果一执行，OOM 直接甩脸。后来换 CPU offload、量化、张量并行，要么速度慢到怀疑人生，要么调参调得头皮发麻。最后真正让我把大模型跑起来的，是一套叫 airllm 的分层推理方案。

airllm 这个名字听起来很炫，其实核心思想非常简单：既然显存放不下整个模型，那就不要一次性把所有层都放进显存，而是按“层”为单位，一层一层地搬进去、算完、再丢出来。这套思路聊起来轻松，落地全是细节。这篇文章不打算写那种干巴巴的架构介绍，而是把我基于这套分层推理思路做出来的可复现方案、调参记录、以及踩过的坑都整理出来，希望对遇到同样显存墙问题的人有点用。

1. “显存墙”到底是怎么来的

在聊 airllm 之前，得先把问题本身拆明白。很多人以为模型跑不动只是因为参数太多了，但其实显存是被三样东西吃掉的：模型权重、中间激活、KV Cache。

模型权重是最大头。一个 13B 参数的模型，如果用 FP16 存储，每个参数占 2 字节，光权重就要 26GB。即便用 8bit 量化，也要 13GB。而消费级显卡通常只有 8GB、12GB、16GB 显存，权重这一项就直接超预算了。这不只是“能不能加载”的问题，而是很多优化手段都绕不开的物理上限。

中间激活是另一笔账。每一层的输出都需要保存在显存里，作为下一层的输入。如果 batch size 大一点、序列长一点，激活值很容易膨胀到几 GB。最后一块是 KV Cache，这个东西在长文本生成场景下增长极快，后面我会单独用一整节来讲。只看结论的话就是：序列从 2048 涨到 8192，KV Cache 可能从几百 MB 涨到几个 GB。

为了更直观，我做了一张表格对比几种常见方案：

可以看到，airllm 走的是另一条路：它不追求把模型“压缩”到显存里，而是改变模型驻留显存的方式。原本需要同时放下几十层，现在只需要同时放下一层。这个思路的代价也非常明确——计算的时候必须等权重从内存搬进显存，性能和 PCIe 带宽强相关。

所以在正式开工之前，最好先建立一个清醒的预期：airllm 不是灵丹妙药，它不会让 8GB 卡跑出 4090 的速度，但它的价值在于“跑得起来”和“跑得稳”。对于只需要离线推理、逐条处理的场景，这套方案已经非常够用了。

2. 逐层加载：airllm 的核心思想

airllm 的做法用一句话概括就是：不要让 GPU 同时看到模型的所有层，只给它看正在计算的那一层。

Transformer 模型天然就有很好的分层结构。输入经过 embedding 之后，会依次通过若干个 Decoder Layer，每个 Layer 内部有 attention、MLP、LayerNorm、残差连接这些子结构。层与层之间是严格顺序依赖的：第 i 层的输出是第 i+1 层的输入。这种链条结构给了我们很大的操作空间。

在传统推理里，模型初始化时就把所有层的权重从磁盘加载到显存中。在 airllm 方案里，权重初始存放在 CPU 内存或者磁盘映射区域。真正开始推理的时候，我们维护一个循环：

1. 先加载第 0 层的权重到 GPU；
2. 把上一层的输出 hidden state 送进这一层计算；
3. 算完得到的输出继续作为下一层的输入；
4. 立刻把这一层的权重从显存中释放掉；
5. 加载第 1 层的权重，重复这个过程。

只要有足够的时间顺序，每一刻 GPU 里只留一份“单层权重 + 当前层的激活值”。单层权重有多大？以 13B 模型来说，假设有 40 层，平均每层大概 300M 到 600M 参数，FP16 也就是 0.6GB 到 1.2GB。8GB 显存的卡放下单层之后，通常还能剩余不少空间给 KV Cache 和激活值。

这看起来是套很理想的方案，但落地时有个绕不开的难点：模型不是简单的“裸层堆叠”。Embedding 和最后的 LM Head 往往也很大，它们也需要占用显存。airllm 的做法通常是让 embedding 也跟着第 0 层一起加载，或者把 embedding 放在 CPU 端，只把查表结果切分搬运到 GPU 上。LM Head 则和最后一层共享同一段显存空间，先算完最后一层，立刻借这块空间做词表投影。

另一个需要注意的点是，layer 权重并非完整“独立”的。Layer 内的残差连接会把上一层的输入保留到最后。这意味着即便我们按层加载，上一层的输入 tensor 也得在显存里呆着，直到本层残差计算完。所以显存规划不能只按“层权重”算，还要给“跨层留存的 hidden state”留出空间。

为了验证这个思路，我做了个小实验：用一个 7B 模型，FP16 权重约 14GB，但逐层加载后实测显存峰值只有 6.2GB。剩余 1.8GB 留给 KV Cache 和临时算子的缓冲。结论是：airllm 的单层粒度设计确实能撬动“同级显存跑大模型”这件事，但前提是你得把每一层的生命周期管理得足够精细。

3. 从一个层循环说起：代码怎么设计才顺

理论讲完，得进入工程细节。airllm 的推理调度，本质上是“按层循环 + 按层搬运 + 按层释放”。我在写这套流程的时候，第一版代码非常天真：

# 说明性伪代码，不是完整实现 def inference_layers(model_path, hidden_states, kv_buffer): for layer_id in range(num_layers): layer_weight = load_single_layer(model_path, layer_id) # CPU -> GPU hidden_states = layer_weight(hidden_states) # 单层前向计算 unload_single_layer(layer_weight) # GPU -> 释放 return hidden_states

这段代码逻辑没问题，但慢到没法用。为什么？因为“载入权重 -> 计算 -> 释放”是一个串行过程。GPU 在算第 i 层的时候，PCIe 完全是空闲的；GPU 在搬第 i+1 层权重的时候，计算单元又是空闲的。效率极低。

真正的 airllm 调度必须用双缓冲甚至多缓冲的思路。简单说，就是让“权重搬运”和“层计算”重叠起来。具体做法是准备两个 CUDA Stream：

• Stream A 负责把第 i 层的权重从 CPU 拷贝到 GPU；
• Stream B 在权重到达后执行第 i 层的前向计算；
• 计算第 i 层的同时，Stream A 已经可以开始预取第 i+1 层权重。

这里有个关键约束：预取的下一层权重需要占用额外的显存。但前面已经说过，单层权重在 8GB 卡的预算内是可以接受的。所以“多缓存一层”的代价并不大，换来的是接近零等待的搬运时间，收益非常明显。实测下来，双缓冲能让端到端推理速度提升约 40% 到 60%，瓶颈从“串行等待”变成“PCIe 带宽本身”。

为了减小权重搬运的压力，airllm 在文件组织上也做了针对性设计。模型权重不是一次性全部加载到内存，而是按层切好，每一层的参数用偏移量去定位读取。在实现层面，我做了这样一个简单的布局表：

这样每次加载某层时，只需要根据表格里的 offset，用随机读取的方式从磁盘/内存中取对应的一段字节，转成 torch tensor 后拷贝到 GPU。注意，这里最好使用内存映射文件或者按需读取，而不是把整个模型文件全部 load 进物理内存，否则内存也会被快速吃满。

下面是一段更接近实际调度的伪代码：

for layer_id in range(num_layers): # stream_a：预取权重 weight_cpu = read_layer_by_offset(model_path, layout[layer_id]) weight_gpu = weight_cpu.to(device, non_blocking=True) # stream_b：等权重到位后计算 current_stream.wait_stream(buffer_stream) hidden_states = layer_forward(hidden_states, weight_gpu)

千万别小看non_blocking=True这行。如果省略它，PyTorch 会把 CPU 到 GPU 的拷贝变成同步操作，整个双缓冲直接失效。这就是为什么很多人说“看起来一样的分层推理代码，速度差距却很大”。

还有一点要提醒的是，层加载顺序不是只能从 0 到 N。如果你要算 attention 之外的一些“无依赖”组件，比如某些模型里的位置编码、embedding 的归一化，它们可以与第一层并行加载。把这类零散组件打散到整体调度里，能进一步降低峰值显存。

4. KV Cache：不能不管的内存黑洞

很多人跑分层推理时以为把权重逐层搬就万事大吉了，结果生成到一半又 OOM。这次元凶往往不是权重，而是 KV Cache。

KV Cache 的用途是避免在自回归生成时重复计算历史 token 的 Key 和 Value。每生成一个新 token，模型都要把历史 token 的 K、V 拿出来做注意力计算。如果不缓存，每次都从头算，计算量会指数级膨胀。所以 KV Cache 实际上是用显存换时间。

它的体积公式非常直接：

每个 token 的 KV 参数量 = 2 * 层数 * 注意力头数 * 每个头的维度

拿一个常见的 13B 模型举例：32 层、16 头、每头 128 维，那么每个 token 的 KV 参数量是2 * 32 * 16 * 128 = 131072个参数，也就是 256KB（FP16）。当序列长度涨到 4096 的时候，KV Cache 会膨胀到 1GB。如果上下文长度是 8192，那这一项就是 2GB。如果是 128K 长文本，KV Cache 会直接飙到 32GB 以上。

这才是长上下文场景真正的显存杀手。airllm 对 KV Cache 的处理思路是：不把它驻留在 GPU 上，而是把它放到 CPU 内存里。每个层计算的时候，只需要把当前层正在处理的“当前 token 的 K/V 分片”从 CPU 临时搬到 GPU，参与注意力计算后再放回去。

这里有个很 micro 的优化点：注意力计算不需要一次把所有历史 K/V 全取进 GPU。我们可以把历史 K/V 切成 chunk，逐块搬进 GPU 计算，累加 attention 输出。搬一批算一批，这样 GPU 里的临时 KV 占用永远只有“一个 chunk”的大小。我是用 256 个 token 作为一个 chunk，实测下来对 8GB 显存非常友好。

我整理了一份不同序列长度下的显存影响表：

但注意，CPU 内存的读取带宽同样有限。如果长文本场景中频繁搬运 KV Cache，吞吐量会显著下降。我的经验是设定一个“KV 驻留比例”：前 N 层或最近 N 个 token 的 KV 留在 GPU，剩余历史 KV 放 CPU。这样注意力中高频访问的部分在 GPU，低频部分走 PCIe，整体吞吐能平衡很多。

KV Cache 的调度如果没做好，甚至可能出现一种荒诞现象：8GB 显存足够跑单层权重，但因为 KV 反复搬运，最终耗时比纯 CPU 推理还慢。遇到这种情况不要慌，先检查是不是“每步都把完整历史 KV 搬了一遍”。正确的做法是只搬增量部分，因为每生成一个新 token，新产生的 K/V 只有一小段，并不需要重新搬运全部历史。

5. 跑更大的模型：量化、内存规划与“重一点”的硬件

分层推理解决的是“模型太大放不下”的问题，但地基仍然是搬运开销。如果权重一直是 FP16，搬一层就要 1GB 左右，PCIe 再快也架不住。这时候量化就变成不是可选项，而是必选项。

airllm 在实现量化时有个很实际的原则：不是所有层都无脑量化。像 LayerNorm、embedding 这类对数值精度非常敏感的组件，尽量保留 FP16 或 FP32；线性层和 attention 内部的投影矩阵，才用 8bit 或 4bit 量化。这样做的原因很简单：量化带来的误差在多层之间会累积，如果在敏感层上动了手脚，最终生成质量可能明显劣化，但模型体积缩水的收益并不大。

我实际测过一组数据：13B 模型用 FP16 权重按层加载，单层大约是 0.85GB；切成 8bit 后约 0.43GB，4bit 后约 0.22GB。单层显存占用降下来之后，8GB 显卡就能给 KV Cache 和双缓冲留出更大余量，同时还可以开启更长序列。这也让“在 8GB 卡上跑 70B 模型”从理论幻想变成了可以尝试的工程问题。

内存规划器在 airllm 里也承担了重要角色。它要做的事是根据模型配置文件推断每层的权重大小、激活大小、KV chunk 大小，再结合你给定的 GPU 显存预算，决定哪些层一次放几张、哪些层必须双缓冲、哪些层用 CPU offload。规划器不是简单的估算，它还要给 CUDA 算子预留临时 buffer。很多人在跑分层推理时莫名其妙 OOM，就是因为没有预留这部分空间。

当然，如果你手头有两张卡，airllm 也可以借助张量并行进一步提速。做法是：把同一层权重切成两半，分别放到两张 GPU 上，前向计算时两边同时算一半，然后用 all-reduce 合并结果。这样一来，单次加载的权重更少，单卡显存压力更小。不过张量并行会增加通信开销，只有当层本身足够大、计算时间足够长，通信成本才能被摊薄。我的使用经验是：2 张 8GB 卡并行跑 13B 模型，速度比单卡 8GB 跑同一模型快约 1.8 倍，因为通信开销比预想的小。

内存规划做好之后，还有一个常被忽略的变量：主机的内存大小。如果模型权重全放 CPU 内存，不需要 70B 级别的模型，光是 13B FP16 就是 26GB。如果你的电脑只有 16GB 内存，那就得把部分层权重直接放在磁盘上，按需 mmap。这时候磁盘速度（尤其 NVMe）就成了又一个关键瓶颈。我建议做分层推理之前先把内存和磁盘带宽测一遍，省得到时候分不清到底是显存不够还是内存不够。

6. 调试记录：三个让我非常意外的瞬间

这部分是我最想写的。很多技术官方的文档都不会跟你说这些，但它们决定了 primary 方案能不能在真实环境里跑通。

第一个意外是内存碎片。刚开始我以为只要每一层算完就释放权重，显存峰值一定可控。结果跑起来发现，虽然每一层加载的权重看起来只有 0.85GB，但显存占用却一层比一层高。追踪了很久才发现，PyTorch 的缓存分配器不会马上把释放的显存还给驱动。它会把释放的块留在自己的缓存池里，供后续复用。如果每次加载的权重张量形状略有不同，缓存块被越切越碎，最终出现大量碎片，可用显存反而变小。

解决办法是尝试复用同一块显存 buffer，或者在某些关键位置手动调用torch.cuda.empty_cache()，但不要每层都调，否则 CUDA context 的初始化开销会让你痛不欲生。我自己是在每 10 层清理一次，既平滑了碎片，又没有明显拖慢速度。

第二个意外是 PCIe 传输的“预取不生效”。我原本以为双缓冲写好之后，搬运和计算一定能重叠。但实测里第一版几乎没有任何加速。排查之后发现原因很简单：我在 CPU 端做权重切分时用了同步方式，导致 Stream A 的前半段和 CPU 计算串行在一起。准确说，是read_layer_by_offset这个函数把数据从文件读到内存时，用了常规的阻塞读写，CPU 要等数据完全读完才开始下一次操作。修复方式是用torch.load配合mmap=True，文件映射不需要把整块数据都读进物理内存，读的时候立刻能被 CUDA 拷贝。这一步改完，整个双缓冲才真正生效。

第三个意外是生成阶段比 prefill 阶段更容易死锁。我之前把注意力集中在 prefill 的逐层加载上，没有特别注意 decode 阶段。结果测试长文本生成时，CPU 内存里的 KV Cache 越来越大，decode 阶段每走一步都要从 CPU 拉一批 chunk 过来。如果 CPU 加载 chunk 和 GPU 算子计算之间没有一个合理的流水线，就会出现 GPU 等数据、CPU 等显存的互相等待局面。后来我把 KV chunk 的搬运也做成双缓冲，并且把 chunk size 调小到 128，情况立刻好转。

调完这三个问题之后，我构建的这套 airllm 分层推理方案终于稳定下来。它不是最快的推理方案，但它是目前我在单卡低显存环境下最顺手的方案：能跑、能续跑、还能在 8GB 显存里跑 13B 模型和一整套长上下文工具链。

最后再分享一个我自己觉得很实用的小技巧：第一次跑新模型的时候，不要直接上最大上下文，先用一个 500 token 的短序列跑通全流程，观察显存峰值和每层耗时。把这两项记下来，再去推算长序列下的资源占用，这样后面调参心里会有底得多。分层推理最忌讳的就是上来就拉满参数，然后被一连串 OOM 打回原型。

如果你想问“这套东西适不适合我”，我的判断很简单：只要你的目标是“在有限显存里把模型跑起来”，并且能接受推理速度不会太快，那 airllm 这条思路绝对值得折腾一遍。它改变的不只是内存占用，还有你对“推理能否运行”这个问题的底层认知。