airllm 分层推理实战：低显存本地跑大模型

1. 大模型本地推理的显存困境与 airllm 的破局思路

把一个大模型塞进消费级显卡，这件事在过去两年里劝退了无数人。一张 24GB 显存的卡，跑 7B 参数量的模型还算从容，一旦上到 70B，哪怕只是做一次前向推理，权重加载阶段就会直接爆显存。很多人第一反应是去租云端算力，但数据隐私、调用成本、网络延迟这三座大山压下来，本地推理的需求反而越来越刚性。airllm 这个项目就是冲着这个痛点来的——它做的事情说起来很朴素：让大模型在显存不够的机器上也能跑起来。

它的核心思路并不复杂，用一句话概括就是“分层加载、按需换入换出”。传统做法是把整个模型的权重一次性加载到显存里，然后开始计算。airllm 换了个思路：模型有几十层，每一层计算的时候才把这一层的权重搬到显存，算完立刻释放，下一层再搬下一层。这样一来，显存里同时存在的权重只有一层，峰值占用从“整个模型”降到“单层模型”。70B 的模型单层权重也就几百 MB 到 1GB 出头，一张 8GB 甚至 4GB 的卡理论上都能扛住。

这个思路听起来像是把“空间换时间”反过来用——用时间换空间。确实如此，airllm 的推理速度会比全量加载慢，因为每一层都要经历一次从内存（甚至磁盘）到显存的搬运。但它的价值在于把不可能变成可能：原本根本跑不起来的硬件配置，现在能跑起来了。对于做研究、做验证、做小批量推理的场景，这个 trade-off 是完全值得的。

我最初接触 airllm 是在一个需要本地跑 70B 模型做文本生成的场景里，手头只有一张 12GB 的卡。试过量化、试过 CPU 推理，效果都不理想——量化损失精度，CPU 推理慢到无法接受。airllm 是当时唯一一个既保住了模型精度、又能在有限显存下跑通的方案。这篇文章我会把 airllm 的设计逻辑、核心机制、实操步骤、参数调优、踩坑经验完整拆一遍，适合所有想在低显存环境下跑大模型的人参考。

2. airllm 核心机制深度拆解

2.1 分层推理到底是怎么工作的

要理解 airllm，得先理解 Transformer 模型的结构。一个标准的 decoder-only 大模型，本质上是一堆结构相同的层堆叠起来的，每一层包含自注意力模块和前馈网络模块。输入经过 embedding 层之后，依次穿过第 1 层、第 2 层……直到第 N 层，最后经过输出层得到 logits。

airllm 做的事情，就是把这个“依次穿过”的过程拆开。它不会在推理开始前把所有层的权重都加载到显存，而是维护一个“当前层”的概念。计算第 k 层的时候，只把第 k 层的权重加载进显存，前向计算完成后，把这一层的输出（也就是 hidden states）保留在显存或内存里，然后释放第 k 层的权重，接着加载第 k+1 层。整个过程像流水线一样，一层一层往前推。

这里有个关键细节：hidden states 的体积远小于权重。以 70B 模型为例，单层权重可能是 800MB 到 1.2GB，但 hidden states 的维度通常是 8192 左右，一个 batch、一个 token 位置的 hidden state 也就几十 KB。所以显存里真正的大头是权重，而权重被 airllm 控制在了“单层”这个量级。这就是它能把显存需求压下来的根本原因。

注意：airllm 的分层加载并不是“把模型切成很多份分别放到不同设备”，它始终是单设备推理，只是通过时间上的串行加载来降低空间占用。这一点和模型并行有本质区别。

2.2 权重存储与加载策略的选择逻辑

airllm 支持从多种来源加载权重，最常见的是从 Hugging Face 格式的模型目录加载。它会把每一层的权重单独保存成文件，推理时按层读取。这个设计带来一个额外好处：你可以把权重放在磁盘上，而不必全部塞进内存。对于内存也不宽裕的机器，这一点很关键。

具体来说，airllm 在首次使用某个模型时，会做一次“分层拆分”的操作。它读取原始的模型权重文件（通常是 pytorch_model.bin 或 safetensors 格式），按照层的边界把权重切分成多个小文件，每个文件对应一层或几层。拆分完成后，后续推理就直接读取这些小文件，不再需要加载完整的模型。

这个拆分过程只需要做一次，之后可以复用。拆分后的目录结构大致是这样的：每一层的权重单独存放，另外还有 embedding 层、输出层、layer norm 等独立文件。airllm 在推理时会按照模型定义的层顺序，依次加载对应的文件。

选择这种策略的理由很直接：减少单次加载的数据量。如果每次推理都从完整的模型文件里读取某一层的权重，就需要先加载整个文件再切片，内存和 IO 开销都很大。拆分成小文件后，每次只需要读取当前层对应的那个文件，IO 效率高很多。

2.3 压缩与量化选项的取舍

airllm 提供了压缩选项，可以在拆分权重的时候对权重进行量化压缩，进一步降低存储和加载开销。常见的选项是 4bit 或 8bit 量化。量化的好处是权重文件变小，加载更快，显存占用也更低；代价是精度会有一定损失。

我的实测经验是：如果显存刚好卡在边界上，量化是必要的；如果显存有余量，优先用原始精度。比如 70B 模型在 12GB 卡上跑，不量化的话单层权重加载后加上 hidden states 和中间激活值，峰值可能接近 10GB，比较紧张；用 4bit 量化后单层权重降到 200MB 左右，显存压力小很多，速度也更快。

但量化对生成质量的影响是存在的，尤其是对数值敏感的任务，比如数学推理、代码生成，量化后的输出可能会出现更多错误。所以我的建议是：先不量化跑一次，看显存峰值和生成质量；如果显存不够再考虑量化，并且对比量化前后的输出差异，确认可接受后再用。

2.4 与同类方案的对比

市面上降低显存占用的方案不止 airllm 一种，常见的有量化加载、CPU offload、模型并行等。我把它们的核心特点整理成了一张表，方便对照选择。

airllm 的定位很清晰：单卡、低显存、可接受速度损失、希望保住精度。它不追求最快，也不追求多卡扩展，就是把“单卡跑大模型”这件事做到极致。如果你的场景是多卡，模型并行可能更合适；如果追求极致速度，全量加载或量化加载更好。airllm 填补的是“硬件不够但想跑”这个空白区间。

3. 从零开始跑通 airllm 的完整实操

3.1 环境准备与依赖安装

airllm 是一个 Python 库，安装方式很直接。我建议用虚拟环境，避免和系统里的其他包冲突。Python 版本建议 3.9 以上，3.10 或 3.11 更稳。

python -m venv airllm_env source airllm_env/bin/activate # Windows 用 airllm_env\Scripts\activate pip install airllm

airllm 本身依赖 PyTorch、Transformers、Accelerate 等库，pip 会自动装。但有几个点需要注意：

• PyTorch 版本要和 CUDA 版本匹配。如果你的机器有 NVIDIA 显卡，建议先装好对应 CUDA 版本的 PyTorch，再装 airllm，避免 airllm 自动拉一个 CPU 版本的 torch 下来。
• 磁盘空间要留够。分层拆分后的权重文件总大小和原始模型差不多，70B 模型大概 130GB 到 140GB（FP16），4bit 量化后大概 35GB 到 40GB。拆分过程中还会临时占用额外空间，所以磁盘至少要有原始模型两倍以上的余量。
• 内存建议 32GB 以上。虽然 airllm 把权重放在磁盘上，但加载单层权重、处理 hidden states、做前向计算都需要内存。16GB 内存跑 70B 会比较吃力，32GB 比较从容。

提示：如果你用的是 Windows，建议在 WSL2 里跑，原生 Windows 下有些依赖的编译和路径处理会有小问题。Linux 环境是最省心的。

3.2 模型分层拆分的操作步骤

第一次使用某个模型时，airllm 会自动做分层拆分。你也可以手动触发，这样能更清楚地看到过程。以加载一个 70B 模型为例：

from airllm import AutoModel model = AutoModel.from_pretrained("meta-llama/Llama-2-70b-hf")

第一次运行这段代码时，airllm 会检查本地缓存目录里有没有拆分好的权重。如果没有，它会开始拆分。拆分过程会打印进度，每一层拆完会保存成独立文件。70B 模型大概有 80 层，拆分时间取决于磁盘速度，机械硬盘可能要半小时以上，SSD 会快很多。

拆分完成后，目录结构大致如下：

~/.cache/airllm/ └── Llama-2-70b-hf/ ├── layer_0/ ├── layer_1/ ├── ... ├── layer_79/ ├── embedding/ └── output/

每个 layer 目录里存放该层的权重文件。后续推理时，airllm 就按这个结构逐层加载。

如果你想用压缩模式，可以在加载时指定压缩参数：

model = AutoModel.from_pretrained( "meta-llama/Llama-2-70b-hf", compression="4bit" )

这样拆分出来的权重文件会更小，加载更快，显存占用更低。但如前所述，精度会有损失，建议先对比再决定。

3.3 推理调用的完整代码与参数说明

拆分完成后，推理调用和普通 Transformers 用法很接近：

from airllm import AutoModel from transformers import AutoTokenizer model_path = "meta-llama/Llama-2-70b-hf" tokenizer = AutoTokenizer.from_pretrained(model_path) model = AutoModel.from_pretrained(model_path) input_text = "请用三句话解释什么是机器学习。" input_ids = tokenizer(input_text, return_tensors="pt").input_ids output = model.generate( input_ids, max_new_tokens=128, do_sample=True, temperature=0.7, top_p=0.9 ) print(tokenizer.decode(output[0], skip_special_tokens=True))

几个关键参数需要说明：

• max_new_tokens：控制生成的最大 token 数。airllm 每生成一个 token 都要把所有层跑一遍，所以这个值直接影响总耗时。建议先设小一点测试，确认流程通了再加大。
• do_sample / temperature / top_p：采样相关参数，和普通推理一致。如果做确定性任务，可以设 do_sample=False。
• batch size：airllm 对 batch 的支持有限，建议先用 batch size 为 1 跑通，再尝试增大。batch 增大会让 hidden states 变大，显存压力上升。

注意：airllm 的 generate 方法和标准 Transformers 的 generate 在行为上基本一致，但内部实现不同。它不会一次性把所有层的 KV cache 都保留，而是每层计算完就释放。这意味着多轮对话场景下，历史上下文的处理需要额外注意，不能直接依赖默认的 KV cache 机制。

3.4 显存与速度的实测数据

我在一张 12GB 显存的卡上跑了几个不同规模的模型，记录了大致的显存峰值和生成速度。数据仅供参考，实际会因硬件、模型、参数不同而有差异。

从数据能看出几个规律：显存峰值基本由单层权重加中间激活决定，和模型总规模关系不大；速度则和模型规模强相关，层数越多、每层计算量越大，生成越慢。70B 在 FP16 下每秒不到 1 个 token，确实慢，但考虑到这是在一张 12GB 卡上跑 70B，这个结果已经很有价值了。

如果换成 SSD 且内存充足，速度还能再提升一些，因为权重加载的 IO 瓶颈会缓解。机械硬盘是明显的短板，建议尽量用 SSD。

4. 实操中的常见问题与排查技巧

4.1 显存仍然不够怎么办

这是最常见的问题。airllm 已经把显存需求压到单层级别，但如果单层权重加上中间激活还是超了，就会 OOM。排查思路按优先级来：

第一，确认是不是真的在分层加载。有时候配置不对，airllm 可能退化成全量加载，显存自然爆。检查日志里有没有逐层加载的记录，或者用 nvidia-smi 观察显存是不是阶梯式上升。

第二，降低 batch size 和序列长度。hidden states 的体积和 batch size、序列长度成正比。把 batch 降到 1，序列长度控制在 512 以内，能显著降低激活值占用。

第三，启用压缩。4bit 量化能把单层权重降到原来的四分之一左右，是最直接的降显存手段。

第四，检查是否有其他进程占用显存。有时候是别的程序没释放干净，导致可用显存比预期少。

4.2 生成速度太慢的优化方向

速度慢是 airllm 的固有代价，但有些优化能缓解：

• 用 SSD 而不是机械硬盘。权重加载是 IO 密集型操作，SSD 的随机读取速度比机械硬盘快一个数量级。
• 增加内存。内存充足时，操作系统会缓存部分权重文件，减少磁盘读取。
• 用压缩模式。权重文件变小，加载更快，整体速度提升明显。
• 减少 max_new_tokens。生成 token 数是线性影响总耗时的，能少生成就少生成。
• 关闭不必要的日志输出。airllm 默认会打印每层加载信息，关掉能减少一点开销。

4.3 模型加载报错的排查清单

加载阶段的报错通常和路径、格式、权限有关。我整理了一个速查表：

提示：airllm 的缓存目录默认在用户主目录下的 .cache/airllm。如果主目录空间不够，可以通过环境变量 AIRLLM_CACHE_DIR 指定到其他磁盘。

4.4 生成质量异常的应对

有时候模型能跑起来，但生成的内容质量不对，比如重复、乱码、答非所问。这类问题通常和量化、参数设置有关：

• 量化导致的精度损失：如果用了 4bit 量化，先换回 FP16 对比。如果 FP16 正常而量化异常，说明该模型对量化比较敏感，考虑换 8bit 或不用量化。
• temperature 和 top_p 设置不当：temperature 太高会乱，太低会重复。一般 0.7 到 0.9 比较稳，top_p 0.9 左右。
• tokenizer 不匹配：确认用的 tokenizer 和模型是配套的，不同版本的 tokenizer 可能导致编码错误。
• 输入格式问题：有些模型对 prompt 格式有要求，比如需要特定的对话模板。检查是否按模型要求组织了输入。

4.5 多轮对话与上下文管理的坑

airllm 在单轮生成上表现稳定，但多轮对话需要额外处理。因为 airllm 不保留跨层的 KV cache，每一轮生成都要重新跑所有层，历史上下文需要手动拼接到输入里。这意味着：

• 对话轮数越多，输入序列越长，每轮耗时越长。
• 输入序列长度受模型最大上下文限制，超了要截断。
• 历史上下文的拼接格式要和模型训练时一致，否则效果会打折。

我的做法是维护一个对话历史列表，每轮把历史和新问题拼成完整 prompt，再送给模型。同时控制历史长度，超过一定 token 数就丢弃最早的轮次。这样虽然每轮都要重跑，但逻辑清晰，不容易出错。

5. 进阶用法与场景延展

5.1 自定义模型适配

airllm 默认支持 Hugging Face 格式的常见模型，但如果你用的是自定义结构的模型，可能需要做一些适配。核心是告诉 airllm 模型的层结构在哪里、每一层怎么访问。airllm 提供了一些钩子和配置项，可以指定层的前缀名、层数、每层的模块路径等。

适配的基本步骤是：先确认模型的层命名规则，比如 model.layers.0、model.layers.1 这样；然后在加载时传入对应的配置，让 airllm 知道去哪里找每一层。如果模型结构比较特殊，可能还需要继承 airllm 的基类，重写加载逻辑。这部分有一定门槛，建议先拿标准模型跑通，再尝试自定义。

5.2 批量推理与任务队列

airllm 单次推理慢，但如果有一批任务要处理，可以做成队列，让模型连续跑。做法是把待处理的输入排成队列，逐个送给模型生成，结果写入文件或数据库。这样虽然总耗时不变，但省去了反复加载模型的开销。

如果任务量大，可以考虑用多进程，每个进程加载一个 airllm 实例，分别处理不同的任务。但要注意显存和内存的竞争，进程数不能太多，否则会互相挤占资源。一般单卡跑一个实例比较稳，多卡可以每卡一个实例。

5.3 与其他工具的配合

airllm 可以和 LangChain、LlamaIndex 这类框架配合，作为本地推理后端。做法是把 airllm 的模型封装成一个符合框架接口的 LLM 类，实现 generate 或 predict 方法。这样就能在框架里直接调用 airllm，享受框架提供的链式调用、检索增强等能力。

配合时要注意接口的输入输出格式对齐。airllm 的 generate 接收 input_ids 返回 output_ids，框架通常期望接收字符串返回字符串。中间需要加一层 tokenizer 的编解码转换。另外，框架可能会有自己的缓存和批处理逻辑，要和 airllm 的分层加载机制协调好，避免冲突。

5.4 硬件选型的经验建议

如果你打算专门配一台机器跑 airllm，我的建议是：

• 显卡：显存 8GB 以上即可，重点是 CUDA 核心数和显存带宽。带宽越高，权重加载和计算越快。
• 内存：32GB 起步，64GB 更稳。内存越大，能缓存的权重越多，IO 压力越小。
• 磁盘：NVMe SSD，容量至少是模型大小的两倍。读写速度直接影响加载效率。
• CPU：不需要特别强，但核心数多一些有助于数据预处理和 IO 调度。

这套配置下来，跑 70B 模型做小批量推理是够用的。如果预算有限，优先保证内存和 SSD，显卡可以适当降级，因为 airllm 对显存的要求本身就不高。

6. 我个人在实际操作中的几点体会

airllm 这个项目最打动我的地方，是它把“不可能”变成了“可能”。在它出现之前，低显存跑大模型基本只有量化和 CPU 推理两条路，前者损精度，后者慢到没法用。airllm 用分层加载这个看似简单的思路，硬是开辟了第三条路。

但也要客观看待它的局限。速度慢是硬伤，70B 模型每秒不到 1 个 token，做交互式应用是不现实的，更适合离线批处理、研究验证、小批量生成这类场景。另外，它对磁盘和内存的要求不低，如果这两项也紧张，体验会打折扣。

我踩过的最大的坑是磁盘空间。第一次跑 70B 的时候没注意，拆分到一半磁盘满了，整个缓存目录损坏，只能删掉重来。后来我养成了习惯：跑大模型之前先 df -h 看一眼剩余空间，确认是模型大小的两倍以上再动手。还有就是 SSD 和机械硬盘的差距真的很大，同样的模型，机械硬盘上加载一层要好几秒，SSD 上不到一秒，整体速度差出好几倍。

最后分享一个小技巧：如果你只是想做验证，不需要完整生成，可以把 max_new_tokens 设成 1，这样只跑一次前向，能快速确认模型能不能加载、显存够不够、输出是否正常。确认没问题再加大生成长度，能省不少调试时间。