显存不够用的日子相信每个在本地折腾大模型的人都经历过。跑推理跑到一半啪一下OOM模型直接崩了前功尽弃。前阵子我在社群里又看到有人问“我的卡是24G显存跑7B模型怎么还会爆显存”这问题一看就知道不是显存真不够而是没做显存规划。这篇文章不聊训练单聊推理场景下怎么把显存抠着用把我自己踩过坑、试过有效的方法全部摊开来讲。先给不熟悉的朋友补个背景本地跑大模型推理显存占用主要分成三块——模型权重、KV Cache键值缓存、以及激活值推理中间过程产生的临时张量。绝大多数人遇到的OOM其实就是KV Cache在作怪因为权重是固定的而KV Cache会随着生成的长度线性增长。你输入越长、生成越长显存占用就越高。明白了这一点优化思路就清晰了要么减小权重要么控制KV Cache要么把部分计算挪到别处去。这篇文章适合那些用消费级显卡跑大模型、经常被OOM困扰、想在不换卡的前提下榨干现有硬件性能的朋友。我会从显存分配的原理讲起逐步给出量化、KV Cache优化、批处理控制、SWAP等一系列实操手段最后附上问题排查实录。1. 先搞懂显存是怎么被吃掉的1.1 推理阶段的显存分配逻辑模型加载进显存的那一刻占用的空间主要就是权重参数。以7B模型为例FP16精度下每个参数占2字节光权重就需要7 × 10^9 × 2 ≈ 14GB显存。你手里如果是一张24G显存的卡看似还能剩下10G但注意这个时候模型还没开始真正干活。一旦开始推理Transformers库里的past_key_values会缓存每个Transformer层在计算Attention时的K和V矩阵。这个缓存的显存消耗公式是2K和V两组 × 层数 × 注意力头数 × 头维度 × 序列长度 × 批次大小 × 2字节。具体数值大家一般记不住我直接说结论对于7B模型序列长度512、批大小1的情况下KV Cache大约占用1-2G当序列长度推到4096时这个数字会膨胀到8-10G。这就是为什么很多人发现“短文本没问题长文本必OOM”的根本原因。激活值则是另一个动态指标。它跟具体的模型结构、隐藏层维度、批次大小都有关系通常在短序列时占比不高但随着batch size增大激活值的增长速度非常夸张。这也是为什么我建议本地推理尽量把batch size设为1的原因之一。1.2 用工具看一眼真实占用与其猜不如直接看。在终端里用nvidia-smi只能看到进程占用的总显存没法看到里面细分。更好的办法是用torch.cuda.memory_summary()或者直接跑一段小代码把权重、KV Cache、激活值分别打出来。import torch from transformers import AutoModelForCausalLM, AutoTokenizer model AutoModelForCausalLM.from_pretrained(your_model_path, torch_dtypetorch.float16).cuda() tokenizer AutoTokenizer.from_pretrained(your_model_path) inputs tokenizer(显存优化测试, return_tensorspt).to(cuda) model.eval() with torch.inference_mode(): outputs model.generate(**inputs, max_new_tokens256, use_cacheTrue) print(torch.cuda.memory_summary())torch.cuda.memory_summary()会输出非常详细的显存分配情况包括PyTorch缓存分配器的内部状态。第一次跑这个命令你会发现自己预留给PyTorch的缓存可能远大于实际用到的这块可以通过环境变量调小。1.3 显存碎片的隐藏杀手还有一个很多人忽略的点显存碎片化。PyTorch的缓存分配器为了避免重复分配的开销会默认缓存一部分显存块结果就是你看到nvidia-smi显示显存占用很高但其实很多是缓存块而不是真实使用。这种情况通常出现在反复加载不同大小模型或者多次执行变长序列推理之后。如果你发现显存明明没有占满但出现了OOM十有八九就是碎片化问题。最简单的解决办法是重启推理进程或者用torch.cuda.empty_cache()把缓存块释放掉。但请注意这个函数只能释放空闲的缓存块不能解决模型本身占用的显存。更彻底的办法是设置环境变量PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True这个选项可以让PyTorch使用可扩展内存段来减少碎片。实测在连续跑多个不同长度文本时这个设置能明显降低碎片导致的虚高占用。2. 权重瘦身量化和精度选择2.1 FP16、INT8、INT4到底选哪个现在主流的本地推理框架llama.cpp、vLLM、Transformers配合bitsandbytes都支持量化。量化的核心逻辑就是降低每个参数占用的比特数把FP16降到INT8权重占用直接减半降到INT4占用再减半。以7B模型为例FP16需要14GBINT8需要7GBINT4只要3.5GB左右。从我的实际体验来看如果你的显存刚好能装下FP16模型那是首选因为精度最高推理表现最稳定。如果FP16装不下或加载后剩余显存太少优先考虑INT8INT4虽然省显存最多但推理质量下降比较明显尤其是在处理中文长文本、逻辑推理强的任务时回答的连贯性会打折扣。这里有个反直觉的经验很多人以为量化只影响显存占用其实它还影响推理速度。INT8和INT4的权重体积更小减少了PCIe传输和显存带宽压力在某些平台上的解码速度反而比FP16更快。当然这跟具体硬件和框架有关建议自己在自己的卡上做一次benchmark。2.2 用bitsandbytes做低精度加载如果你继续用Transformers库最简单的量化方式是bitsandbytes的load_in_8bit或load_in_4bit配置。它的好处是代码改动极小坏处是因为它做了逐层量化模型加载时间会变长而且首次加载时会把一部分量化计算放在CPU上进行。from transformers import AutoModelForCausalLM, BitsAndBytesConfig import torch quantization_config BitsAndBytesConfig( load_in_4bitTrue, bnb_4bit_compute_dtypetorch.float16, bnb_4bit_quant_typenf4, bnb_4bit_use_double_quantTrue, ) model AutoModelForCausalLM.from_pretrained( your_model_path, quantization_configquantization_config, device_mapauto, )这套参数里bnb_4bit_use_double_quantTrue会启用二次量化把量化常数也压一遍大约能再省0.5-1GB显存。bnb_4bit_compute_dtypetorch.float16表示计算时反量化到FP16这算是在显存和精度之间的折中选择。2.3 GGUF的量化格式值不值得用如果你用过llama.cpp应该对GGUF格式不陌生。GGUF内置了多种量化方案最常用的是Q4_K_M、Q5_K_M、Q6_K这些。以7B模型为例Q4_K_M大小约为4.08GBQ5_K_M约为4.45GBQ6_K约为5.53GB。我的建议是如果你不需要跑复杂的中文上下文依赖任务可以用Q5_K_M作为日常使用版本如果显存比较紧张就用Q4_K_M。这两个量化级别在绝大多数实际对话场景下质量差距很小。但是Q2、Q3级别我不太建议日常使用那个质量下降肉眼可见除非你就是想拿一块4G老卡体验一下本地模型。3. KV Cache最容易被忽视的大户3.1 KV Cache到底怎么控制上一节提到KV Cache是推理过程中动态增长的显存消耗。它的大小主要取决于四个参数模型层数、上下文窗口长度、批次大小、以及是否使用GQAGrouped Query Attention。前三个都好理解第四个需要单独解释。MHAMulti-Head Attention里每一层的每个注意力头都会各自缓存一份K和V显存开销最大。GQA分组查询注意力是让多个查询头共享一组K和V可以显著减少KV Cache的显存占用。比如从MHA切换到GQAKV Cache的显存可以减半甚至更多。这个能力不是模型现有的就得靠框架支持而是跟模型架构一起决定的。你在引入模型时需要确认它是否用了GQA很多新版模型如Llama 3、Mistral已经用了GQA这也是它们能在消费级显卡上跑长文本的原因之一。如果推理框架支持你可以手动限制KV Cache的最大长度或参与计算的上下文长度。比如Transformers里设置max_new_tokens并不能直接控制KV Cache真正相关的是max_length或 generation config里的max_position_embeddings。把最大输入长度限制在4096而不是直接拉满到32KKV Cache占用会差好几倍。3.2 开启use_cache与关闭use_cache的博弈很多教程会告诉你关闭use_cache可以省显存。这确实没错但代价是推理速度大幅下降。因为关闭KV Cache后模型每生成一个token都需要重新计算前面所有token的注意力值复杂度从线性退化到平方级。你会在长文本生成时明显感觉到速度变慢慢到让你怀疑是不是卡了。我的看法是除非你只是在做极短的对话测试否则不要关闭use_cache。更好的思路是配合vLLM这类高性能推理框架它们对KV Cache的管理更精细引入了PagedAttention可以按页分配显存利用率高得多且吞吐量更大。在vLLM里你可以通过--max-model-len和--gpu-memory-utilization来显式控制显存占用比例实操中我会把gpu-memory-utilization设为0.85到0.9剩下一点空间给CUDA context和其他进程避免直接顶满后出幺蛾子。3.3 Sliding Window和位置编码的隐性影响有些模型通过滑动窗口注意力Sliding Window让KV Cache只保留窗口范围内的token这能大幅限制Cache增长。Mistral 7B就用了这个机制所以它在长文本场景下的显存表现优于同样大小的旧模型。如果你选模型时特别在意长文本下的显存优先考虑带窗口注意力或GQA的架构。位置编码类型也会间接影响显存占用。比如使用RoPE旋转位置编码的模型在推理时需要在Q和K上追加位置信息这部分虽然本身不大但在长上下文时会增加一部分计算和临时显存开销。你要是把上下文长度拉到几万这部分同样不可忽视。但这些算优化细节普通场景下先关注前三点就够了。4. 实践调优从加载到生成的完整链路4.1 加载阶段的显存省钱法除了量化之外加载阶段还有几个容易踩的坑。第一个是使用device_mapauto让框架自动分配模型层到不同的设备GPU/CPU/MPS对于显存不宽裕的机器让一部分层驻留在CPU内存里GPU只处理计算密集的层。这样做的好处是显存不爆坏处是速度变慢因为CPU和GPU之间要频繁搬运权重。第二个是torch_dtype的选择。设定torch_dtypetorch.float16而不是不设因为加载FP32的模型权重会把显存占用翻倍。很多人OOM其实就是FP32导致的。第三个是关闭不需要的梯度计算。推理模式下记得用torch.inference_mode()包裹它会禁用梯度跟踪减少中间张量保存带来的显存开销。4.2 生成阶段的参数控制生成阶段我最常调的几个参数是max_new_tokens、temperature、top_p和num_beams。其中num_beams是显存杀手贪心解码只维护一个候选序列而beam search会同时维护beam_width个候选序列KV Cache也跟着翻倍。如果你本地推理显存紧张beam search最好设置为1也就是直接用贪心采样。很多朋友一上来就习惯性用beam search调个4或5稍微长一点的文本就直接OOM这不是模型问题是参数问题。temperature和top_p虽然不直接占显存但它们决定了采样过程的随机性。对追求稳定性的本地使用来说temperature0.7、top_p0.9是比较通用的区间。另外max_new_tokens千万不要设得太大很多人设置成512看着不够就设2048结果生成长文必爆。先设小一点分次生成并拼接结果也是绕开单次KV Cache峰值的好办法。4.3 一个7B模型在8G显存上的配置实战这里给一个可复现的配置是我在GTX 3060 8G上跑7B模型的常用方案模型Q4_K_M格式的GGUF模型通过llama.cpp加载上下文--ctx-size 2048不直接拉满批处理--batch-size 256短输入够用线程数--threads 8GPU部分用--n-gpu-layers 999尽可能多地把层放到GPU上KV Cachellama.cpp按上下文长度自动分配2048时大约2G左右实测这套配置下7B模型大概占6-7G显存剩余一点空间给系统和其他应用生成速度在CPU和GPU搭配下能保持每秒10-15 token日常问答完全够用。如果你的显存比8G更小那就适当减小ctx-size到1024同时只加载部分层到GPU让CPU帮跑一部分。别小看这种“轻重混跑”它比直接OOM强太多了。4.4 显存不足时的兜底方案CPU Offload和Swap显存实在不够时还有一个兜底方案是使用CPU Offload。Transformers库的device_mapauto已经做了这件事但你可以更细粒度地控制哪些层放在GPU、哪些层放在CPU。常用模式是Embedding层和最后几层放在GPU中间的Transformer层部分放CPU。因为Embedding和LM Head的计算量大但参数量不一定最大放GPU能保住速度。中间的层虽然多但CPU也能算就是慢一点。llama.cpp也支持--n-gpu-layers参数来指定把前N层放到GPU剩下的在CPU跑。我建议从N模型总层数的一半开始试然后逐渐增大找到一个显存刚好不满的临界值。这个临界值会随上下文长度变化所以改完上下文后最好重新试。用这种方式跑大模型速度会比我前面说的全GPU方案慢不少但总比跑不起来好对不对等你后面换了显卡直接把这个层级拉满就行。5. 常见OOM问题与排查技巧实录5.1 问题一加载就爆还没推理就OOM这通常是权重精度问题。检查一下是不是默认用了FP32加载把模型转成FP16或INT4就能解决。还有一个可能CUDA context本身会占用几百MB到1G不等的显存如果机器显存真的很小比如6G以下这部分占比很高。可以在加载模型前先torch.zeros(1).cuda()预热CUDA看看空载占用然后根据剩余空间调整模型精度。5.2 问题二短文本没问题长文本必爆这是最典型的KV Cache增长导致的。解决方案优先考虑缩短上下文窗口然后考虑量化或换用带GQA的模型。如果你的业务确实需要长上下文可以考虑给模型升级到更大显存的机器或者使用支持PagedAttention的框架这样KV Cache按页管理长上下文的显存利用效率会好很多。5.3 问题三同样配置过一会儿就OOM这个通常有两个原因。一个是进程里残留了多个推理上下文比如在notebook里反复跑cell旧的推理结果没释放解决方法是重启kernel或显式del后调用gc.collect()和torch.cuda.empty_cache()。另一个原因是累积生成了过长的对话历史每次生成都在对话历史的KV Cache基础上继续增长。如果对话不能中断可以定期把历史消息截断只保留最近几轮。这个经验很实用很多把模型当聊天机器人用的朋友聊久了必爆就是因为历史越积越长。5.4 问题四多进程跑测试显存互相挤占如果你一边跑本地模型一边开着浏览器或视频剪辑软件显存被吃掉了不少。nvidia-smi能帮你看到谁在占显存但这种场景解决方式很简单先把无关应用关了。如果你是在做自动化测试需要并行的多个推理进程建议给每个进程设置不一样的PYTORCH_CUDA_ALLOC_CONF和CUDA_VISIBLE_DEVICES把负载分散到多张卡上。单卡的话就别并行跑多个模型了老老实实做排队。5.5 快速排查清单我给自己做过一张排查表遇到OOM就按这个查一遍排查项操作效果权重精度检查是否FP32考虑改为FP16/INT8/INT4立减50%-75%权重显存上下文长度降低ctx-size或max_length控制KV Cache上限解码策略确认beam size为1避免多个候选序列的KV翻倍模型加载方式使用device_mapauto做CPU Offload分散显存压力历史对话截断并只保留最后几轮避免KV累积到爆显存碎片调用empty_cache或重启进程释放缓存碎块这套表从我个人的使用体验来看能解决90%以上的本地推理OOM。剩下10%属于显卡本身确实装不下这个规模的模型那就得考虑进一步量化或者换模型而不是硬撑着。6. 进阶优化框架选型和硬件调优6.1 vLLM和llama.cpp怎么选如果你只跑单轮或少次对话llama.cpp最省事下载即用CPU和GPU混合推理也做得很成熟。它的GGUF量化格式选择多而且社区活跃遇到问题很容易找到解决方案。如果要做并发比较高的推理服务比如搭一个给团队用的API接口vLLM更合适。它的PagedAttention能让你在有限的显存里塞进更大的吞吐量而且连续批处理能显著提高GPU利用率。但vLLM对显存的最小要求比较高小显存卡上跑不太开。我自己在8G卡上优先选llama.cpp在有20G以上显存的机器上跑服务时就用vLLM。6.2 GPU驱动与CUDA版本的坑有一部分OOM是驱动层的问题。比如某些老版本驱动对特定PyTorch版本的显存管理不太好导致显存释放不彻底。建议在可控范围内把CUDA toolkit和PyTorch的版本对齐至少在解决OOM问题时不要排除这一项。nvidia-smi显示驱动支持的最高CUDA版本torch.version.cuda显示当前PyTorch内置的CUDA版本这两个可以快速对照一下。6.3 模型选择才是最大的显存优化说句实在话如果你总是在压榨显存最省心的“优化”其实是选一个参数规模更合适的模型。7B模型跑不动就看看1.5B-4B的小模型很多场景下表现并不差。比如一些简洁问答、格式整理类的任务4B模型和7B模型差距很小但显存占用差了一半。把一个模型用到极致值得敬佩但没必要死磕量力而行选模型也是内容质量控制的一部分——毕竟模型输出质量取决于你能跑得起哪种模型。选模型还有一个可以直接参考的点同一系列模型的量化版本大小已知用显存总容量去匹配模型文件大小KV Cache余量即可。7B的Q5_K_M约4.45GBKV Cache留2G那8G卡基本就是稳的如果显存只有6G那就选Q4_K_M并把上下文压到1024。写在最后OOM这件事本地跑大模型的朋友十有八九都会遇到每次遇到都觉得很上头的点在于它往往不是显存真的满了而是某一项配置没做好。我自己从最开始只会满世界找“更小模型”的版本到现在能根据手头机器参数快速给出量化等级、上下文窗口、KV Cache策略的搭配组合中间也就是把这些细节一个一个调明白的过程。如果你手头也有个“老显存焦虑症”建议从今天就开始做三件事第一装一个带显存检测的工具跑一下自己的真实占用第二试试把量化等级降一级看看输出质量你能不能接受第三给当前常跑的模型配一个固定的显存配置模板。三件事做完你再回头看你遇到的OOM多半已经能自己诊断了。