先问一句你手头是不是也有一张显存“刚好不够用”的显卡我最近在折腾本地跑 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 涨到 8192KV Cache 可能从几百 MB 涨到几个 GB。为了更直观我做了一张表格对比几种常见方案方案显存需求主要代价适合场景全量权重驻留显存模型完整权重 KV 激活显存必须够大高性能部署CPU Offload只把部分层放 GPU卡间 PCIe 传输频繁速度慢显存极小但内存充足4bit/8bit 量化权重大幅缩小精度损失部分算子兼容性差显存中等追求速度张量并行每张卡只放一部分权重需要多卡通信通信开销大多卡环境分层推理airllm只需单层权重 单层激活加载层权重有开销带宽成为瓶颈单卡 / 低显存可以看到airllm 走的是另一条路它不追求把模型“压缩”到显存里而是改变模型驻留显存的方式。原本需要同时放下几十层现在只需要同时放下一层。这个思路的代价也非常明确——计算的时候必须等权重从内存搬进显存性能和 PCIe 带宽强相关。所以在正式开工之前最好先建立一个清醒的预期airllm 不是灵丹妙药它不会让 8GB 卡跑出 4090 的速度但它的价值在于“跑得起来”和“跑得稳”。对于只需要离线推理、逐条处理的场景这套方案已经非常够用了。2. 逐层加载airllm 的核心思想airllm 的做法用一句话概括就是不要让 GPU 同时看到模型的所有层只给它看正在计算的那一层。Transformer 模型天然就有很好的分层结构。输入经过 embedding 之后会依次通过若干个 Decoder Layer每个 Layer 内部有 attention、MLP、LayerNorm、残差连接这些子结构。层与层之间是严格顺序依赖的第 i 层的输出是第 i1 层的输入。这种链条结构给了我们很大的操作空间。在传统推理里模型初始化时就把所有层的权重从磁盘加载到显存中。在 airllm 方案里权重初始存放在 CPU 内存或者磁盘映射区域。真正开始推理的时候我们维护一个循环先加载第 0 层的权重到 GPU把上一层的输出 hidden state 送进这一层计算算完得到的输出继续作为下一层的输入立刻把这一层的权重从显存中释放掉加载第 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 在搬第 i1 层权重的时候计算单元又是空闲的。效率极低。真正的 airllm 调度必须用双缓冲甚至多缓冲的思路。简单说就是让“权重搬运”和“层计算”重叠起来。具体做法是准备两个 CUDA StreamStream A 负责把第 i 层的权重从 CPU 拷贝到 GPUStream B 在权重到达后执行第 i 层的前向计算计算第 i 层的同时Stream A 已经可以开始预取第 i1 层权重。这里有个关键约束预取的下一层权重需要占用额外的显存。但前面已经说过单层权重在 8GB 卡的预算内是可以接受的。所以“多缓存一层”的代价并不大换来的是接近零等待的搬运时间收益非常明显。实测下来双缓冲能让端到端推理速度提升约 40% 到 60%瓶颈从“串行等待”变成“PCIe 带宽本身”。为了减小权重搬运的压力airllm 在文件组织上也做了针对性设计。模型权重不是一次性全部加载到内存而是按层切好每一层的参数用偏移量去定位读取。在实现层面我做了这样一个简单的布局表层序号参数字节数8bit 量化后文件偏移量embedding1.6GB / 2 0.8GB0x0000_0000layer 00.85GB0x1000_0000layer 10.85GB0x2000_0000.........这样每次加载某层时只需要根据表格里的 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_blockingTrue) # stream_b等权重到位后计算 current_stream.wait_stream(buffer_stream) hidden_states layer_forward(hidden_states, weight_gpu)千万别小看non_blockingTrue这行。如果省略它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个参数也就是 256KBFP16。当序列长度涨到 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 显存非常友好。我整理了一份不同序列长度下的显存影响表上下文长度KV Cache 总大小FP16每 chunk 占用256 token20480.5GB0.06GB40961GB0.06GB81922GB0.06GB327688GB0.06GB13107232GB0.06GB但注意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.43GB4bit 后约 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配合mmapTrue文件映射不需要把整块数据都读进物理内存读的时候立刻能被 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 这条思路绝对值得折腾一遍。它改变的不只是内存占用还有你对“推理能否运行”这个问题的底层认知。