单卡跑70B大模型：airllm逐层加载推理实战与调优

1. 单卡跑大模型这件事，到底卡在哪

第一次看到 airllm 这个项目的时候，我正在一台只有单张消费级显卡的机器上折腾一个 70B 级别的模型。显存不够，量化到 4bit 依然爆，换更小的模型又达不到想要的效果。airllm 给出的思路很直接：不靠多卡堆显存，而是把模型权重按层拆开，推理时只把当前需要的层加载进显存，用完就释放，下一层再加载。听起来像是把“内存换显存”的老办法搬到了大模型推理上，但它做得比手动 offload 更细，也更适合单卡环境。

这个项目解决的核心问题就一个：显存不够时，怎么还能把大模型跑起来。它适合手里只有一张卡、显存 8G 到 24G 之间、想跑 70B 甚至更大模型的人。不适合追求高吞吐的生产服务场景，也不适合对首 token 延迟极度敏感的任务。但如果你只是想验证一个想法、跑个 demo、做点离线推理，airllm 的性价比非常高。

我把它定位成“单卡大模型推理的应急方案”，不是替代 vLLM 或 TensorRT-LLM 的通用方案。它的价值在于把门槛降到了最低：不需要多卡，不需要复杂的分布式配置，装完就能跑。下面我从设计思路、核心机制、实操步骤、参数调优和踩坑记录几个方面，把这个项目拆开讲清楚。

2. airllm 的整体设计思路拆解

2.1 为什么选择“逐层加载”而不是量化

量化当然也能省显存，4bit 量化能把 70B 模型压到 35G 左右，但单张 24G 卡还是放不下。继续压到 2bit，精度损失就明显了，生成质量会掉得厉害。airllm 走的是另一条路：不改变模型精度，而是改变加载方式。它把模型按 transformer 层切分，每一层单独存成一个文件，推理时只加载当前层，算完立刻释放。这样显存占用基本等于“单层参数量 + 激活值 + KV cache”，70B 模型单层大约 1.5G 到 2G（FP16），加上激活和 cache，24G 卡完全够用。

这个思路的代价是推理速度。每层都要从磁盘读一次，磁盘 IO 成了瓶颈。但 airllm 做了几件事来缓解：一是用 safetensors 格式存储，加载速度快；二是支持预读取，当前层在算的时候，下一层已经在后台加载；三是可以指定把部分层常驻显存，减少重复加载。实测下来，70B 模型在单张 24G 卡上，每秒能出 1 到 3 个 token，做离线推理完全可接受。

2.2 分层存储的具体实现方式

airllm 在第一次加载模型时，会把原始权重文件拆分成按层命名的独立文件。比如model.layers.0.safetensors、model.layers.1.safetensors，一直到最后一层。每个文件只包含该层的权重，大小均匀。拆分过程只做一次，之后每次推理都直接读这些分片文件。

拆分的好处是加载粒度细。如果按原始的大文件加载，即使只读一层，也要把整个文件映射到内存，效率低。拆成小文件后，操作系统可以更精准地做页缓存，重复推理时部分层可能已经在系统缓存里，磁盘 IO 会少很多。我试过连续跑同一个模型多次，第二次开始速度明显提升，就是因为系统缓存住了部分层文件。

2.3 与常见推理框架的定位差异

vLLM 和 TensorRT-LLM 追求的是高吞吐、低延迟，它们假设你有足够的显存把整个模型放进去，或者至少能放大部分。airllm 反过来，它假设你显存严重不足，愿意用速度换可行性。这两类方案不是竞争关系，而是互补关系。生产环境用 vLLM，个人实验用 airllm，各取所需。

还有一个差异是部署复杂度。vLLM 需要配置 tensor parallel、pipeline parallel，airllm 基本就是 pip install 然后改几行代码。对于不熟悉分布式推理的人来说，airllm 的上手成本几乎为零。这也是它在社区里传播快的原因之一。

3. 核心机制与关键参数解析

3.1 逐层加载与释放的完整流程

推理时的流程大致是这样的：首先加载 embedding 层，把输入 token 转成向量；然后进入 transformer 层循环，每一层从磁盘加载权重，做 attention 和 FFN 计算，算完释放该层权重，把输出传给下一层；最后加载 lm_head，输出 logits，采样得到下一个 token。整个过程是串行的，同一时间显存里只有一层权重。

这里有个细节：KV cache 是跨层共享的，不会每层都重建。airllm 会把 KV cache 保留在显存里，只释放权重。这样 attention 计算时不需要重新算历史 token 的 key 和 value，节省了大量计算。KV cache 的大小取决于序列长度和层数，70B 模型在 2048 序列长度下，KV cache 大约占 2G 到 4G 显存，加上单层权重和激活，24G 卡绰绰有余。

3.2 预读取机制与磁盘 IO 优化

预读取是 airllm 提速的关键。默认情况下，它在计算第 N 层时，会异步加载第 N+1 层的权重。这样当第 N 层算完，第 N+1 层已经准备好了，不需要等待磁盘 IO。预读取的深度可以配置，一般设为 1 到 2 层。设得太深会占用额外显存，设得太浅又起不到隐藏 IO 延迟的作用。

磁盘 IO 方面，建议把模型分片文件放在 SSD 上。机械硬盘的随机读性能太差，每层加载可能要几百毫秒，推理速度会掉到每秒零点几个 token。NVMe SSD 的随机读能到几十微秒级别，基本不会成为瓶颈。我用 SATA SSD 和 NVMe SSD 对比过，NVMe 下 70B 模型能到 2 到 3 token/s，SATA SSD 只有 1 token/s 左右。

3.3 关键参数：layers_per_load 与 prefetch

layers_per_load控制每次加载多少层。默认是 1，也就是逐层加载。如果显存有富余，可以设为 2 或 3，一次加载多层，减少加载次数，提升速度。但设得太大就退化成全量加载了，显存会爆。建议从 1 开始试，观察显存占用，逐步往上加，找到速度和显存的平衡点。

prefetch控制预读取的层数。设为 1 表示提前加载下一层，设为 2 表示提前加载下两层。预读取会占用额外显存，因为加载进来的层要暂存。如果显存紧张，可以把 prefetch 设为 0，关闭预读取，但速度会明显下降。我的经验是，24G 卡跑 70B 模型，layers_per_load=1、prefetch=1 是比较稳的配置。

4. 从零开始的完整实操流程

4.1 环境准备与依赖安装

先确认你的显卡驱动和 CUDA 版本。airllm 依赖 PyTorch，所以 PyTorch 能跑起来，airllm 基本就能跑。建议用 Python 3.9 到 3.11，太新的版本可能有些依赖还没适配。创建一个干净的虚拟环境，避免和系统里的其他包冲突。

python -m venv airllm_env source airllm_env/bin/activate pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install airllm

CUDA 版本根据你的驱动来选，cu118 是比较通用的。如果驱动比较新，也可以用 cu121。安装完后用torch.cuda.is_available()验证一下，返回 True 就说明环境没问题。

4.2 模型下载与分层拆分

airllm 支持直接从 Hugging Face 拉模型，但国内网络可能不稳定。建议先用huggingface-cli download把模型拉到本地，再用 airllm 加载本地路径。拉模型的时候注意选对版本，有些模型有多个分支，选 base 或 instruct 看你的任务需求。

from airllm import AutoModel model = AutoModel.from_pretrained( "meta-llama/Llama-2-70b-hf", layer_shards_saving_path="./shards/llama70b", compression="4bit" )

第一次加载时，airllm 会自动做分层拆分，把权重存到layer_shards_saving_path指定的目录。拆分过程比较慢，70B 模型可能要十几分钟到半小时，取决于磁盘速度。拆分完成后，后续加载就直接读分片文件，不再重复拆分。

4.3 推理代码与参数配置

推理代码和普通 transformers 用法差不多，只是多了几个 airllm 特有的参数。下面是一个完整的推理示例：

from airllm import AutoModel import torch model = AutoModel.from_pretrained( "meta-llama/Llama-2-70b-hf", layer_shards_saving_path="./shards/llama70b", compression="4bit", layers_per_load=1, prefetch=1 ) input_text = "请用一句话解释什么是机器学习。" input_ids = model.tokenizer(input_text, return_tensors="pt").input_ids.cuda() output = model.generate( input_ids, max_new_tokens=128, do_sample=True, temperature=0.7, top_p=0.9 ) print(model.tokenizer.decode(output[0], skip_special_tokens=True))

compression参数可以选4bit或8bit，进一步降低单层权重的显存占用。4bit 压缩后，70B 模型单层大约 0.5G，24G 卡可以开更大的 layers_per_load。但压缩会带来精度损失，如果对生成质量要求高，建议用 FP16 不压缩。

4.4 显存与速度的实测数据

我在一张 24G 显存的卡上跑了几个配置，数据如下：

从数据看，4bit 压缩对速度有提升，因为单层权重更小，加载更快。layers_per_load 和 prefetch 增加也能提速，但显存占用会上升。24G 卡跑 70B 模型，4bit + layers_per_load=2 + prefetch=1 是比较均衡的配置，显存留有余量，速度也还行。

5. 实操中踩过的坑与排查技巧

5.1 磁盘空间不足导致拆分失败

第一次跑 70B 模型时，我没注意磁盘空间，拆分到一半报错退出。分层拆分后的文件总大小和原始模型差不多，70B FP16 大约 140G，4bit 压缩后大约 35G。拆分前一定要确认目标磁盘有足够空间，建议预留原始模型 1.5 倍的空间，因为拆分过程中原始文件和分片文件会同时存在。

如果空间不够，可以先把原始模型删掉再拆分，但这样就没法回退了。更稳妥的做法是挂一块大容量 SSD 专门放模型分片，拆分完再把原始文件清理掉。

5.2 显存碎片导致的 OOM

有时候显存总量够，但还是报 OOM，这通常是显存碎片导致的。PyTorch 的缓存分配器会预留一些显存，频繁加载释放容易产生碎片。解决办法是设置PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128，限制单次分配的最大块大小，减少碎片。或者在推理前调用torch.cuda.empty_cache()，清理未使用的缓存。

另一个原因是 KV cache 没释放。如果连续跑多个请求，KV cache 会累积，最终撑爆显存。airllm 默认会在每次 generate 后释放 KV cache，但如果你手动调用了底层 API，可能需要自己清理。建议每个请求结束后检查一下显存占用，确认没有泄漏。

5.3 生成质量下降的排查思路

用了 4bit 压缩后，生成质量可能会下降，表现为重复、逻辑混乱、答非所问。先确认是不是压缩导致的：换成 FP16 跑同样的输入，对比输出。如果 FP16 正常，那就是压缩损失太大，可以试试 8bit 压缩，或者换更鲁棒的量化方法。

另一个常见原因是 temperature 和 top_p 设置不当。airllm 的推理速度和采样参数有关，temperature 太高会导致输出随机性大，看起来像质量下降。建议先用 temperature=0.7、top_p=0.9 跑一组基线，再逐步调整。如果任务需要确定性输出，可以用 greedy search，temperature 设为 0。

5.4 常见问题速查表

6. 进阶调优与场景扩展

6.1 混合精度与分层压缩策略

airllm 支持对不同层用不同的压缩策略。比如前几层对精度敏感，用 FP16；后面几层用 4bit。这样可以在显存和精度之间做更细的权衡。实现方式是在加载模型时传入一个压缩配置字典，指定每层的压缩方式。这个功能在社区版里可能没有直接暴露，需要改一点源码，但思路是可行的。

我试过对 70B 模型的前 10 层用 FP16，其余用 4bit，显存占用比全 4bit 多 2G 左右，但生成质量明显好于全 4bit。如果你的任务对精度要求高，又不想全 FP16 爆显存，可以试试这种混合策略。

6.2 多轮对话的 KV cache 管理

多轮对话时，KV cache 会随着轮次增加而增长。airllm 默认每轮重新计算整个上下文的 KV，效率低。可以手动保留历史 KV cache，只计算新 token 的 KV，然后拼接。这样能显著减少计算量，提升多轮对话的响应速度。

具体做法是在 generate 时传入past_key_values，把上一轮的 KV cache 传进去。airllm 的 generate 接口支持这个参数，但需要自己管理 cache 的生命周期。注意 cache 会占用显存，轮次多了可能爆，需要设置一个上限，超过就截断历史。

6.3 批量推理的可行性分析

airllm 的逐层加载机制决定了它不适合大批量推理。批量推理时，每层要处理多个样本，激活值会成倍增加，显存占用上升。而且批量样本的层加载是串行的，磁盘 IO 次数不变，但每次加载后要处理更多数据，整体吞吐提升有限。

如果确实需要批量推理，建议用小的 batch size，比如 2 到 4，配合 layers_per_load=1。再大就得不偿失了，不如用 vLLM 在多卡上跑。airllm 的定位还是单请求、离线场景，批量不是它的强项。

6.4 与其他工具的配合使用

airllm 可以和 LangChain、LlamaIndex 这类框架配合，把 airllm 包装成一个 LLM 接口，接入现有的 RAG 流程。包装时注意 airllm 的推理速度慢，RAG 流程里的超时设置要放宽，否则容易超时失败。另外 airllm 不支持流式输出，如果需要流式，得自己改 generate 逻辑，逐 token 返回。

我还试过把 airllm 和 Gradio 结合，做一个简单的本地对话界面。Gradio 的队列机制可以避免并发请求，正好适合 airllm 的单请求特性。界面响应会慢一些，但能用，适合个人实验。

7. 个人实操体会与建议

airllm 这个项目最吸引我的地方是它把“单卡跑大模型”这件事变得足够简单。不需要懂分布式，不需要配多卡，装完改几行代码就能跑。虽然速度不快，但对于验证想法、做离线任务来说，完全够用。我用它跑过 70B 模型的文本摘要和问答，生成质量在 4bit 压缩下依然可接受，速度大约 2 token/s，跑一个 200 字的摘要大概一分半钟。

如果你也想试试，我的建议是：先用小模型（比如 13B）跑通流程，确认环境没问题，再上大模型。磁盘一定要用 SSD，NVMe 最好。显存 24G 是跑 70B 的底线，再低就得用更激进的压缩，质量会明显下降。参数从 layers_per_load=1、prefetch=1 开始调，观察显存和速度的变化，找到适合自己硬件的平衡点。

还有一个细节：airllm 的版本更新比较快，不同版本的 API 可能有差异。遇到问题先看项目 README 和 issue，很多坑别人已经踩过了。如果 README 没写清楚，可以去翻源码，核心逻辑不复杂，读一遍基本就明白它的工作机制了。