首页 / 资讯中心 / 文章详情

AirLLM实战:把70B大模型塞进24GB显存

AirLLM实战:把70B大模型塞进24GB显存 ★ FEATURED ARTICLE
最近我一直在折腾一件事把参数量达到70B级别的开源大模型放到一块只有24GB显存的普通显卡上跑推理。听上去像天方夜谭但用AirLLM确实做到了。这篇就聊聊它的原理、实操路径以及我踩过的各种坑。如果你手里只有消费级显卡又想体验大模型而不是到云上花钱租卡AirLLM是一个值得关注的项目。它不是把模型改小而是用更聪明的调度方式让模型“挤”进现有显存代价是速度变慢。但它解决了从0到1的问题先跑起来再谈优化。1. 一个看似不可能的推理需求1.1 大模型参数规模与显存需求的换算先算一笔最简单的账。假设一个模型有70B参数也就是700亿个参数。如果用16位浮点数保存每个参数占2个字节光是权重就要占掉约140GB的显存。算上推理过程中的中间激活值、KV Cache实际需要的内存峰值还会更高。消费级显卡的显存普遍是8GB、12GB或24GB哪怕是目前比较高端的产品线24GB也已经算很大了跟140GB比还差着数量级。所以常规思路下想跑70B模型就要上多张显卡做并行或者去租一台显存几十GB以上的服务器。对于个人开发者和业余爱好者来说这个门槛确实高。AirLLM解决的就是这个痛点让低显存设备也能跑超大参数模型。它的核心做法不复杂但实现细节相当有意思。1.2 常规方案的局限在聊AirLLM之前先说说大家通常怎么处理显存不足。第一个思路是模型并行把不同层放到不同显卡但普通机器没有那么多卡第二个思路是整体卸载把模型权重放到CPU内存每次计算只拷贝部分数据但工程上很难做到精细调度第三个思路是量化把16位参数压成8位甚至4位显存占用确实降下来了但仍存在容量天花板比如24GB显存跑70B四比特量化后的35GB权重依然放不下。量化可以和AirLLM结合使用但不能替代AirLLM。AirLLM的思路更接近一种“层级别的流水线”不把整个模型一次性放进显存而是每次只加载计算需要的若干层算完就走把显存让给下一层。这种思路把“模型需要140GB显存”这个硬指标拆解成了“模型每一层只需要几百MB显存”的可执行任务。2. 拆解AirLLM的核心思想2.1 为什么大模型可以“分层”跑想理解AirLLM得先理解Transformer的推理结构。无论是哪一代开源大模型底层结构基本都有大量重复的Transformer层。每一层做的事情是接收前一层的输出经过自注意力机制和前馈网络再产生新的输出给下一层。这里的关键是层与层之间是串行计算的而且是严格的先后依赖。第一层不计算完第二层就拿不到输入。既然如此就没必要同时把所有层的参数都驻留在显存里。AirLLM的做法就是按层做调度。举个不恰当但很直观的例子做一道需要10道工序的菜你不需要把十口锅同时摆满灶台只需要一口锅做完一道工序洗干净再做下一道。显存就是那口锅模型层就是工序。代价是切菜和洗锅需要时间速度会慢但灶台要求大大降低了。一个重要细节是Transformer采用了残差结构每一层的输出要跟输入相加后再传给下一层。这意味着中间激活值不能随便丢弃至少要保留完整的残差连接数据。AirLLM是把这些激活值放在CPU内存还是显存里有一套精细的取舍逻辑后面会讲。2.2 按层调度显存、内存、磁盘三级交换AirLLM把存储分成了三层显存、CPU内存、磁盘。在GPU推理的每个时间点显存里只需要驻留极少量数据当前正在计算的这一层权重、输入激活值以及必要的KV Cache。剩下的所有层权重都存放在CPU内存或者磁盘中。当计算到第N层时把第N层的权重从内存加载进显存算完后立刻释放再加载第N1层。这样显存压力不再取决于模型总参数而取决于“单层最大值”加上“激活值加上KV Cache”的大小。对于常见的70B模型单层参数在几十MB到几百MB之间远小于整模型。这也是24GB显存能跑起来的关键。CPU内存是第二层缓冲。如果内存足够大模型权重可以全部映射到内存里加载速度取决于内存带宽如果内存也不够AirLLM会利用操作系统把文件映射到磁盘借助虚拟内存机制实现第三层存储。这种场景下速度会明显下降因为要反复读硬盘但总比跑不起来强。这里要注意一个容易被忽略的点CPU内存和磁盘的占用也要提前评估。70B模型即使量化到4位权重也有35GB左右。所以我建议准备至少64GB内存再配合足够大的固态硬盘。实测下来内存不足导致的系统崩溃远比显存OOM更难受。2.3 用时间换显存到底值不值AirLLM不是没有代价最大代价就是速度。跨层调度需要反复传输数据GPU实际上有一多半时间在等数据真正计算的时间占比并不高。我在24GB显存上跑70B模型一个token的生成时间是秒级甚至更慢和云端多卡动辄每秒几十token的体验完全没法比。但你得想清楚自己的使用场景。如果是高并发在线服务别说AirLLM任何低显存方案都很难满足需求如果是个人调试、离线批量推理、研究模型输出或者只是好奇想体验一下那这个速度完全可以接受。我把它当成“穷人的实验室装备”用来研究模型行为、跑特定离线任务价值很大。3. 实操笔记把70B模型塞进24G显存3.1 环境准备和安装先把环境说清楚。我使用的是一台普通Linux工作站CPU内存64GB显卡显存24GB硬盘是NVMe固态。系统安装好显卡驱动以及对应的深度学习计算环境Python版本在3.8以上。安装AirLLM非常简单直接用包管理器拉取即可pip install airllm依赖的底层Python包会自动装上不需要额外手工配置太多东西。不过有一点要注意如果你的环境里已经装了深度学习框架的老版本AirLLM对版本有要求建议在干净环境中安装避免依赖冲突。模型文件需要下载到本地。AirLLM支持从常见的模型仓库加载权重也可以加载你自己下载好的本地目录。我更推荐先把模型下载到本地再跑因为第一次加载时会做格式转换和索引构建如果边下边跑非常容易中断。准备好模型目录后建议先跑一个小模型测试环境比如7B级别确认链路通畅后再切换到大模型。我见过不少朋友一上来就怼70B结果环境问题都没排查清楚最后很难判断是显存问题还是安装问题。3.2 最小推理示例下面这段代码是我在某个版本上实测跑通的注释都写在里面了。由于模型目录结构不同具体参数可能略有差异你拿到手要按实际路径调整。from airllm import AutoModel, AutoConfig # 这里换成你下载好的本地模型目录 model_path /data/models/your_70b_model config AutoConfig.from_pretrained(model_path) model AutoModel.from_pretrained(model_path, configconfig) # 用与模型配套的分词工具把文本变成token编号 # 这里我封装成load_tokenizer实际使用时用你常用的分词器即可 tokenizer load_tokenizer(model_path) prompt 用三句话解释什么是大模型推理 input_ids tokenizer(prompt).input_ids # 只生成一小段降低显存压力 output_ids model.generate( input_ids[input_ids], max_new_tokens32, do_sampleFalse, ) print(tokenizer.decode(output_ids[0]))跑起来之后你会看到终端里不断刷日志每加载一层都会有提示。正常情况下一开始会比较慢因为要加载和分配资源等真正开始输出之后会稳定在一个速度上。3.3 针对不同显存规模的调参建议我没有办法给出固定的配置因为这跟模型架构、量化方式、序列长度都有关系。但可以分享几个经验参考显存大小建议可尝试的模型规模注意事项8G7B左右配合四比特量化max_new_tokens控制在64以内不要开过长的序列12G10B到13B量化激活值和大batch容易溢出batch固定为116G20B到34B量化需要足够CPU内存做支撑24G70B级别量化后配AirLLM内存至少64G保证固态硬盘读写性能序列长度对显存占用的影响非常大尤其是KV Cache。KV Cache会随着生成token数量线性增长如果你要一次生成几百个token显存压力会迅速上升。建议在调试阶段把max_new_tokens调小先把流程跑通。4. 实操中躲不开的性能问题4.1 慢的底层原因数据搬运速度跟不上计算跑通了AirLLM之后最大的感受就是慢。这种慢不是没有原因的。大模型的单层计算量并不是特别大但参数权重只有在显存里才能被GPU高效访问。AirLLM需要把权重从内存搬到显存这个搬运通道的速度成了系统性能的命门。你可以理解成GPU是一个手速极快的厨师但食材放在仓库里每次只取一盘菜。厨师炒菜只要10秒取菜却要1分钟整条流水线就卡在取菜环节。PCIe带宽、内存带宽、磁盘带宽任何一环不够都会拖后腿。这也是为什么固态硬盘重要。模型权重如果要从磁盘读取机械硬盘几乎不能用固态硬盘比内存慢一个数量级但至少能接受。如果内存够大尽量让系统把权重都缓存到内存里减少磁盘访问。4.2 可以落地的6条提速技巧虽然速度上限受硬件限制但多数情况下并没有发挥到最优。下面几个技巧是我实际用下来比较有效的。第一量化一定要做。AirLLM支持加载量化后的权重量化把70B模型的传输量从140GB降到了35GB左右IO压力骤降速度提升非常明显。第二控制batch size。AirLLM是按层计算的batch越大每层算得越久中间激活值也越大。个人调试时batch直接设为1没必要追求batch性能。第三限制生成长度。KV Cache是显存杀手如果你只需要模型回答50个字就把max_new_tokens设为64而不是512。长文本生成过程中后半段会明显越来越慢。第四减少其他程序占用。AirLLM需要大量CPU内存做权重的临时存储浏览器开一堆页面、编译器后台跑任务都可能导致内存不足触发磁盘交换。跑推理时最好关掉不必要的进程。第五对输入序列做裁剪。输入太长会导致激活值累积显存压力随之增大。先做文本清洗和截断对日常任务影响不大。第六选用合适的调度策略。有些版本提供了自定义层组的配置意思是每轮加载的层数。加载层数越多层间切换次数越少但当前峰值显存越高。你可以试着把几层合并成一组加载找到一个峰值和速度的平衡点。4.3 如何判断瓶颈在磁盘、内存还是显存性能出问题时先别急着怪AirLLM先定位瓶颈。打开系统监控如果磁盘读取一直处于高位说明内存不够权重在反复走磁盘如果内存接近满但磁盘却很清闲说明权重已经在内存里而内存带宽不够如果显存使用率一直很高那说明当前层参数加上激活值已经压着显存极限需要减小batch或序列长度。AirLLM日志也会输出每层加载的时间如果加载耗时占比很高就该考虑换更好的硬盘或加内存。有一次我跑13B模型很慢查了半天发现是固态硬盘剩余空间太少写入缓存也受影响清理出几十GB空间后明显改善。5. 常见问题速查表我把这阵子遇到最多的问题整理成一张表大家可以对照着排查。问题现象可能原因解决办法启动加载时报错模型目录不完整或格式不支持确认权重完整用官方预转换格式或先转成AirLLM支持的格式显存OOM激活值或KV Cache过大减小max_new_tokensbatch设为1关闭其他显存占用程序系统内存耗尽内存不足以缓存全部权重增加物理内存或用虚拟内存兜底但速度会下降速度极慢CPU占用高权重频繁在磁盘和内存之间交换换成NVMe固态关掉占内存的程序尽量让权重驻留内存输出质量差量化力度太大或加载了不匹配的权重试不同量化级别或回退到高精度版本生成到一半崩溃KV Cache持续增长导致OOM限制生成长度或改用流式处理节省缓存这里有一个容易忽略的细节AirLLM并不是所有模型都支持它针对主流Transformer结构做过适配。如果加载某个模型报“不支持的架构”不要硬改配置大概率是这个库还没适配。先换同架构的其他开源模型或者关注版本更新。另外如果你的显存连一个小单层都装不下比如只有4G显存那即便用AirLLM也很难跑较大的模型因为单层参数加激活值也要占用空间。AirLLM把显存需求从整模型压缩到了单层级别但没有压缩到零。6. 关于AirLLM的一些真实体会把70B模型真正跑起来的那一刻我并没有觉得它快反而觉得整个过程像一场赌局盯着日志逐层加载生怕哪个环节内存爆掉。等第一条回答生成出来确实有种说不出的成就感。AirLLM最大的价值是把大模型推理的门槛从“买卡/租卡”降到了“动手调参”。它让你手里的普通显卡不再只是跑小模型而是能触碰大模型的边界。虽然速度不适合生产环境但用来学习、做实验、跑离线任务非常有意义。如果你也想试我的建议是先准备一台内存和硬盘都很宽裕的机器再选一个中等规模的模型跑通流程最后再挑战大参数模型。别一上来就奔着70B去容易消磨掉热情。AirLLM不是银弹它是在硬件限制下找到的一种平衡。理解了它的调度逻辑你也就理解了大模型推理的显存瓶颈到底卡在哪里。最后分享一个小技巧跑之前尽量预热内存。先在系统里把模型文件读取一遍让文件缓存生效再用AirLLM加载第一次加载速度会快不少。这个操作很小但省下来的时间相当可观。
阅读完成 · 觉得有帮助?
咨询建站