☰
单卡24GB显存跑70B大模型:airllm分层加载推理实战
2026/10/15 5:42:25 网站建设 项目流程

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 airllm

PyTorch 的版本要和你的 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 OOMKV 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 不同规模模型的配置参考

模型规模最低显存推荐显存最低内存推荐内存存储要求
7B6GB8GB16GB32GBSATA SSD
13B8GB12GB16GB32GBSATA SSD
30B12GB16GB32GB64GBNVMe SSD
70B16GB24GB32GB64GBNVMe SSD
130B24GB48GB64GB128GBNVMe 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 和内存,不在计算。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询