airllm 实践:分层加载如何让 8GB 显存跑 70B 大模型
前几天有个朋友问我一台只有 8GB 显存的老笔记本能不能跑 70B 级的大模型换在以前我会直接说别折腾了老老实实去用云端接口。但现在我会认真告诉他可以试试 airllm 这种分层加载的思路。airllm 的核心不是黑魔法而是把模型一层一层拆开算完一层就换下一层用“时间换空间”让显存需求从整个模型降成单层模型加中间激活。对显存不大、又想在本地跑开源大模型的人来说这是目前最务实的一条路。这篇文章不打算复读文档我想结合自己实际踩坑的过程把 airllm 的思路、部署参数、调优手段一次说清楚。1. airllm 到底解决了什么问题1.1 本地推理的第一瓶颈为什么是显存先说一个最简单的估算方法。一个 70B 参数的稠密模型如果用 2 字节的精度保存权重那么光权重就得占 140GB 左右。这是什么概念很多人家里的整台电脑内存都没这么大更不用说显存了。哪怕把精度压到 4bit这个数也还要 35GB 上下普通桌面显卡根本放不下。还有人会想显存不够不是还有内存吗问题在于传统推理流程里模型权重必须常驻显存因为每个 token 的前向计算都要遍历所有层。内存里的数据要反复拷贝到显存带宽和延迟都扛不住。主流深度学习框架默认会把整个模型搬到计算设备上于是显存墙就出现了。除了权重还有两个隐形成本容易被忽略一是每一层前向计算时的中间激活值二是自回归生成时不断增长的 KV cache。模型越大、上下文越长这两个东西就越占显存。所以本地跑大模型的瓶颈往往不是算力而是容量。airllm 选择绕开这道墙干脆不要求模型完整地待在显存里。1.2 airllm 的思路把“常驻”改成“流转”airllm 的核心思路非常朴素Transformer 的层本来就是一层一层按顺序执行的那我为什么一定要把所有层同时放在显存里只把当前需要计算的那一层放到显存算完这一层就释放掉再加载下一层完全可以做到同样结果。这个改动把显存需求从“整个模型”降到了“单层权重 当前激活 KV cache”。一个 70B 模型哪怕有 80 层拆分后单层权重也就 1GB 左右普通显卡能轻松装下。代价是每一层都要经历一次“从存储到显存”的搬运推理速度会慢不少。我用一个生活类比辅助理解以前看一本厚字典你得把整本放在桌上翻哪页都方便但桌子必须足够大。airllm 的做法是只留一页在桌上看完再翻下一页。桌子不用大但翻页会花时间。所以 airllm 的核心取舍很清楚——用推理延迟换硬件容量。1.3 和量化不是二选一而是可以叠加很多人一听“低显存跑大模型”第一反应是量化。量化确实能压缩模型体积但它是通过降低数值精度来换空间的通常会让推理精度有一些损失。airllm 则完全不需要量化原始权重它可以加载完整精度的模型在推理结果上更接近模型本来面貌。更关键的是这两种方案并不冲突。airllm 解决的是“总量放不下”的问题量化解决的是“单块数据太大、搬运太慢”的问题。把层拆分到普通显存后如果单层仍然太大或者内存带宽吃紧可以再对每层做量化进一步减小体积。airllm 加 4bit 量化是我在实际部署中觉得最稳的组合精度损失相对可控速度又比纯全精度快不少。下面这张表可以参考方案核心手段优点主要代价适合场景全量常驻显存大显存硬扛速度快、实现简单显存成本极高有服务器级显卡普通量化降低权重精度显存需求明显下降精度损失、需要校准显存中等、追求速度airllm 分层加载按层卸载/加载显存门槛极低可保持高精度层搬运导致延迟普通个人电脑airllm 量化两者结合显存门槛低搬运量进一步降低精度和延迟需要平衡预算有限的离线推理2. airllm 的工作原理与关键设计2.1 一次前向推理的层流转过程如果抛开工程细节airllm 的前向过程可以简化成下面这段伪代码# 伪代码airllm 分层加载核心逻辑 def forward_one_token(input_ids, model_info): hidden embed(input_ids) for layer_id in range(model_info.num_layers): layer load_layer(model_info, layer_id) layer.to(compute_device) hidden, kv_cache layer(hidden, kv_cache) layer.to(storage_device) release_intermediate() return lm_head(hidden)每次只加载一层运算结束后立刻把这一层的参数挪回存储设备同时释放中间缓存。hidden state 是要保留的它会在层与层之间传递KV cache 也会不断累积因为它记录了前面的注意力信息后续层还要用。需要注意一个细节虽然层权重可以换来换去但 hidden state 和 KV cache 一般还是会留在显存里。所以显存占用不是“单层权重”这么简单还要加上这两部分。如果上下文太长KV cache 涨起来最后瓶颈就不在权重的层加载而在 KV cache 的持续累积。2.2 为什么按“层”粒度切分最划算把模型拆开有很多切法可以按张量切可以按模块切也可以按层切。airllm 选择按层切我认为有三个很实际的原因。第一层是 Transformer 的天然执行边界。每一层的输入输出结构一致切分后不需要改动注意力实现也不需要改前向计算逻辑只要按顺序调度就行通用性非常高。第二按层切可以在计算完当前层后立刻释放这一层的中间变量显存复用效率最高。第三层内部虽然还可以继续拆比如把注意力部分和 MLP 部分分开卸载但那样会增加大量调度节点工程复杂度高收益却不一定明显。有人会问为什么不按“几个 block”切也可以但块越大显存峰值越高能跑的模型就越小。层粒度是在工程复杂度和显存收益之间最舒服的平衡点。设备好一点的可以一次加载多层相当于把“层粒度”扩大成“块粒度”灵活性反而更高。2.3 计算设备、内存设备、存储设备的分工airllm 本质上是在构建一套非常轻量的“存储分级调度系统”。最上层是显存容量最小但速度快中间层是内存容量大很多但带宽受限最底层是磁盘容量最大但速度最慢。推理时权重主要在磁盘和显存之间流转内存往往扮演临时缓存角色。这里有一个容易被忽略的点如果模型权重以普通文件方式读取每次加载层都要从头读盘哪怕只读 1GB 数据也会产生大量随机 IO。airllm 的做法通常是把权重文件以内存映射方式打开配合操作系统按需读取让层数据在真正被访问时才进入内存。这样预取和缓存都交给系统调度工程上省了很多事。所以airllm 对磁盘性能和内存带宽很敏感。同样的模型配 NVMe 固态和普通 SATA 固态首 Token 延迟能差好几倍内存是单通道还是双通道也会影响层搬运效率。这一点我后面实操部分会专门说。3. 实操在显存有限的设备上部署 airllm3.1 环境准备与需要避开的坑先说我手头这台测试设备的情况8GB 显存、32GB 内存系统盘是普通 NVMe 固态。目标是一个约 70B 参数的开源稠密模型权重使用 safetensors 分片格式存放在本地目录。准备环境时最该注意的就是不要一上来就调用框架默认的“自动设备映射”。很多深度学习框架默认会尽可能把模型塞进显存塞不下了才放内存。这个行为和 airllm 的按层调度是冲突的如果开启了自动设备映射airllm 的层缓存策略就很难生效甚至会出现显存爆掉的情况。其次权重文件尽量用分片格式不要把整个模型打包成一个巨型二进制文件。分片的好处是加载单层时可以精确定位到对应分片不必把整个文件读一遍。如果只有单文件权重我可以先把权重按层重新切分或者用内存映射方式打开但效果都不如直接用分片格式来得干净。最后要确认存储空间。70B 模型原始权重按 2 字节计算大概 140GB再加上预取缓存和临时文件我建议至少留出 200GB 可用空间。空间不够的时候程序可能不是直接报错而是反复读取导致卡死排查起来很费劲。3.2 核心参数照着改就能少踩一半坑配置可以整理成下面这样我用它作为基准model: path: /data/models/70b-safetensors precision: bf16 max_length: 2048 engine: offload_dir: /fast_ssd/airllm_cache prefetch_layers: 2 keep_embed_in_memory: true keep_lm_head_in_memory: true max_loaded_layers: 1逐个说下关键项。precision优先用 bf16。bf16 的指数范围和 float32 一样但只占 2 字节数值稳定性比普通 float16 好尤其适合大模型。如果你的设备不支持 bf16可以退回 float16但要注意生成时可能出现 NaN 或乱码。offload_dir要放在最快的磁盘上最好是 NVMe 固态。它是层数据的临时落脚点如果放机械硬盘每次层交换都能让你等到失去耐心。prefetch_layers控制预取层数我习惯设 2。意思是计算当前层的时候提前把后面两层的数据读到缓存减少等待。设太大反而会让缓存占用过高甚至因内存不足而触发系统回收。keep_embed_in_memory和keep_lm_head_in_memory很值得打开。embedding 和 lm_head 虽然参与计算但参数总量远小于中间层常驻内存后能明显减少首 Token 的冷启动延迟。max_loaded_layers在基准配置里设成 1这是最保守的做法。如果显存还有余量可以适当调大比如设成 2 或 4速度会提升但显存占用也会线性增长。3.3 首跑实测一个 70B 级模型的参考数据我用上面这套配置跑了一个 70B 级模型prompt 长度控制在 512 token生成 200 个 token。首 Token 延迟在 35 秒左右后续每个 token 大约 1.8 到 2.5 秒峰值显存稳定在 6.5GB 附近。这个速度说不上快但对于本地跑 70B 模型来说已经完全可用了。我又试了把max_loaded_layers调成 2首 Token 延迟降到了 22 秒单 token 速度提升到 1.4 秒左右但峰值显存也涨到了 8GB 临界点。再往上调就会开始吃紧所以我最终在“能稳定运行”的前提下还是保留了单层加载。如果只生成短文本比如 50 个 token 以内这个方案体验还可以如果用来做长对话或者长文档生成每秒 2 token 的速度确实会让人着急。所以 airllm 更适合离线任务比如批量总结、代码审查、数据标注辅助而不是实时聊天。4. 性能调优让 airllm 从“能跑”到“好用”4.1 先算好显存、内存、缓存盘之间怎么分账性能调优的第一步不是改参数而是算清楚容量账。显存里要放三样东西当前层权重、hidden state、KV cache。单层权重大约是 0.7 到 1.5GBhidden state 和 batch size 成正比KV cache 则和上下文长度成正比。我在部署时先按“单层权重 2GB 保底”估算显存然后再看剩余空间能不能容纳 KV cache。内存这边至少要有“模型总量”的一部分作为映射页缓存。如果你开了预取内存占用还会更高。我的经验是物理内存至少要是显存的好几倍8GB 显存的设备推荐 32GB 内存16GB 显存的设备推荐 64GB 内存。内存不足时系统会开始频繁交换内存页到磁盘整个推理会慢到几乎不可用。缓存盘其实就是 offload 目录。它不需要一次性装下整个模型但至少要能容纳经常被换入换出的那部分层数据。速度优先级是NVMe 固态 SATA 固态 机械硬盘。如果你只有机械硬盘建议把 offload 目录放到内存盘上但那样会牺牲内存空间需要自己权衡。4.2 让“搬层”不再拖后腿的两个方向第一个方向是减少搬运次数。最直接的办法是加大max_loaded_layers一次多放几层。但这不是无脑调大的要观察显存余量。另一种思路是把层内计算分成“必须驻留”和“可以临时加载”两部分。比如注意力层的投影矩阵可以先驻留MLP 中的扩展矩阵再按需加载这算更细粒度的优化工程量会大不少。第二个方向是隐藏搬运延迟。prefetch_layers就是在计算当前层时提前加载后面的层让搬运和计算重叠起来。之前我设 2 已经能覆盖大部分延迟如果你的存储比较慢可以尝试调高到 4 或 5但要注意内存占用。还有一个容易忽略的优化把 embedding 和 lm_head 常驻显存。这两个模块虽然也要参与计算但参数不多不会挤占太多显存却能省去每轮生成里最频繁的参数切换。我实测过这个小改动首 Token 延迟能降低 15% 左右。4.3 长文本生成时的显存防洪KV cache 是 airllm 模式下最容易爆显存的部分。我跑 512 token 上下文时KV cache 占用的显存不超过 2GB但把上下文拉到 2048KV cache 可能直接翻两倍以上。由于 KV cache 不能像权重一样卸载只能留在显存里所以它成了长期生成的硬约束。最稳妥的方法是把max_length设短比如 1024 或 2048让模型在生成过程中主动截断超长提示。如果业务真的需要长文本可以在输入侧做摘要压缩把历史对话压缩成更短的概括再拼接到新 prompt 里。还有一种做法是把旧 token 的 KV 状态转回 CPU但这样做会引入大量内存拷贝速度反而更难接受。另一个建议是控制 batch size。airllm 的层加载模式下batch size 变大hidden state 和 KV cache 都会成倍增长但速度并不会线性提升。离线任务尽量用 batch size 1在线任务如果必须批处理就要同步调小max_loaded_layers给 KV cache 留出余量。5. 常见问题与排查技巧实录5.1 首 Token 延迟特别高airllm 第一次推理时冷启动延迟高是正常的因为所有层都要从头读一遍。但如果首 Token 延迟到了几分钟甚至十来分钟大概率不是算力问题而是存储速度问题。我遇到过两种情况。第一种是 offload 目录放在机械硬盘上层切换时随机读取速度只有几十 MB/s一个 1GB 的单层要读十几秒。第二种是系统内存不足层数据被放到交换文件里内存和磁盘之间反复倒腾。排查时我一般先看磁盘利用率再看内存余量。如果磁盘长时间满负荷优先换 NVMe 或调整prefetch_layers。还可以做一次“预热”加载模型后先用一段很短的 prompt 跑一次把 embedding、lm_head 和一部分常用层页面提前读入内存正式推理时首 Token 会快很多。5.2 跑到一半 OOM 或者直接被系统杀掉如果程序直接显示“显存不足”通常是 KV cache 增长超过预期。先把max_length调小或者降低max_loaded_layers。如果进程被系统杀掉没有像样的报错那就更可能是内存被耗尽而不是显存。此时要检查打开的其他程序尤其浏览器和编译工具它们会吃掉大量内存。还有一个隐蔽的坑多个进程同时使用同一个 offload 目录。比如你同时起了两个推理脚本它们会对同一批缓存文件做读写轻则互相等待重则导致缓存文件损坏。我后来规定每个进程必须使用独立的 offload 目录问题立刻消失。如果内存还有富余也可以考虑把offload_dir放到内存盘上用内存当缓存盘。这样显存、内存、暂存三层都很快代价是可用内存变少。适合模型不太大、但存储很慢的场景。5.3 生成结果不对劲乱码、重复、逻辑断裂airllm 的分层计算本身不应该改变模型输出结果异常往往是数值精度或状态管理出了问题。最常见的是精度类型混用某一层用 float16另一层用 bf16虽然看起来差别不大但在大模型里会累积成明显误差。我通常会把全模型精度固定成 bf16只在确实不支持 bf16 的设备上用 float16。还有一个问题出在 KV cache 的传递上。如果某层的 KV 被错误地放到了 CPU 或其他设备后续层读取时会出现隐性拷贝轻则速度变慢重则产生随机错误。排查时可以把 KV cache 的 device 和 hidden state 的 device 打出来确认它们保持一致。需要说明的是airllm 在 CPU 和 GPU 之间反复搬运层权重受驱动和运行库影响偶尔会出现层加载不完整的情况。如果发现某次生成从某个 token 开始突然乱码可以尝试把offload_dir清空强制重建缓存通常能解决。5.4 并发访问同一模型文件的隐患airllm 按层加载模型时会频繁读取模型目录里的分片文件。如果你在多个脚本里同时加载同一个模型路径操作系统会面临大量并行读操作虽然数据本身不会写坏但磁盘 IO 会被拖垮所有实例一起变慢。我的做法是要么串行跑多个任务要么复制模型到不同磁盘路径让每个实例有自己的模型副本。如果你只改了显存配置但没换路径两个进程的性能会互相拖累排查起来还容易误判成 airllm 本身太慢。6. 个人体会与后续扩展6.1 airllm 适合什么场景、不适合什么场景我在实际使用中发现airllm 最适合的是三类任务离线批量处理、隐私敏感的数据分析、以及教学演示。数据不需要离开本机模型越大越能体现价值。你要处理的文本量大但不要求实时返回等几十秒完全可接受。它不适合实时聊天、高并发接口和超长文本生成。聊天对首 Token 延迟要求很高并发场景下层调度会互相抢资源超长文本则会被 KV cache 拖垮。如果这些是你的核心需求建议还是走正常的多卡推理或优化过的高吞吐服务。另外层数越多的模型airllm 的单 token 延迟往往越高因为每个 token 都要把所有层过一遍。所以“能用 airllm 跑 70B”不等于“跑得很舒服”它对模型的深度非常敏感。同样参数规模下层数少、宽度大的模型在 airllm 下反而更有优势。6.2 后续可以继续做的方向airllm 的模式其实还有很大扩展空间。比如结合投机解码用一个小模型先预测多个候选 token再用大模型去验证这样大模型可以一次验证多个 token相当于用更少的层遍历次数完成了相同长度的生成。另一个方向是更智能的预取。现在的预取顺序是固定的“下一层”但如果结合注意力权重或者历史访问模式完全可以预测未来几层哪些数据会被高频使用提前放到更快的缓存里。还可以做更细粒度的层内拆分。把一层的注意力部分和 MLP 部分拆开按需调度显存占用甚至可以降到单层参数的一半。不过这需要针对具体模型做适配通用性会差一些。最后如果能把 airllm 和简易的任务队列结合给离线任务做排队和进度管理它就能从“一个推理脚本”变成“本地大模型服务”。我自己跑 airllm 这段时间最大的体会是它把大模型的门槛从“买一张天价显卡”降到了“有一台内存足够大的普通电脑”。如果你也想在本地跑大模型却一直被显存卡着不妨先按这套思路试一遍。第一次看到几十亿参数的模型在自己的机器上吐出完整句子时那种感觉确实很值。