1. 单卡跑大模型这件事到底卡在哪第一次听说 airllm 这个概念的时候我正在一台只有单张消费级显卡的机器上折腾一个 70B 参数级别的模型。当时显存只有 24GB模型权重按 FP16 算下来需要接近 140GB差了将近六倍。直觉告诉我这事不可能但 airllm 这个思路恰恰就是冲着不可能去的——它要解决的核心问题只有一个当显存装不下整个模型时怎么还能把推理跑起来而且不用量化到面目全非。airllm 本质上是一套**分层推理layer-by-layer inference**的实现方案。它的核心主张是大模型的推理过程天然是逐层串行的第 1 层的输出喂给第 2 层第 2 层喂给第 3 层直到最后一层出结果。既然任意时刻只有一层在真正参与计算那为什么非要把所有层的权重同时塞进显存把权重放在硬盘上用到哪一层就把哪一层加载进显存算完就释放显存占用就能从整个模型降到单层模型。这个思路听起来简单但真正落地要解决一堆工程问题硬盘读取速度够不够、层加载和计算怎么重叠、KV Cache 怎么管理、精度怎么保持。airllm 的价值就在于它把这些脏活累活都封装好了让你用几行代码就能在一个显存远小于模型体积的机器上跑起推理。它适合谁适合那些手里只有单卡、又不想把模型量化到 4bit 以下、还想保持 FP16 精度跑大模型的人。如果你追求的是极致吞吐量或者生产级并发那它不是你的菜但如果你是想在本地验证一个大模型的效果、做做实验、跑跑 demo这套方案能帮你省下一大笔硬件钱。我后面会从设计思路、核心机制、实操步骤、踩坑经验几个维度把 airllm 这套东西拆开讲清楚。不管你是刚接触大模型推理的新手还是已经玩过 vLLM、llama.cpp 的老手都能从中找到能直接抄作业的部分。2. airllm 的整体设计与思路拆解2.1 为什么是分层加载而不是量化压缩面对显存不足业界主流有两条路一条是量化把 FP16 压成 INT8、INT4 甚至更低用精度换空间另一条是分层加载权重精度不动靠时间换空间。airllm 选了后者这个选择背后有很明确的取舍逻辑。量化的好处是显存占用直接砍半甚至砍到四分之一推理速度还更快因为低精度计算本身就更省算力。但量化有个绕不开的问题精度损失。4bit 量化在大多数任务上表现还行但在一些对数值敏感的任务上——比如长链推理、数学计算、代码生成——质量下降会很明显。而且量化本身需要校准数据不同模型、不同任务的最优量化策略不一样调起来很烦。分层加载则完全避开了精度问题。权重还是 FP16计算还是 FP16输出结果和全量加载理论上完全一致实际会有极小的数值差异后面会讲。代价是速度每一层都要从硬盘读进显存硬盘 IO 成了瓶颈。但 airllm 通过预读取prefetch和计算重叠把这个代价压到了可接受的范围。提示如果你的场景对延迟极度敏感比如在线服务要求 P99 在几百毫秒内那分层加载基本不适合。但如果是离线批处理、本地实验、单次生成这个速度完全可以接受。2.2 分层推理的核心假设与边界airllm 能成立依赖一个关键假设Transformer 的推理是严格逐层串行的。第 N 层的计算只依赖第 N-1 层的输出不依赖任何未来层的信息。这个假设对标准 Transformer 解码器成立所以 airllm 对 LLaMA、Qwen、ChatGLM 这类自回归模型都能用。但有几类模型它会失效或者效果打折。第一类是MoE混合专家模型因为 MoE 层内部有路由机制不同 token 激活不同专家层的加载策略会复杂很多。第二类是带跨层连接的模型如果某一层需要引用前面好几层的输出那算完就释放的策略就不成立了。第三类是编码器-解码器结构编码器部分可以分层但解码器每步都要访问编码器输出内存管理更麻烦。airllm 官方主要支持的是 LLaMA 系和类似的自回归解码器架构。你在选模型之前最好先确认一下它的结构是不是纯解码器、有没有跨层引用。这个判断不难看一眼模型的config.json里architectures字段和num_hidden_layers就能大致判断。2.3 时间换空间的账怎么算我们来算一笔账看看分层加载到底划不划算。假设一个 70B 模型FP16 精度总权重大约 140GB共 80 层平均每层 1.75GB。一台机器有 24GB 显存硬盘是 NVMe SSD顺序读取速度约 3GB/s。全量加载需要 140GB 显存直接出局。分层加载的话任意时刻显存里只需要放 1 到 2 层当前层 预取层加上 KV Cache 和激活值显存占用大概在 5-8GB24GB 绰绰有余。时间上每层读取 1.75GB按 3GB/s 算约 0.58 秒。80 层就是 46 秒的纯读取时间。如果计算和读取能完全重叠那总时间约等于 max(读取总时间, 计算总时间)。70B 模型单层计算在单卡上大概几十毫秒80 层计算总时间可能只有几秒所以瓶颈完全在 IO总时间约 46 秒生成一个 token 序列的完整前向。等等这里要澄清一个容易误解的点分层加载的开销是每个 token 都要付一次的。因为自回归生成是逐 token 的每生成一个新 token都要把所有层重新走一遍。所以如果生成 100 个 token理论上要读 100 遍权重那就是 4600 秒一个多小时。这显然不能接受。airllm 的优化就在这里它把权重常驻在内存RAM里而不是每次都从硬盘读。硬盘只在第一次加载时读一遍之后每层权重都在内存中每生成一个 token 只需要把层从内存搬到显存。内存带宽比硬盘高一个数量级DDR4 大约 20-50GB/s这样每层搬运 1.75GB 只要 0.05 秒左右80 层 4 秒生成 100 个 token 约 400 秒。这个速度虽然还是慢但至少进入了能用的范畴。注意这就是为什么 airllm 对内存容量有要求。70B 模型 FP16 需要约 140GB 内存你的机器内存得够大。如果内存也不够那就只能退回到硬盘流式读取速度会慢到难以忍受。3. 核心机制与关键细节解析3.1 权重分片与层索引的建立airllm 第一步要做的是把原始模型权重按层切分并建立索引。原始模型通常是一个或几个大的.safetensors或.bin文件里面所有层的权重混在一起。airllm 会读取模型的层结构信息把每一层的权重单独提取出来存成独立的文件或者建立偏移索引。这个过程叫模型分片sharding。分片的好处是加载时可以精确地只读某一层不用把整个文件读进来再丢弃。如果模型本来就是分片的比如 HuggingFace 上很多大模型已经按层或按块分好了airllm 可以直接复用这些分片省去重新切分的步骤。分片时有个细节要注意层的边界要和模型结构严格对应。一个 Transformer 层通常包含自注意力self-attention和前馈网络FFN两部分每部分又有若干权重矩阵Q、K、V、O、gate、up、down 等。分片时要保证同一层的所有权重被归到一起否则加载时会漏掉或者多读。airllm 内部通过解析模型的state_dict键名来识别层归属键名里通常带有layers.{i}.这样的模式按这个模式分组就行。3.2 预读取与计算重叠的实现分层加载最大的性能杀手是串行等待读第 N 层 → 算第 N 层 → 读第 N1 层 → 算第 N1 层。如果读取和计算不重叠总时间就是读取时间加计算时间。airllm 通过**预读取prefetch**把这两者重叠起来。具体做法是在计算第 N 层的同时后台异步读取第 N1 层的权重到显存。这样当第 N 层算完时第 N1 层已经准备好了可以直接开始计算。理想情况下总时间约等于 max(总读取时间, 总计算时间)而不是两者之和。实现预读取需要用到异步数据传输。在 CUDA 里这通常通过cudaMemcpyAsync配合多个流stream来实现。一个流负责计算另一个流负责数据传输两者通过事件event同步。airllm 在底层封装了这些操作你不需要自己写 CUDA 代码但理解这个机制有助于你调优。提示预读取的深度可以调。预读 1 层是最基本的预读 2 层能进一步掩盖延迟但会多占显存。显存紧张时预读 1 层显存充裕时可以预读 2-3 层。这个参数在 airllm 的配置里通常叫prefetch_layers或类似名字。3.3 KV Cache 的管理策略自回归生成里KV Cache 是显存占用的另一个大头。每生成一个 token都要把当前 token 的 Key 和 Value 缓存下来供后续 token 的注意力计算使用。序列越长KV Cache 越大。在分层加载的场景下KV Cache 的管理有个特殊之处KV Cache 是跨层共享的。第 N 层的 KV Cache 在第 N 层计算时产生但后续所有层在计算注意力时都要用到前面所有层的 KV Cache 吗不是的。每一层有自己的 KV Cache层与层之间的 KV Cache 是独立的。第 N 层的注意力只用到第 N 层自己的 KV Cache。所以 airllm 需要为每一层维护独立的 KV Cache。当层被释放时它的 KV Cache 也要保留因为后续 token 生成还要用这就意味着 KV Cache 的总量是所有层的 KV Cache 之和和全量加载时一样大。这部分显存是省不掉的。对于长序列生成KV Cache 可能比单层权重还大。这时候可以考虑KV Cache 量化或者分页管理把不常用的 KV Cache 换出到内存。airllm 本身对 KV Cache 的优化有限如果你的场景序列特别长可能需要额外的手段。3.4 精度保持与数值一致性前面提到分层加载理论上和全量加载结果一致但实际会有微小差异。差异来源主要有两个一是计算顺序全量加载时所有层权重都在显存计算是连续的分层加载时层与层之间有数据传输中间可能有精度转换。二是累加误差FP16 本身精度有限多次搬运和计算可能引入舍入误差。实测下来这些差异通常在 1e-3 到 1e-4 量级对生成质量基本没有可见影响。但如果你做的是需要严格复现的实验建议用 FP32 或者做数值对比验证。airllm 支持指定计算精度默认跟随模型原始精度。4. 实操过程与核心环节实现4.1 环境准备与依赖安装先说一下我的测试环境方便你对照一台单卡机器显卡显存 24GB内存 128GB硬盘是 NVMe SSD系统是常见的 Linux 发行版。这个配置跑 70B 模型的分层推理是够的跑 13B 或 30B 会更轻松。安装 airllm 本身很简单它是个 Python 包用 pip 就能装pip install airllm但它依赖 PyTorch、Transformers、Accelerate 这些库如果环境里没有会自动装。建议先确认 PyTorch 版本和你的 CUDA 驱动匹配否则后面加载模型时会报错。可以用下面这行验证import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果cuda.is_available()返回 False说明 CUDA 环境没配好先解决这个再往下走。4.2 模型加载与分片配置airllm 的用法非常简洁核心就是一个AutoModel类。以加载一个 70B 模型为例from airllm import AutoModel model AutoModel.from_pretrained(your-model-path-or-name)第一次加载时airllm 会检查模型是否已经分片。如果没有它会自动做分片这个过程可能比较久取决于模型大小和硬盘速度。70B 模型分片大概要十几分钟到半小时。分片结果会缓存下来下次加载就快了。分片缓存的位置默认在用户目录下的某个隐藏文件夹如果硬盘空间紧张可以指定到其他盘model AutoModel.from_pretrained( your-model-path-or-name, layer_shards_saving_path/path/to/your/cache )注意分片缓存会占用和原始模型差不多大小的空间。70B 模型 FP16 约 140GB分片缓存也要 140GB 左右。提前确认硬盘空间够用。4.3 推理调用与参数设置加载完之后推理调用和普通 Transformers 用法几乎一样input_text 请用一句话解释什么是分层推理。 input_tokens model.tokenizer( [input_text], return_tensorspt, return_attention_maskFalse, truncationTrue, max_length128 ).input_ids.cuda() generation_output model.generate( input_tokens, max_new_tokens64, use_cacheTrue, return_dict_in_generateTrue ) output model.tokenizer.decode(generation_output.sequences[0]) print(output)这里有几个参数值得说。max_new_tokens控制生成长度越长耗时越久因为每个 token 都要把所有层走一遍。use_cacheTrue开启 KV Cache能显著加速长序列生成但会占显存。max_length是输入截断长度输入太长会拖慢首 token 的生成。实测下来70B 模型在 24GB 显存机器上生成 64 个 token 大概需要几分钟。这个速度做交互式对话肯定不行但做批量实验、跑评测、验证效果是够用的。4.4 显存与内存的监控调优跑起来之后建议开一个终端监控资源占用watch -n 1 nvidia-smi同时用htop或free -h看内存。正常情况下显存占用应该稳定在 5-10GB 之间波动内存占用接近模型总大小。如果显存持续上涨直到 OOM可能是 KV Cache 没释放或者预读取层数设多了。如果内存不够系统会开始用交换分区速度会断崖式下跌这时候要么加内存要么换更小的模型。调优的方向主要是两个预读取层数和KV Cache 策略。预读取层数增加能提升速度但占显存减少则相反。KV Cache 如果支持量化开启后能省不少显存但可能影响长序列的生成质量。5. 常见问题与排查技巧实录5.1 加载时报错找不到层这是最常见的问题通常出现在模型结构和 airllm 预期不一致的时候。报错信息一般是KeyError或者Layer not found。排查思路是先看模型的config.json确认architectures字段和num_hidden_layers。然后检查分片缓存目录看生成的层文件数量是否和num_hidden_layers一致。如果不一致说明分片逻辑没正确识别层结构可能需要手动指定层名模式或者换一个 airllm 版本。5.2 生成速度慢到无法接受速度慢的原因可能有好几个。先确认权重是不是常驻内存了如果每次生成都从硬盘读那速度肯定慢。检查方法是看生成过程中硬盘 IO 是否持续很高。如果是说明内存不够权重被换出了。其次是预读取有没有生效可以试着把预读取层数调大一点。最后是输入长度输入越长首 token 越慢尽量把输入控制在必要范围内。5.3 生成结果重复或乱码这通常不是 airllm 的问题而是模型本身或者采样参数的问题。检查一下temperature、top_p、repetition_penalty这些采样参数。如果参数正常但结果还是差可能是精度问题试着用 FP32 跑一次对比。还有一种可能是 KV Cache 管理有 bug关掉use_cache试试如果关掉就正常了那就是 KV Cache 的问题。5.4 显存溢出OOMOOM 的原因通常是预读取层数太多、KV Cache 太大、或者同时跑了别的占显存的程序。解决方法是降低预读取层数、缩短生成序列、关掉其他程序。如果都不行考虑用更小的模型或者更低的精度。下面这张表把常见问题和排查方向整理了一下方便你快速定位问题现象可能原因排查方向加载报错找不到层模型结构不匹配检查 config.json 和分片缓存生成速度极慢权重未常驻内存检查内存占用和硬盘 IO结果重复乱码采样参数或精度问题调整采样参数对比 FP32显存溢出预读取过多或 KV Cache 过大降低预读取缩短序列内存不足模型太大加内存或换小模型提示遇到问题时先把max_new_tokens设成 1只生成一个 token这样能快速判断是加载阶段的问题还是生成阶段的问题。加载阶段的问题通常在第一个 token 生成前就会暴露。5.5 几个我踩过的坑第一个坑是分片缓存路径没设对结果缓存写到了系统盘把系统盘撑爆了。后来养成习惯加载大模型前先确认缓存路径和剩余空间。第二个坑是内存估算错误。我以为 128GB 内存跑 70B 模型够了但实际上系统本身要占一部分加上 Python 进程开销和 KV Cache实际可用内存没那么多。后来跑 70B 时把其他程序都关了才稳定。第三个坑是预读取层数设太大。一开始为了追求速度把预读取设成 4 层结果显存直接爆了。后来发现预读取 1-2 层就够了再多了收益递减还占显存。第四个坑是忽略了输入长度。有一次输入了一段很长的文本首 token 生成等了快十分钟我还以为程序卡死了。后来才知道输入越长首 token 的注意力计算越重时间线性增长。现在我会主动控制输入长度必要时先做摘要再喂给模型。6. 这套方案还能怎么扩展airllm 的分层推理思路其实可以延伸到很多场景。比如多卡场景可以把不同层分配到不同显卡上每张卡只负责一部分层层间通过 PCIe 或 NVLink 通信。这样单卡显存要求进一步降低但通信开销会增加。再比如CPU 推理如果完全没有显卡可以把层加载到内存用 CPU 计算速度会慢很多但至少能跑。还有一个方向是和量化结合。分层加载解决的是显存容量问题量化解决的是显存和算力问题两者不冲突。你可以先量化到 INT8再用分层加载这样显存占用更低速度也可能更快。airllm 本身对量化的支持有限但你可以用其他工具先量化模型再用 airllm 加载量化后的权重。最后说个我自己的体会airllm 这类方案的价值不在于性能而在于降低了实验门槛。以前想跑 70B 模型要么租多卡云服务器要么买专业卡成本很高。现在用一张消费级卡加足够的内存就能跑起来虽然慢但能让你先验证想法、调通流程等确认有价值了再上更好的硬件。这个先跑通再优化的思路在实际项目里比一上来就追求极致性能要务实得多。