资讯动态

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

发布时间:2026/10/17 15:11:53 来源:尧图企业网站定制
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 到 2GFP16加上激活和 cache24G 卡完全够用。这个思路的代价是推理速度。每层都要从磁盘读一次磁盘 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 parallelairllm 基本就是 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 层时会异步加载第 N1 层的权重。这样当第 N 层算完第 N1 层已经准备好了不需要等待磁盘 IO。预读取的深度可以配置一般设为 1 到 2 层。设得太深会占用额外显存设得太浅又起不到隐藏 IO 延迟的作用。磁盘 IO 方面建议把模型分片文件放在 SSD 上。机械硬盘的随机读性能太差每层加载可能要几百毫秒推理速度会掉到每秒零点几个 token。NVMe SSD 的随机读能到几十微秒级别基本不会成为瓶颈。我用 SATA SSD 和 NVMe SSD 对比过NVMe 下 70B 模型能到 2 到 3 token/sSATA SSD 只有 1 token/s 左右。3.3 关键参数layers_per_load 与 prefetchlayers_per_load控制每次加载多少层。默认是 1也就是逐层加载。如果显存有富余可以设为 2 或 3一次加载多层减少加载次数提升速度。但设得太大就退化成全量加载了显存会爆。建议从 1 开始试观察显存占用逐步往上加找到速度和显存的平衡点。prefetch控制预读取的层数。设为 1 表示提前加载下一层设为 2 表示提前加载下两层。预读取会占用额外显存因为加载进来的层要暂存。如果显存紧张可以把 prefetch 设为 0关闭预读取但速度会明显下降。我的经验是24G 卡跑 70B 模型layers_per_load1、prefetch1 是比较稳的配置。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 airllmCUDA 版本根据你的驱动来选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, compression4bit )第一次加载时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, compression4bit, layers_per_load1, prefetch1 ) input_text 请用一句话解释什么是机器学习。 input_ids model.tokenizer(input_text, return_tensorspt).input_ids.cuda() output model.generate( input_ids, max_new_tokens128, do_sampleTrue, temperature0.7, top_p0.9 ) print(model.tokenizer.decode(output[0], skip_special_tokensTrue))compression参数可以选4bit或8bit进一步降低单层权重的显存占用。4bit 压缩后70B 模型单层大约 0.5G24G 卡可以开更大的 layers_per_load。但压缩会带来精度损失如果对生成质量要求高建议用 FP16 不压缩。4.4 显存与速度的实测数据我在一张 24G 显存的卡上跑了几个配置数据如下模型规模压缩方式layers_per_loadprefetch显存峰值生成速度70BFP161118G1.5 token/s70B4bit1112G2.2 token/s70B4bit2116G2.8 token/s70B4bit2220G3.1 token/s13BFP16118G6 token/s从数据看4bit 压缩对速度有提升因为单层权重更小加载更快。layers_per_load 和 prefetch 增加也能提速但显存占用会上升。24G 卡跑 70B 模型4bit layers_per_load2 prefetch1 是比较均衡的配置显存留有余量速度也还行。5. 实操中踩过的坑与排查技巧5.1 磁盘空间不足导致拆分失败第一次跑 70B 模型时我没注意磁盘空间拆分到一半报错退出。分层拆分后的文件总大小和原始模型差不多70B FP16 大约 140G4bit 压缩后大约 35G。拆分前一定要确认目标磁盘有足够空间建议预留原始模型 1.5 倍的空间因为拆分过程中原始文件和分片文件会同时存在。如果空间不够可以先把原始模型删掉再拆分但这样就没法回退了。更稳妥的做法是挂一块大容量 SSD 专门放模型分片拆分完再把原始文件清理掉。5.2 显存碎片导致的 OOM有时候显存总量够但还是报 OOM这通常是显存碎片导致的。PyTorch 的缓存分配器会预留一些显存频繁加载释放容易产生碎片。解决办法是设置PYTORCH_CUDA_ALLOC_CONFmax_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 太高会导致输出随机性大看起来像质量下降。建议先用 temperature0.7、top_p0.9 跑一组基线再逐步调整。如果任务需要确定性输出可以用 greedy searchtemperature 设为 0。5.4 常见问题速查表问题现象可能原因解决方法拆分时报磁盘空间不足目标磁盘剩余空间不够清理磁盘或换更大分区推理时 OOM显存碎片或 KV cache 累积设置 max_split_size_mb每次请求后清 cache生成速度极慢磁盘 IO 瓶颈换 NVMe SSD开启 prefetch输出重复或无意义压缩损失或采样参数不当换 FP16 或 8bit调整 temperature加载模型时报格式错误模型文件不完整或版本不匹配重新下载模型确认 transformers 版本多请求并发时崩溃显存不足或线程冲突串行处理请求避免并发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_load1。再大就得不偿失了不如用 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跑通流程确认环境没问题再上大模型。磁盘一定要用 SSDNVMe 最好。显存 24G 是跑 70B 的底线再低就得用更激进的压缩质量会明显下降。参数从 layers_per_load1、prefetch1 开始调观察显存和速度的变化找到适合自己硬件的平衡点。还有一个细节airllm 的版本更新比较快不同版本的 API 可能有差异。遇到问题先看项目 README 和 issue很多坑别人已经踩过了。如果 README 没写清楚可以去翻源码核心逻辑不复杂读一遍基本就明白它的工作机制了。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑