1. 单卡跑大模型这件事,到底卡在哪
第一次看到 airllm 这个项目名的时候,我正被显存问题折磨得够呛。手头只有一张 24GB 显存的卡,想跑一个 70B 参数量的模型做推理测试,结果加载到一半就爆显存,进程直接被系统干掉。相信很多做模型推理的朋友都遇到过类似的场景:模型权重文件动辄几十上百 GB,而消费级显卡的显存就那么点,想跑大模型要么租多卡机器,要么眼睁睁看着 OOM 报错。
airllm 这个项目解决的正是这个核心矛盾——用单张显卡、有限显存,跑起远超显存容量的大模型推理。它的思路不是去压缩模型精度或者做量化蒸馏,而是从推理过程的显存占用机制入手,把"一次性全部加载"改成"按需分层加载"。这个项目适合所有手头显卡显存不够、但又想本地跑大模型做推理验证的开发者,尤其是做模型效果测试、Prompt 调试、小批量推理任务的人。它不需要你懂分布式训练,也不需要多卡环境,一张卡就能跑起来。
我花了大概两周时间把这个项目的核心逻辑摸了一遍,从源码结构到实际跑通 70B 模型的完整流程都走了一遍,中间踩了不少坑,也总结了一些文档里没写的经验。下面我把整个项目的设计思路、核心机制、实操步骤和避坑要点完整拆解一遍,尽量让不同基础的朋友都能看懂、能复现。
2. 项目整体设计与核心思路拆解
2.1 为什么"分层加载"能解决显存瓶颈
要理解 airllm 的设计,得先搞清楚大模型推理时显存到底被什么占满了。一个 Transformer 结构的模型,推理时的显存占用主要分三块:模型权重、KV Cache、中间激活值。其中模型权重是大头,比如一个 70B 参数的模型,如果用 FP16 存储,光权重就要 140GB 左右,这还没算推理过程中的额外开销。
传统做法是把整个模型加载到显存里再推理,显存不够就没办法。airllm 的核心思路是:Transformer 是逐层计算的,第 N 层的计算只依赖第 N-1 层的输出,不需要同时把所有层都放在显存里。所以完全可以只把当前正在计算的那一层加载进显存,算完就释放,再加载下一层。这样一来,显存占用就从"整个模型的大小"降到了"单层模型的大小"。
这个思路听起来简单,但实现起来有几个关键问题要解决。第一是加载速度,如果每算一层都要从磁盘读一次权重,IO 会成为瓶颈;第二是层与层之间的数据传递,上一层的输出要保留在显存里传给下一层;第三是 KV Cache 的管理,自回归生成时 KV Cache 会随着生成长度增长。airllm 对这几个问题都做了针对性的处理,后面会详细讲。
2.2 方案选型:为什么不用量化或蒸馏
市面上解决显存不足的方案其实不少,量化和蒸馏是最常见的两条路。量化是把 FP16 权重压成 INT8 或 INT4,显存直接减半甚至减到四分之一;蒸馏是训练一个小模型去模仿大模型的行为。airllm 没有走这两条路,而是选择了分层加载,这个选择背后有它的道理。
量化的优势是显存占用小、推理速度快,但代价是精度损失。对于某些对精度敏感的任务,量化后的模型效果会明显下降,而且量化本身也需要额外的校准过程。蒸馏的问题更明显——你得先有一个大模型,还得有足够的算力和数据去训练小模型,成本很高,而且蒸馏出来的模型和原模型终究不是一回事。
分层加载的好处是精度无损,你跑的就是原始权重的模型,输出结果和全量加载完全一致。代价是推理速度会慢一些,因为多了磁盘 IO 和层加载的开销。但对于那些"显存不够但又想要原汁原味结果"的场景,这个取舍是值得的。airllm 的定位很清晰:不是追求极致速度,而是追求在有限硬件上跑起大模型的可能性。
2.3 核心模块划分与数据流
从代码结构上看,airllm 主要分几个模块:模型加载器、层管理器、推理调度器、KV Cache 管理器。模型加载器负责从磁盘读取权重文件并做格式转换;层管理器负责按需加载和释放每一层的权重;推理调度器负责协调整个前向传播过程;KV Cache 管理器负责在生成过程中维护注意力机制的键值缓存。
数据流是这样的:输入 token 经过 embedding 层后,进入第一层 Transformer,层管理器把第一层权重加载进显存,计算完成后释放权重但保留输出;输出传给第二层,重复这个过程直到最后一层;最后一层的输出经过 LM Head 得到 logits,采样出下一个 token;然后进入下一轮生成,此时 KV Cache 会复用之前计算的结果,避免重复计算。
这个流程里最关键的优化点是权重预加载。如果严格按"用完就释放、要用再加载"来做,每层计算前都要等磁盘 IO,速度会非常慢。airllm 的做法是在计算当前层的时候,异步预加载下一层的权重,这样等当前层算完,下一层已经准备好了,IO 和计算可以重叠起来。
3. 核心细节解析与实操要点
3.1 权重文件的组织与读取策略
airllm 对权重文件的组织方式有特定要求。它需要模型权重是按层分开存储的,这样才能做到按层加载。如果你拿到的是一个整合在一起的权重文件,需要先做拆分。项目里提供了拆分脚本,可以把 Hugging Face 格式的模型权重按层拆成单独的文件。
拆分后的目录结构大概是这样的:每一层对应一个文件,文件名里包含层号,embedding 层和 LM Head 单独存放。这样做的好处是加载时可以精确地只读需要的那一层,不用把整个文件读进来再截取。我实测下来,拆分后的文件数量会比较多,一个 70B 模型大概有 80 多个层文件,加上 embedding 和 LM Head,总共近百个文件。
注意:拆分权重的时候要确保文件完整性,如果某个层文件损坏或缺失,推理到那一层就会报错。建议拆分完成后做一次校验,对比拆分前后权重的哈希值。
读取策略上,airllm 用了内存映射的方式。内存映射的好处是不需要把整个文件读进内存,而是按需读取文件的不同部分。对于大文件来说,这能显著降低内存占用。不过内存映射对文件系统的要求比较高,在某些网络文件系统上可能会有性能问题,建议把权重文件放在本地 SSD 上。
3.2 层加载与释放的时机控制
层加载和释放的时机控制是 airllm 性能的关键。加载太早会占用显存,加载太晚会让计算等 IO。airllm 的策略是双缓冲:维护两个层的显存空间,一个用于当前计算,一个用于预加载下一层。当前层算完后,释放当前层,把预加载的层切换为当前层,同时开始预加载再下一层。
这个策略的显存占用是两层权重的空间,比单层多一倍,但换来的是 IO 和计算的重叠。对于显存特别紧张的场景,可以关掉预加载,改成单层模式,但速度会明显下降。我建议如果显存允许,还是开着预加载,体验会好很多。
释放时机上有个细节要注意:释放层权重的时候要确保这一层的计算已经完全结束,包括反向传播(如果做训练的话)。airllm 主要面向推理场景,所以只需要考虑前向传播。但即便如此,也要确保 CUDA 的异步操作已经完成,否则可能出现释放了还在用的显存的情况。项目里用了 CUDA 的同步机制来保证这一点。
3.3 KV Cache 的显存管理
KV Cache 是自回归生成时的显存大户。每生成一个 token,都要把当前 token 的 Key 和 Value 缓存起来,供后续 token 的注意力计算使用。生成长度越长,KV Cache 越大。对于一个 70B 模型,如果生成 2048 个 token,KV Cache 可能就要占好几 GB 显存。
airllm 对 KV Cache 的管理有几个策略。首先是分层存储,每一层的 KV Cache 单独管理,这样在层切换时可以精确控制哪些 Cache 需要保留。其次是动态扩展,KV Cache 的显存按需分配,不是一开始就分配最大长度。最后是可选的 CPU 卸载,如果显存实在不够,可以把部分 KV Cache 放到 CPU 内存里,需要时再换回显存。
实操心得:KV Cache 的 CPU 卸载会显著降低生成速度,因为每次注意力计算都要在 CPU 和 GPU 之间传输数据。如果显存够用,尽量不要开这个选项。我试过在 24GB 卡上跑 70B 模型,生成 512 个 token 以内,KV Cache 放显存里是够的,超过这个长度就要考虑卸载了。
3.4 精度与性能的平衡取舍
airllm 默认用 FP16 精度做推理,这是精度和性能的平衡点。FP16 相比 FP32 显存减半,速度更快,精度损失在大多数场景下可以接受。如果你的显卡支持 BF16,也可以用 BF16,动态范围更大,不容易出现溢出问题。
但 FP16 有个坑要注意:某些模型在 FP16 下会出现数值不稳定,表现为生成结果乱码或者重复。这通常是因为模型权重里有超出 FP16 表示范围的数值。解决办法是用 BF16,或者对特定层做 FP32 保留。airllm 提供了混合精度的选项,可以对 embedding 层和 LM Head 用 FP32,中间层用 FP16,这样能兼顾稳定性和显存占用。
性能方面,分层加载的额外开销主要来自磁盘 IO。我实测下来,用 NVMe SSD 的话,70B 模型每生成一个 token 大概比全量加载慢 30% 到 50%。如果用 SATA SSD,慢一倍以上。如果用机械硬盘,基本没法用,IO 会成为绝对瓶颈。所以硬件配置上,SSD 是必须的,NVMe 更好。
4. 实操过程与核心环节实现
4.1 环境准备与依赖安装
先把环境搭起来。airllm 是 Python 项目,依赖 PyTorch 和 Transformers 等库。我建议用 conda 建一个独立环境,避免和系统里的其他 Python 包冲突。
conda create -n airllm python=3.10 conda activate airllm pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 pip install transformers accelerate sentencepiece pip install airllmPyTorch 的版本要和你的 CUDA 版本匹配。我用的 CUDA 11.8,所以装的是 cu118 版本的 PyTorch。如果你用的是 CUDA 12.x,把 index-url 换成对应的版本。装完之后可以跑一下python -c "import torch; print(torch.cuda.is_available())"确认 GPU 能被识别。
airllm 本身依赖不多,主要是 transformers 和 accelerate。但实际跑的时候可能还需要 bitsandbytes 做量化支持(如果你要用量化的话),以及 safetensors 做权重加载。这些可以按需安装。
4.2 模型权重的拆分与准备
airllm 需要按层拆分的权重文件。如果你从 Hugging Face 下载的是整合权重,需要先拆分。项目里提供了拆分工具,用法大概是这样的:
from airllm import AutoModel model = AutoModel.from_pretrained( "meta-llama/Llama-2-70b-hf", layer_shards_saving_path="./llama70b_shards", profiling_mode=False )这段代码会下载模型权重并自动拆分成按层存储的文件。拆分过程比较耗时,70B 模型大概要跑十几分钟到半小时,取决于磁盘速度。拆分完成后,./llama70b_shards目录下就是按层组织的权重文件。
注意:拆分需要足够的磁盘空间,大概是原始模型大小的 1.1 倍左右。70B 模型 FP16 大概 140GB,拆分后需要 150GB 以上的空闲空间。提前清理好磁盘。
如果你已经有拆分好的权重,可以直接指定路径加载,跳过拆分步骤。拆分后的权重可以复用,不用每次推理都重新拆。
4.3 推理代码的编写与参数配置
加载模型和推理的代码其实很简洁,airllm 把复杂性都封装在内部了。一个基本的推理脚本大概长这样:
from airllm import AutoModel import torch model = AutoModel.from_pretrained( "./llama70b_shards", max_seq_len=512, layer_shards_saving_path=None, compression=None, profiling_mode=False ) 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=128, use_cache=True, return_dict_in_generate=True ) output = model.tokenizer.decode(generation_output.sequences[0]) print(output)几个关键参数说明一下。max_seq_len是最大序列长度,决定了 KV Cache 的预分配大小,设太大浪费显存,设太小生成会被截断。compression是量化选项,设成 None 就是不做量化,用原始精度。profiling_mode打开后会打印每层的加载和计算耗时,调试性能时有用。
generate里的max_new_tokens控制生成的最大 token 数,use_cache控制是否使用 KV Cache。建议开着 use_cache,不然每生成一个 token 都要重新计算整个序列,速度会慢很多。
4.4 实际跑通 70B 模型的完整记录
我在一张 24GB 显存的卡上跑 Llama-2-70B,完整记录一下过程和观察到的现象。模型拆分花了大概 25 分钟,拆分后的权重目录有 150GB 左右。加载模型的时候,airllm 先加载了配置文件和 tokenizer,然后开始按层加载权重。第一层加载的时候有明显等待,后面因为预加载机制,等待感就不明显了。
生成第一个 token 大概花了 3 秒左右,后续每个 token 大概 0.8 到 1.2 秒。生成 128 个 token 总共用了不到 3 分钟。显存占用峰值在 18GB 左右,主要是两层权重加 KV Cache。CPU 内存占用在 8GB 左右,主要是内存映射的文件缓存。
对比一下全量加载:同样一张卡,全量加载 70B 模型直接 OOM,根本跑不起来。所以 airllm 的价值就在于让"跑不起来"变成了"能跑起来",虽然慢一些,但至少能用了。
实操心得:生成速度受磁盘 IO 影响很大。我一开始把权重放在机械硬盘上,每个 token 要 5 秒以上,换成 NVMe SSD 后降到 1 秒左右。如果要做交互式调试,建议把权重放在最快的盘上。
5. 常见问题与排查技巧实录
5.1 显存溢出与内存不足的排查
跑 airllm 最常见的报错就是 CUDA out of memory。虽然 airllm 的设计就是为了省显存,但如果配置不当还是会 OOM。排查思路是这样的:先看显存占用峰值出现在哪个阶段,如果是加载权重时 OOM,说明预加载的层数太多,可以关掉预加载或者减少预加载层数;如果是生成过程中 OOM,说明 KV Cache 太大,可以减小 max_seq_len 或者开启 KV Cache 的 CPU 卸载。
还有一个容易被忽略的点是 PyTorch 的显存缓存机制。PyTorch 默认会缓存已分配的显存,不会立即还给系统。这会导致显存看起来一直占用很高,但实际上可以复用。如果遇到 OOM,可以试试torch.cuda.empty_cache()手动清理缓存。不过频繁调用这个函数会影响性能,建议只在关键时刻用。
CPU 内存不足也会导致问题。airllm 用内存映射读权重文件,操作系统会把文件内容缓存在内存里。如果物理内存不够,系统会频繁换页,速度急剧下降。建议至少 32GB 内存,70B 模型最好 64GB 以上。
5.2 生成结果异常的原因分析
生成结果异常通常表现为乱码、重复、截断或者答非所问。乱码最常见的原因是精度问题,FP16 下的数值溢出会导致输出异常。解决办法是换 BF16 或者对敏感层用 FP32。重复问题通常是采样参数设置不当,可以调整 temperature 和 repetition_penalty。
截断问题一般是 max_seq_len 或者 max_new_tokens 设得太小。检查一下输入长度加上生成长度是否超过了 max_seq_len。如果超了,要么增大 max_seq_len(需要更多显存),要么截断输入。
答非所问可能是模型本身的问题,也可能是 tokenizer 配置不对。检查一下 tokenizer 的 special tokens 设置,确保 padding 和 eos token 配置正确。有些模型需要特定的 prompt 格式,格式不对会导致效果很差。
5.3 速度慢的优化方向
速度慢是分层加载的固有代价,但可以通过一些手段优化。首先是硬件层面,用 NVMe SSD 替代 SATA SSD 或机械硬盘,IO 速度能提升好几倍。其次是软件层面,开启预加载让 IO 和计算重叠,调整预加载的层数找到平衡点。
还有一个优化点是减少精度转换的开销。airllm 加载权重时可能需要做精度转换(比如从 FP32 转 FP16),这个转换有计算开销。如果权重文件本身就是 FP16 格式,可以跳过转换,加载更快。
批处理也能提升吞吐量。虽然单次生成的速度不变,但一次处理多个请求可以摊薄层加载的开销。airllm 支持 batch 推理,但要注意 batch size 增大会增加 KV Cache 的显存占用,需要权衡。
5.4 常见问题速查表
| 问题现象 | 可能原因 | 排查方法 | 解决方案 |
|---|---|---|---|
| CUDA OOM | 预加载层数过多 | 查看显存峰值出现阶段 | 关闭预加载或减少层数 |
| CUDA OOM | KV Cache 过大 | 检查 max_seq_len 设置 | 减小 max_seq_len 或开启 CPU 卸载 |
| 生成乱码 | FP16 数值溢出 | 换 BF16 测试 | 使用 BF16 或混合精度 |
| 生成重复 | 采样参数不当 | 检查 temperature 和 penalty | 调整采样参数 |
| 生成截断 | 序列长度限制 | 检查输入输出总长度 | 增大 max_seq_len |
| 速度极慢 | 磁盘 IO 瓶颈 | 查看磁盘类型和占用 | 换 NVMe SSD |
| CPU 内存不足 | 物理内存不够 | 查看内存占用和换页 | 增加内存或减小模型 |
| 加载报错 | 权重文件损坏 | 校验文件哈希 | 重新拆分或下载权重 |
避坑技巧:第一次跑的时候先用小模型(比如 7B)验证流程,确认环境没问题再上大模型。大模型拆分和加载都很耗时,流程有问题的话浪费时间。
6. 工具选型与硬件配置建议
6.1 显卡选择的考量因素
airllm 对显卡的要求主要是显存容量和 CUDA 支持。显存越大,能预加载的层数越多,速度越快。24GB 显存可以比较流畅地跑 70B 模型,16GB 会紧张一些,8GB 只能跑小模型或者开各种省显存选项。
显卡的计算能力也有影响,但相比显存,计算能力不是瓶颈。因为分层加载的瓶颈在 IO,GPU 计算单元大部分时间在等数据。所以一张计算能力中等但显存大的卡,可能比计算能力强但显存小的卡更适合 airllm。
我试过在 3090(24GB)和 4090(24GB)上跑,速度差异不大,因为瓶颈都在 IO。如果用 A100(80GB),可以预加载更多层,速度会快一些,但性价比不高。对于个人开发者,24GB 的消费级卡是比较合适的选择。
6.2 存储配置的优化建议
存储是 airllm 性能的关键。权重文件要放在 SSD 上,NVMe 最好。我实测过几种存储介质的速度差异:NVMe SSD 每 token 约 1 秒,SATA SSD 约 2 秒,机械硬盘约 5 秒以上。这个差异在长文本生成时会非常明显。
如果权重文件太大,SSD 放不下,可以考虑用外置 NVMe 硬盘盒。USB 3.2 Gen2 的接口带宽大概 10Gbps,比内置 NVMe 慢一些,但比 SATA 快。雷电接口的外置硬盘速度接近内置,但价格贵。
内存方面,建议至少 32GB,70B 模型建议 64GB 以上。内存的作用是缓存权重文件,减少磁盘读取次数。内存越大,缓存命中率越高,速度越稳定。
6.3 不同规模模型的配置参考
| 模型规模 | 最低显存 | 推荐显存 | 最低内存 | 推荐内存 | 存储要求 |
|---|---|---|---|---|---|
| 7B | 6GB | 8GB | 16GB | 32GB | SATA SSD |
| 13B | 8GB | 12GB | 16GB | 32GB | SATA SSD |
| 30B | 12GB | 16GB | 32GB | 64GB | NVMe SSD |
| 70B | 16GB | 24GB | 32GB | 64GB | NVMe SSD |
| 130B | 24GB | 48GB | 64GB | 128GB | NVMe SSD |
这个表是基于 FP16 精度、开启预加载、生成 512 token 以内的场景估算的。如果生成长度更长或者 batch size 更大,显存需求会相应增加。
7. 个人实操体会与后续扩展方向
用了这段时间,我最大的体会是 airllm 这类工具的价值不在于性能,而在于降低了硬件门槛。以前想跑 70B 模型,要么租多卡云服务器,要么买专业卡,成本很高。现在一张消费级卡就能跑起来,虽然慢一些,但对于做效果验证、Prompt 调试、小批量推理这些场景,完全够用了。
踩过的坑里,最值得说的是存储配置。我一开始图省事把权重放在机械硬盘上,结果慢到怀疑人生,后来换了 NVMe SSD 才体会到正常速度。所以如果你打算认真用这个工具,存储上的投入是值得的,一块 1TB 的 NVMe SSD 现在也不贵。
后续扩展方向有几个。一是支持更多模型架构,目前 airllm 主要支持 Llama 系和类 Llama 结构,其他架构的适配还在完善中。二是优化 KV Cache 管理,比如用 PagedAttention 的思路做更细粒度的显存管理。三是支持多卡并行,把不同层分配到不同卡上,进一步提升速度。
如果你也在用类似方案跑大模型,建议先从 7B 或 13B 模型入手,把流程跑通,再逐步上更大的模型。遇到问题多看看项目的 issue 区,很多坑别人已经踩过了。硬件上,SSD 和内存的投入比显卡更划算,因为瓶颈在 IO 和内存,不在计算。