单卡消费级显卡跑大模型：airllm分层加载与显存优化实战

1. 为什么要在单张消费级显卡上跑大模型

大模型推理这件事，过去两年最大的门槛其实不是算法，而是显存。一个 70B 参数的模型，如果以 FP16 精度加载，光权重就要吃掉大约 140GB 显存，这还没算 KV Cache 和中间激活值。这意味着什么？意味着你手里那张 24GB 的卡连门都进不去，甚至连 7B 模型在长上下文场景下都会捉襟见肘。很多人第一反应是量化，INT8 砍一半，INT4 再砍一半，但量化本身有精度损失，而且 4-bit 的 70B 依然需要 35GB 左右，单张 24GB 卡还是放不下。

airllm 这个项目切入的角度就不一样了。它不去动模型的数值精度，而是从内存调度这个维度做文章。核心思路一句话概括：把模型按层切分，推理时只把当前需要的层加载进显存，用完就释放，下一层再加载。这样显存占用从"整个模型"降到"单层 + 少量缓冲"，理论上 70B 模型在 4GB 显存的卡上都能跑起来。代价是推理速度会慢，因为每层都要从磁盘或内存重新搬运，但至少它让"能不能跑"这个问题变成了"跑得慢不慢"。

我第一次看到这个思路的时候，直觉是"这不就是把 CPU offload 做到极致吗"。确实，HuggingFace Accelerate 早就有device_map="auto"把部分层放到 CPU 内存，但那个方案要求 CPU 内存足够大，而且层与层之间的搬运是静态的。airllm 更激进的地方在于，它连 CPU 内存都不要求能装下整个模型，而是把模型权重以分片文件的形式存在磁盘上，按需读取。这就把硬件门槛压到了极低——你需要的只是一块能跑 CUDA 的显卡、一块够大的硬盘、以及足够的耐心。

适合谁来参考这个方案？我梳理了三类人。第一类是个人开发者或学生，手头只有一张 8GB 或 12GB 的游戏卡，想本地跑大模型做实验、做 demo，不想花钱租云卡。第二类是做模型部署验证的工程师，需要在低配环境里快速验证一个模型能不能跑通、输出质量如何，不追求吞吐。第三类是对推理优化感兴趣的学习者，airllm 的代码结构清晰，是理解"分层加载 + 显存复用"这套机制的绝佳样本。如果你追求生产级的高并发推理，那这个方案不适合你，它的定位就是"低资源可用性"。

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

2.1 分层加载的核心机制

要理解 airllm 为什么能省显存，得先搞清楚标准推理流程里显存都被什么占了。以 Transformer 为例，一次前向传播，显存里同时存在的东西包括：所有层的权重、当前层的激活值、KV Cache、以及 PyTorch 自己的缓存分配器留下的碎片。其中权重是大头，70B 模型 140GB 里，权重占了 130GB 以上。airllm 的做法是：权重不常驻显存，只保留当前正在计算的那一层。

具体来说，模型被保存成多个分片文件，每个文件对应若干层。推理时，第 0 层从磁盘加载到显存，算完，把这一层的输出（激活值）留在显存里，然后释放第 0 层的权重，加载第 1 层，以此类推。这样显存峰值就是"单层权重 + 激活值 + KV Cache"。单层权重对于 70B 模型来说，假设 80 层，每层大约 1.75GB（FP16），加上激活和 KV，4GB 到 6GB 显存就能兜住。

这里有个关键设计点：为什么按层切而不是按张量切。按层切的好处是层与层之间是串行依赖的，天然适合流水线式的加载释放，不需要复杂的依赖管理。如果按张量切，同一层内部的多个权重矩阵之间有并行计算关系，你得保证它们同时在场，调度复杂度会高很多。airllm 选择按层，是用"调度简单"换"加载次数多"，这个取舍在低资源场景下是合理的，因为低资源场景本来就不追求速度。

2.2 与量化方案的对比取舍

很多人会问，既然量化也能省显存，为什么还要用 airllm 这种"搬运"方案？我把两者的差异整理成一张表，方便对照。

从表里能看出来，airllm 和量化不是替代关系，而是可以叠加的。你可以先把模型量化到 INT4，再用 airllm 分层加载，这样单层权重更小，加载更快，显存占用更低。但要注意，量化后的模型分片需要重新保存，airllm 默认加载的是原始精度分片，叠加量化需要额外处理。

我个人的判断是：如果你的显存缺口在 2 倍以内，优先考虑量化；如果缺口在 5 倍以上，airllm 这类方案才真正有价值。比如 7B 模型 FP16 需要 14GB，你有 8GB 卡，量化到 INT4 变成 3.5GB，直接就能跑，没必要上 airllm。但 70B 模型 FP16 需要 140GB，你只有 8GB，量化到 INT4 还要 35GB，还是放不下，这时候 airllm 就是唯一出路。

2.3 磁盘 IO 成为新瓶颈的应对

分层加载把瓶颈从显存转移到了磁盘 IO。每跑一层就要读一次磁盘，70B 模型 80 层，一次前向传播就是 80 次读取，每次读 1.75GB，总共 140GB 的读取量。如果你的硬盘是机械盘，顺序读大概 150MB/s，那光读权重就要 900 多秒，这还没算计算时间。所以 airllm 对存储介质是有要求的，强烈建议放在 NVMe SSD 上，顺序读能到 3GB/s 以上，140GB 大约 45 秒，虽然还是慢，但至少能接受。

airllm 在代码层面也做了一些缓解 IO 压力的设计。比如它支持预读取，在当前层计算的同时，后台线程已经开始加载下一层，把 IO 和计算重叠起来。另外它会把分片文件做内存映射（mmap），避免每次读取都走完整的文件系统调用。这些优化不能消除 IO 瓶颈，但能把等待时间压下来不少。实测下来，在 NVMe 上跑 70B 模型，单 token 生成时间大概在 3 到 8 秒，取决于你的 CPU 和磁盘性能。

注意：如果你的模型放在网络存储或外接 USB 硬盘上，IO 延迟会高到无法忍受，基本不具备可用性。这一点在动手之前一定要确认。

3. 核心细节解析与实操要点

3.1 模型分片的保存格式

airllm 要能按层加载，前提是模型已经被保存成"按层分片"的格式。它内部用的是 safetensors 格式，每个分片文件对应模型的一部分层。为什么要用 safetensors 而不是 PyTorch 原生的.bin？两个原因。第一，safetensors 的加载是零拷贝的，可以直接 mmap 到内存，读取速度快。第二，safetensors 不执行任意代码，安全性比 pickle 格式的.bin高，加载时不会有反序列化风险。

分片的粒度是可以配置的。默认情况下，airllm 会按照模型的层结构自动切分，每个分片包含若干层。你可以通过参数控制每个分片的大小，比如设置成每片 2GB，这样在加载单层时，实际读取的数据量会更接近单层大小，减少浪费。但分片太小也有问题，文件数量多了之后，文件系统的元数据开销会上升，而且预读取的粒度变细，调度开销增加。我的经验是，分片大小控制在 1GB 到 4GB 之间比较合适，既能保证单次读取效率，又不会让文件数量爆炸。

保存分片的过程本身需要一点时间。以 70B 模型为例，从原始权重转换成分片格式，在 NVMe 上大概要 10 到 20 分钟，取决于 CPU 和磁盘速度。这个过程只需要做一次，之后推理就直接加载分片。如果你要换模型，得重新做一次转换。这里有个小技巧：转换时把分片目录放在和推理时相同的磁盘上，避免转换完还要跨盘拷贝，白白多花时间。

3.2 显存复用的关键参数

airllm 在显存管理上有几个关键参数，理解它们对调优很重要。第一个是layer_shards_saving_path，指定分片保存的位置，这个前面说过了。第二个是compression，它支持对分片做压缩，比如用 4-bit 或 8-bit 压缩存储，加载时再解压。这个压缩和量化不是一回事，它只影响存储和 IO，不影响计算精度，因为解压后还是 FP16。压缩的好处是磁盘占用小、IO 量小，坏处是解压要消耗 CPU。如果你的 CPU 比较强、磁盘比较慢，开压缩是划算的。

第三个是profiling_mode，打开之后会输出每一层的加载时间和计算时间，方便你定位瓶颈。我建议第一次跑的时候一定打开这个，看看时间到底花在 IO 上还是计算上。如果 IO 占了 80% 以上，说明磁盘是瓶颈，考虑换 NVMe 或者开压缩。如果计算占了大部分，那说明你的卡算力不够，这时候换更快的卡比优化 IO 更有效。

还有一个容易被忽略的点是KV Cache 的管理。分层加载省的是权重显存，但 KV Cache 是随序列长度增长的，长上下文场景下 KV Cache 可能比单层权重还大。airllm 默认会把 KV Cache 留在显存里，如果你跑很长的序列，显存还是会被撑爆。这时候要么限制最大序列长度，要么把 KV Cache 也做分层管理（部分放 CPU 内存），但后者 airllm 支持得不算完善，需要自己改代码。

3.3 推理速度的预期管理

用 airllm 之前，一定要对速度有个合理预期，不然很容易失望。我拿几个典型配置做个估算。假设你用一张 8GB 显存的卡，NVMe SSD，跑 70B 模型，单层权重 1.75GB，读取时间按 3GB/s 算大约 0.6 秒，计算时间按卡的算力算大约 0.3 秒，加起来单层接近 1 秒，80 层就是 80 秒一个 token。这显然太慢了，实际中因为有预读取重叠，能压到 3 到 8 秒一个 token。生成 100 个 token 就是 5 到 13 分钟。

这个速度能干什么？能做离线批处理，比如晚上跑一批数据，第二天看结果。能做低频交互，比如你问一个问题，等几分钟看答案。但绝对做不了实时对话，用户等 3 秒没响应就跑了。所以 airllm 的定位要清楚：它是"能跑"的方案，不是"好用"的方案。如果你需要实时性，要么加钱上大显存卡，要么用 API。

提示：如果你的场景是"验证模型效果"而不是"提供服务"，airllm 完全够用。跑几个 prompt 看看输出质量，几分钟的等待是可以接受的。

4. 实操过程与核心环节实现

4.1 环境准备与依赖安装

动手之前，先把环境理清楚。airllm 是 Python 包，依赖 PyTorch、transformers、safetensors、accelerate 这几个库。Python 版本建议 3.9 以上，PyTorch 建议 2.0 以上，因为新版本对显存管理和 CUDA 图的支持更好。CUDA 版本要和你显卡驱动匹配，这个不用多说。

安装 airllm 本身很简单，一条命令：

pip install airllm

但这里有个坑：airllm 会默认安装它依赖的 transformers 版本，可能和你环境里已有的版本冲突。如果你之前装过其他推理框架，建议用虚拟环境隔离，避免依赖打架。我一般用 conda 建一个干净环境：

conda create -n airllm_env python=3.10 conda activate airllm_env pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install airllm

装完之后验证一下：

import airllm print(airllm.__version__)

能打印出版本号就说明装好了。如果报 CUDA 相关的错，检查一下 PyTorch 能不能正常调用 GPU：

import torch print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))

这两行输出True和你的显卡型号，才算环境就绪。

4.2 模型分片转换的完整流程

环境好了之后，第一步是把模型转换成分片格式。airllm 提供了AutoModel类，加载时会自动处理分片。但第一次加载原始模型时，它会先做转换，这个过程比较慢，而且需要足够的磁盘空间（大约是模型大小的 1.2 倍，因为转换过程中原始权重和分片会同时存在）。

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

这段代码做了几件事：下载原始模型（如果本地没有）、按层切分、压缩存储到指定路径。compression="4bit"表示分片用 4-bit 压缩存储，加载时解压回 FP16。转换完成后，./shards/llama70b目录下会有一堆 safetensors 文件，每个对应若干层。

转换时间取决于模型大小和磁盘速度。7B 模型大概 2 到 5 分钟，70B 模型 15 到 30 分钟。转换过程中可以去看日志，它会打印每一层的处理进度。如果中途中断了，重新跑会从断点继续，不用从头来。

注意：转换时确保磁盘剩余空间足够。70B 模型 FP16 原始权重 140GB，分片压缩后可能 40GB 左右，但转换过程中两者并存，峰值需要 180GB 以上。空间不够会直接失败。

4.3 推理调用的代码实现

分片准备好之后，推理就简单了。airllm 的接口设计得和 transformers 很像，基本可以无缝迁移：

from airllm import AutoModel import torch model = AutoModel.from_pretrained( "meta-llama/Llama-2-70b-hf", layer_shards_saving_path="./shards/llama70b", compression="4bit" ) input_text = "请用一句话解释什么是机器学习。" input_tokens = model.tokenizer( input_text, return_tensors="pt", return_attention_mask=False, truncation=True, max_length=128 ).input_ids.cuda() generation_output = model.generate( input_tokens, max_new_tokens=64, use_cache=True, return_dict_in_generate=True ) output = model.tokenizer.decode(generation_output.sequences[0]) print(output)

这段代码里几个参数值得说明。max_new_tokens=64控制生成的最大长度，别设太大，因为每个 token 都要走一遍全部层，64 个 token 就是 64 次完整的前向传播，时间会线性增长。use_cache=True开启 KV Cache，这样生成第二个 token 时不用重新计算前面所有 token 的 KV，能省不少时间。truncation=True和max_length=128限制输入长度，防止输入太长把显存撑爆。

实测下来，第一次调用会慢一些，因为要加载分片、初始化 CUDA 上下文。后续调用会快一点，因为部分分片可能还在系统缓存里。但别指望有质的飞跃，分层加载的本质决定了它快不起来。

4.4 性能调优的实操记录

我在一台配置为 8GB 显存显卡、32GB 内存、NVMe SSD 的机器上跑过 70B 模型，记录了一些调优过程。初始配置下，单 token 生成时间 7.2 秒，打开 profiling 发现 IO 占了 5.8 秒，计算只占 1.4 秒。明显 IO 是瓶颈。

第一个优化是开压缩。把compression从None改成"4bit"，分片体积降到原来的四分之一，IO 时间降到 2.1 秒，但解压消耗了 0.8 秒 CPU，净收益是单 token 降到 4.5 秒。这个优化是划算的。

第二个优化是调预读取。airllm 默认的预读取深度是 1，也就是只预读下一层。我改成 2，让它提前读两层，IO 和计算重叠得更充分。单 token 降到 3.8 秒。但预读取深度不能太大，因为预读的层要占内存，深度太大内存会爆。32GB 内存下，深度 2 是安全的，深度 3 就开始吃紧了。

第三个优化是限制序列长度。原来我设的max_length=512，KV Cache 占了 1.5GB 显存，导致单层权重加载时显存不够，触发了显存和内存之间的换页，反而更慢。改成max_length=128后，KV Cache 降到 0.4GB，显存充裕了，单 token 降到 3.2 秒。

最终这套配置下，生成 64 个 token 大约 3.5 分钟。这个速度做离线验证是够用的，做交互就不行了。

5. 常见问题与排查技巧实录

5.1 加载失败与显存不足的排查

最常见的问题就是加载到一半报显存不足。这个问题的根源往往不是单层权重太大，而是显存碎片。PyTorch 的缓存分配器在反复加载释放之后，会留下很多小块空闲显存，虽然总量够，但没有一块连续的大块能放下单层权重。解决办法是在加载前设置环境变量：

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

这个配置让分配器把大块切小，减少碎片。另一个办法是每次推理前手动清缓存：

import torch torch.cuda.empty_cache()

但empty_cache有开销，别每层都调，在每次 generate 之前调一次就行。

如果确认显存总量确实不够，那就得降配置。把compression开到"4bit"，把max_length降到 64，把max_new_tokens降到 32，一步步试，找到能跑通的组合。别一上来就追求高配置，先跑通再优化。

5.2 推理结果异常的处理

有时候模型能跑通，但输出是乱码或者重复。这种情况我遇到过几次，原因各不相同。第一次是分片转换不完整，某个分片文件损坏了，加载时没报错但权重是错的。解决办法是删掉分片目录重新转换，转换完检查文件数量是否和预期一致。

第二次是 tokenizer 配置不对。airllm 默认用模型自带的 tokenizer，但有些模型的 tokenizer 需要额外配置，比如设置padding_side或者add_special_tokens。如果输出全是特殊 token，检查一下 tokenizer 的配置。

第三次是精度问题。虽然 airllm 本身不损失精度，但如果你开了压缩，解压过程中如果有 bug，可能导致数值异常。这时候把compression关掉，用原始精度跑一遍，如果正常了，说明是压缩的问题，换个压缩算法或者升级 airllm 版本。

5.3 速度慢到无法接受的优化方向

如果速度慢到你觉得不可用，先别急着放弃，按下面的顺序排查。第一步，确认磁盘类型。用lsblk或者df -h看模型分片在哪个盘上，如果是机械盘，换到 SSD 上，速度能提升 5 到 10 倍。第二步，确认 CPU 不是瓶颈。开压缩的情况下，解压是 CPU 密集的，如果 CPU 太老，解压时间可能比 IO 还长。用htop看推理时 CPU 占用，如果一直 100%，说明 CPU 是瓶颈，关掉压缩反而更快。

第三步，确认显卡驱动和 CUDA 版本匹配。驱动太老会导致 CUDA 内核启动慢，间接拖慢整体速度。第四步，考虑换更小的模型。70B 跑不动，试试 13B 或者 7B，速度会快很多，虽然效果差一些，但至少能用。我个人的经验是，在低资源场景下，选对模型比优化推理更重要。一个 13B 的模型在 airllm 上跑，体验可能比 70B 好得多。

5.4 长期运行的稳定性注意事项

如果你打算让 airllm 长时间跑批处理任务，有几个稳定性问题要注意。第一是内存泄漏。airllm 在反复加载释放分片的过程中，Python 的垃圾回收可能跟不上，导致内存缓慢增长。解决办法是每隔一段时间重启进程，或者手动调gc.collect()。我一般每跑 100 次推理重启一次，能避免大部分泄漏问题。

第二是磁盘温度。NVMe SSD 在持续高负载读取下会发热，温度过高会降速。如果你的任务要跑几个小时，给 SSD 加个散热片，或者中间留点间隔让盘凉一凉。第三是CUDA 上下文丢失。长时间运行后，偶尔会遇到 CUDA 上下文丢失的报错，这通常是驱动或者硬件的问题，重启进程能恢复。如果频繁出现，检查一下显卡散热和电源功率。

提示：批处理任务建议写成脚本，加上重试逻辑。遇到 CUDA 错误就重启进程，从上次中断的地方继续，别让一个错误毁掉整批任务。

6. 这套方案还能怎么扩展

airllm 的代码结构不复杂，核心就是"分层加载 + 显存复用"，在这个基础上可以做不少扩展。我试过的一个方向是多卡并行。airllm 默认是单卡，但你可以把不同层分配到不同卡上，每张卡负责一部分层的加载和计算，层与层之间通过 PCIe 或者 NVLink 传激活值。这样显存总量上去了，单卡压力小了，但卡间通信会成为新瓶颈。如果卡间是 PCIe 4.0 x16，带宽 32GB/s，传激活值（通常几 MB 到几十 MB）开销不大，是可行的。

另一个方向是和量化叠加。前面提过，airllm 加载的是 FP16 分片，如果你先用 GPTQ 或者 AWQ 把模型量化到 INT4，再保存成分片，单层权重会小很多，加载更快，显存占用更低。但量化后的模型结构可能和原始模型不完全一致，airllm 的自动分片逻辑需要调整。这个我还在摸索，目前的做法是手动指定层映射关系，把量化模型的层对应到 airllm 的分片结构上。

还有一个方向是缓存热点层。观察发现，模型的前几层和最后几层被调用的频率最高（因为每次生成都要走完整流程），中间层相对固定。可以把前几层和最后几层常驻显存，只对中间层做分层加载。这样能省掉一部分 IO，速度有提升。airllm 没有直接支持这个，但可以通过修改加载逻辑实现，把常驻层的分片预加载到显存，加载时跳过。

最后说个实际体会。airllm 这类方案的价值，不在于它跑得多快，而在于它把"能不能跑"的门槛降到了极低。我见过太多人因为显存不够就放弃了本地跑大模型的想法，其实只要愿意接受慢一点，办法是有的。airllm 就是这样一个"兜底"方案，它不完美，但在资源受限的场景下，它给了你一个可用的选择。如果你手头只有一张小显存卡，又想折腾大模型，不妨试试这个思路，踩几个坑之后，你会对推理的内存管理有更深的理解。