AirLLM:逐层搬运突破显存墙,小显卡也能跑大模型
1. 为什么大模型跑不进你的小显卡显存墙的基本事实1.1 一个 7B 模型到底要吃多少显存先算一笔账。所谓 7B是指模型有大约 70 亿个参数市面上常见的开源大模型在保存权重时默认用 fp16/bf16 精度也就是每个参数占 2 个字节。光权重就是 70 亿 × 2 字节 ≈ 14GB。这还只是权重没算计算过程中产生的激活值、KV Cache以及框架自身的一堆额外开销。所以哪怕 7B 只是“入门规模”要把完整模型装进显存16GB 起步是跑不掉的13B 需要 26GB 以上70B 更是直接奔着 140GB 去。很多人的游戏本显卡是 6GB 到 8GB 显存连 7B 全精度都装不下。这不是什么丢人的事是数学决定的你手上这块卡物理上限就在那里。与其抱怨硬件不如换一种思路去看“显存不够”这件事。1.2 量化不是银弹显存不够时大家的第一反应通常是量化。把模型从 16bit 压到 8bit显存需求砍一半再压到 4bit7B 模型的权重只剩 3.5GB 左右看起来 8GB 的卡也能跑了。量化确实有效但它不是万能的。第一个问题模型越大4bit 量化后的体积依然可观。70B 模型量化到 4bit权重仍有 35GB 以上一张 24GB 的卡照样放不下。第二个问题量化本质是有损压缩。主流方案已经能把质量损失压得很低但在长文本推理、多步数学推理、工具调用这类任务里量化引入的微小误差会被一轮轮放大模型看起来在认真回答实际上已经在“编”了。你会发现单靠“把模型变小”这条路走到 70B 这个量级基本就堵死了。1.3 AirLLM 的思路不是把模型变小而是让显存只装一层AirLLM 换了一个完全不同思路我不跟你谈压缩我谈“轮换”。大模型推理时Transformer 的各层是一个接一个顺序执行的第 0 层算完才会算第 1 层所有层并不会同时参与计算。AirLLM 抓住这个特点把模型的全部权重放回磁盘或者 CPU 内存里用到哪一层的权重就把哪一层临时搬进显存算完立刻搬走接着搬下一层。这样一来显存里同时只需要容纳“一层的权重 激活值”6GB、甚至 4GB 显存去跑 70B 全精度模型在数学上就完全说得通了。代价当然也有每一轮前向都要来回搬权重推理速度会肉眼可见地变慢。这篇文章的后半部分我会把“到底慢多少、怎么让它快一点、什么时候不该用它”这些问题逐个讲清楚。2. AirLLM 的工作原理按层搬运的流水线2.1 Transformer 天然就是一串顺序执行的层要理解 AirLLM先得理解 Transformer 的结构。一个标准的开源大模型主体由几十层结构相同的 Transformer 层堆叠而成。每层内部大致包含两块注意力部分QKV 投影、注意力计算、输出投影和前馈部分通常两到三个线性变换加激活函数外面再包上 LayerNorm 和残差连接。关键点在于这些层之间存在严格的数据依赖第 0 层的输出就是第 1 层的输入第 1 层的输出又是第 2 层的输入链条式往下走。也就是说推理过程天然是“按层推进”的某一时刻真正在 GPU 上参与计算的只有当前这一层。AirLLM 把这种天生的顺序性利用到了极致既然同一时间只需要一层那我为什么要把几十层全部塞在显存里等着落灰2.2 一次前向到底做了什么AirLLM 的一次前向推理实际执行流程是这样的输入 token 在 CPU 侧准备好先把 embedding 表的权重临时搬上 GPU把 token 序列转成向量序列算完就把 embedding 权重释放掉。把第 0 层 Transformer 的全部权重从磁盘或内存搬到 GPU。在 GPU 上完成第 0 层的前向计算把输出的 hidden state 以及该层需要的 KV Cache 转回 CPU 侧管理。立刻释放 GPU 上第 0 层的权重。重复第 2 到第 4 步逐层处理完所有 Transformer 层。最后临时加载末层 LayerNorm 和 LM Head把 hidden state 映射到词表上得到 logits采样出下一个 token。生成下一个 token 时把整条流水线再跑一遍只是输入序列变长了一点。这个方案对 GPU 显存极其友好但对“搬运”这件事提出了苛刻要求。每一步都在磁盘、内存、显存之间来回倒腾整个系统的瓶颈从“算力不够”变成了“搬得不够快”。2.3 为什么慢到底慢在哪里慢的核心原因只有一个搬运带宽远低于计算带宽。一块中端显卡的显存带宽能做到每秒几百 GB但把权重从磁盘读进内存再通过 PCIe 总线拷进显存实际能跑到的速度也就是每秒几 GB 的量级NVMe 快一点机械硬盘慢到让人怀疑人生。算一笔更直观的账7B 模型通常有 32 层每层权重大约 400 到 500MB。如果从 NVMe 直读按 2.5GB/s 算读一层要 200ms 左右再算上 PCIe 拷贝和框架调度开销一层奔着 300ms 去32 层就是接近 10 秒。所以全精度 7B 模型用 AirLLM 跑实测速度通常就在 0.5 到 1.5 token/s 的水平。这不是 AirLLM 的实现有问题是物理瓶颈摆在那。好在这个瓶颈有优化空间。一是把权重优先缓存进 CPU 内存内存到显存的拷贝比磁盘直读快好几倍二是做流水线预取算第 N 层的同时预先把第 N1 层搬上来把搬运时间和计算时间部分重叠。AirLLM 这类经过打磨的 offload 方案都会在这些细节上下功夫。2.4 为什么它能跑峰值显存的账很多人第一次听说 AirLLM 会觉得很玄学其实把峰值显存算一遍就明白了。以 7B 模型为例32 层结构下单独一层的权重大概是注意力四张矩阵加前馈两三张矩阵加起来约 500MB激活值在 batch 为 1、序列长度 1024 的情况下也就几十 MB 到一百多 MBKV Cache 按层计算每层很小32 层全部加在一起在短序列下也就几百 MB。也就是说满打满算4GB 显存就能转起来。70B 模型同理虽然总权重高达 140GB但它有 80 层每层权重摊下来也就 1.6GB 到 1.8GB。只要显存能装下一层权重加激活值和必要的 KV Cache整个模型就能跑。AirLLM 的“反直觉”之处就在这里它不问模型总共有多少参数只问你的显存能不能装得下“最重的那一层”。3. 从零跑通 AirLLM环境、安装、第一个 Demo3.1 环境准备一张检查清单AirLLM 的目标平台是有一张普通显卡的 Linux 环境如果你用的是 Windows强烈建议装 WSL2 再继续纯 Windows 底下会遇到各种路径和动态库问题没必要硬刚。动手之前先对照下面这张清单确认硬件项目最低要求推荐配置操作系统Linux 或 WSL2Ubuntu 20.04 以上GPU 显存4GB6GB 到 8GB 最舒服磁盘模型体积 10%NVMe 固态空间越大越好内存8GB16GB 以上能缓存更多权重Python3.93.10 或 3.11 的独立虚拟环境磁盘这块特别提醒一下7B 全精度模型下载下来要 14GB 左右70B 要 140GB 左右千万别在只剩 20GB 的盘上硬跑。内存也别省AirLLM 会把用过的权重尽量留在内存里内存越大第二次跑的时候越接近“直接从内存搬”体验完全不一样。3.2 安装两条命令的事先建一个干净的虚拟环境然后用 pip 安装python -m venv airllm-env source airllm-env/bin/activate pip install --upgrade pip pip install airllmpip 会自动把 PyTorch、transformers 这些依赖装好。如果后续想配合量化使用再补一个pip install bitsandbytes这里不建议在系统 Python 里直接装。我吃过一次亏系统环境里之前装过一个老版本 transformers跟 AirLLM 依赖的版本冲突报错报得不知所云。虚拟环境隔离之后整个世界都清净了。模型默认下载到用户目录下的 .cache 里想存到别的盘就提前设置环境变量例如export HF_HOME/data/hf_cache3.3 跑起第一个 Demo安装完直接上代码AirLLM 的 API 风格和 transformers 很像用起来没有太多学习成本import time from airllm import AutoModel MAX_SEQ_LEN 1024 MODEL_ID qwen2.5-7b # 具体写法以项目支持列表为准 model AutoModel.from_pretrained(MODEL_ID, max_seq_lenMAX_SEQ_LEN) prompt 请用三句话解释什么是梯度下降。 inputs model.tokenizer(prompt, return_tensorspt) start time.time() output model.generate( inputs.input_ids, max_new_tokens64, temperature0.7, top_p0.9, ) elapsed time.time() - start gen_len output.shape[1] - inputs.input_ids.shape[1] print(f生成 {gen_len} 个 token耗时 {elapsed:.1f} 秒) print(model.tokenizer.decode(output[0], skip_special_tokensTrue))第一次运行会自动把模型权重从平台拉下来7B 模型大概要下载十几 GB取决于你的网速。跑通后你会发现 AirLLM 的模型名写法和 transformers 不完全一样每个开源模型系列在项目里都有自己的映射名动手之前先看一眼支持列表别拿一个没适配的模型硬试。3.4 怎么算才是真实速度衡量这类方案的真实速度不要用“首 token 延迟”来算。第一次前向会把输入序列的所有层都跑一遍属于 prefill 阶段耗时天然偏长生成每个新 token 时又要重新逐层搬运全部权重。所以正确的测速方式是用“总耗时 ÷ 生成的 token 数”来算单 token 延迟并且至少连续跑两轮第二轮再记录数据避开磁盘冷缓存和显存尚未预热的影响。结合常见硬件组合我整理了一个可参考的量级表注意这些数字只代表量级不是精确承诺硬件组合模型规模实测速度量级4GB 显存 NVMe7B0.4 到 1.0 token/s6GB 显存 NVMe7B0.8 到 1.5 token/s8GB 显存 NVMe13B0.3 到 0.6 token/s8GB 显存 机械硬盘13B0.1 到 0.2 token/s8GB 显存 NVMe70B0.05 到 0.15 token/s看到这个速度你大概能理解为什么我说“能跑”和“好用”是两回事。AirLLM 解决的是“有没有”的问题不是“快不快”的问题。4. 踩坑实录我实际遇到的四个问题4.1 第一次 OOM居然不是权重撑爆的我一开始理所当然地以为AirLLM 会在显存上非常节制结果第一轮测试就遇到了 CUDA out of memory。排查了很久才意识到炸显存的根本不是层权重而是激活值和 KV Cache。默认配置下输入序列一旦拉到几千 token中间每一层都要在显存里留一份 hidden state 和注意力中间结果再加上 KV Cache 随序列线性增长6GB 显存很容易被这条隐形通道挤爆。解决办法很简单把 MAX_SEQ_LEN 调小。6GB 显存跑 7B 模型序列长度控制在 1024 以内比较稳8GB 显存可以放到 2048 左右。另外跑之前用 nvidia-smi 看一眼有没有别的进程占着显存我遇到过一次浏览器硬件加速吃掉了 1GB 显存白排查了半天。4.2 硬盘空间被悄悄吃光AirLLM 有个很隐蔽的坑模型下载缓存默认放在系统盘用户目录下而下载过程远比想象中占空间。70B 全精度模型下载要 140GB下载中途还会留临时分片文件实际占用比最终体积再多出一截。我有一次在剩余空间 30GB 的盘上尝试下载 14B 模型跑了半小时后直接报 No space left on device最麻烦的是缓存文件已经写了一半还得手动清理。从那以后我养成了两个习惯第一动手之前用 df -h 确认磁盘空间第二提前把 HF_HOME 指到大容量盘上。这两个动作三十秒就能做完能省下大把时间。4.3 慢到怀疑人生时的排查顺序如果你跑出来的速度比预期慢一个数量级先别急着骂实现按这个顺序排查确认权重是不是真的在 NVMe 上。机械硬盘跑 AirLLM 就是灾难13B 模型直接掉到 0.1 token/s 以下这不是 bug是磁盘带宽撑不住。确认系统内存没有被塞满。内存不足时Linux 会开始用 swap磁盘既要读模型权重又要做内存换页整个系统像陷在泥潭里。确认显卡没有因为省电策略降频。笔记本上尤其常见插电、把电源模式切到性能档能明显感觉到速度变化。确认跑的不是第一轮。第一轮冷缓存要把权重从磁盘全量读一遍后面几轮如果内存足够速度会快不少。我之前在笔记本上排查了整整一个晚上最后发现是电源模式被切到节能显卡一直以一半的频率在跑。这种低级错误写代码的人最容易忽略。4.4 与 transformers 生态的隐性冲突AirLLM 依赖的 transformers 版本和你的其他项目容易互相踩脚。常见症状是单独跑 AirLLM 一切正常一进某个已存在的项目环境就报算子不兼容或者 tokenizer 加载失败。我后来就固定在一个专属虚拟环境里跑 AirLLM绝不和其他项目混装。如果遇到模型不支持的报错先看看是不是 transformers 版本太新或太老再考虑模型本身是否在适配列表里。还有一点AirLLM 不负责微调它的定位非常纯粹就是推理。如果想在这个环境里同时加载 LoRA 或者做训练趁早换一套方案别在 offload 推理框架里硬凑。5. AirLLM 和其他方案怎么选一张表讲清楚5.1 主流方案对比我见过太多人一上来就问“AirLLM 和某某哪个好”这个问题其实没有标准答案取决于你的显卡、任务和耐心。我把几个常见路线放一张表里方案最小显存7B 速度量级精度上手难度适合场景AirLLM4GB0.5 到 1.5 token/s无损中等小卡跑大模型、精度敏感llama.cpp 系含 Ollama5GB 到 8GB5 到 20 token/s有损低日常对话、快速原型Accelerate 磁盘卸载1GB 甚至更低0.1 到 0.5 token/s无损低已有 HF 生态、偶尔跑一次专业推理引擎24GB 起几十到几百 token/s视配置高生产环境、高并发llama.cpp 系通过量化加多线程优化在 7B 这个档位明显更快日常聊天体验好得多。但到 70B 这个量级它要么需要超大内存做 CPU 推理要么需要多张大显存卡做 offload普通用户的硬件依然撑不住。AirLLM 的价值恰恰在“小显存 超大模型”这个交叉点上。5.2 什么场景真的该用 AirLLM判断标准很简单如果你的需求同时满足“显卡只有几 GB”和“模型想上十几 B 甚至几十 B”AirLLM 几乎是唯一兼顾无损精度和可实施性的选项。具体来说这几类人很适合手里只有老显卡、学生机、游戏本但想实际跑一跑大模型的开发者对量化精度损失有顾虑想在有限硬件上先验证模型真实输出的研究者数据敏感强需求本地推理能接受离线等几分钟出结果想给学生或者同事直观演示 Transformer 逐层执行过程的教学场景。除了实际跑模型AirLLM 还有一个隐藏价值它把“层”这个概念变成了一个可观测的单元跑一遍你就能清楚看到模型一层一层在做什么对理解大模型内部机制非常有帮助。5.3 什么时候别用它反过来如果你要的是响应及时的对话体验或者要支撑线上业务请求AirLLM 现阶段确实不适合。1 token/s 意味着说一句话要等一分钟用户早就关页面了。这类需求要么老老实实上大显存卡要么接受量化方案换取速度。另外如果你手里已经有 24GB 以上的显卡也别迷信 AirLLM直接在显存里完整加载模型任何 offload 都会让速度变慢没有例外。6. 进阶优化与实用建议6.1 把 CPU 内存当成第二级缓存AirLLM 这类 offload 架构的命门在磁盘带宽而 CPU 内存就是它最好的缓冲垫。原理很简单磁盘到显存要走两段路磁盘到内存、内存到显存如果权重已经躺在内存里就是一步直通。内存带宽和 PCIe 拷贝速度比 NVMe 快好几倍系统内存够大的时候AirLLM 的运行速度会明显上一个台阶。实操上你几乎不需要额外配置只要内存足够系统会把热数据留在 page cache 或进程内存里。要注意的是别同时开一堆吃内存的服务把内存挤没了系统开始频繁 swap速度还不如直接从磁盘读。我一般跑 13B 模型时会顺手关掉浏览器的大标签页给模型腾出十六七 GB 的内存空间。6.2 序列长度和生成长度要抠着用用 AirLLM 一定要养成“抠长度”的习惯。输入序列越长prefill 阶段逐层跑的耗时越长max_new_tokens 越大后面重复逐层搬权重的次数越多。这两个参数直接决定总时长所以每次跑任务之前先想清楚是不是真的需要那么长的上下文。我常用的策略是验证提示词效果时max_new_tokens 只给 32 到 64先快速看个开头确认思路对了再跑完整长度。这就像写代码先跑单元测试再跑全量回归能省下大量等待时间。6.3 多卡场景别硬套 AirLLMAirLLM 的设计目标非常明确就是单机单卡小显存。如果你手上有两张小卡不要指望它自动帮你把层均摊到两张卡上这不是它的强项。更实际的做法是让两张卡各跑各的进程分别处理不同的任务互不干扰总吞吐反而更高。想真正吃满多卡资源做高吞吐推理还是得依赖专门为多卡设计的方案。AirLLM 解决的是“能不能跑”的问题不是“跑得够不够快”的问题两者边界一定要拎清楚。6.4 最后说点个人体会用 AirLLM 跑了几轮大模型之后我对“显存不够”这件事的看法彻底变了。早些年我也觉得没有好显卡就不配玩大模型实际上很多问题并不是算力不够而是“一次性把所有东西都装进显存”这个惯性思维在作祟。AirLLM 提醒我推理本身是流水线不是堆栈逐层搬运的代价是速度换来的是门槛的极大降低。如果你现在正攥着一张 6GB 老显卡犹豫要不要折腾我的建议是不要犹豫直接装个环境跑个 7B 模型试试。第一次看到生成文字在终端里一个个蹦出来的时候你会觉得那些等待都值了。等验证完思路再根据实际需求决定要不要上量化、上专业推理引擎那时候你已经知道每一步消耗在哪里不会再做无头苍蝇式的试错了。