★ FEATURED ARTICLE 本地跑大模型这件事从去年到现在热度一直没降过。我身边不少做算法的朋友平时在云端调API调得挺顺手一旦要把模型搬到自己的开发机上问题就全冒出来了显存不够、权重加载报错、推理速度慢得离谱、tokenizer对不上号。Gemma这个系列的模型开源之后我第一时间在自己的机器上折腾了一轮PyTorch集成踩的坑比预想的多不少。这篇就把整个从零到跑通的过程拆开讲清楚包括权重怎么拿、环境怎么配、显存怎么省、推理怎么加速以及那些文档里不会写但实际一定会遇到的问题。不管你是刚接触本地推理的新手还是已经跑过其他开源模型想换Gemma试试的老手下面这些内容应该都能直接用上。1. 为什么值得在本地用PyTorch跑Gemma1.1 本地推理的真实价值在哪先说清楚一件事本地跑模型不是为了替代云端服务两者解决的是不同问题。云端API适合快速验证产品逻辑但一旦涉及数据隐私、离线场景、高频调用成本控制本地推理的优势就出来了。Gemma作为开放权重的模型系列参数量从2B到7B不等这个规模区间恰好卡在一个很舒服的位置——7B的模型经过量化之后一张消费级显卡就能扛住2B的版本甚至能在笔记本上跑起来。用PyTorch做集成而不是直接用现成的推理框架核心原因在于可控性。你需要改模型结构做实验、需要拿到中间层的激活值、需要自定义采样策略这些在封装好的高层接口里都很难做到。PyTorch给你的是完全透明的计算图每一层的输入输出你都能插手。代价就是配置工作量大依赖关系得自己理清楚。我自己的使用场景是做一个离线的文本处理工具需要对一批敏感文档做摘要和分类数据不能出本机。这种情况下云端方案直接排除本地PyTorch推理是唯一选择。跑通之后单次推理延迟在可接受范围内批量处理一晚上能跑完几千条完全够用。1.2 Gemma模型的结构特点与集成难点Gemma的架构基于Transformer decoder但和Llama系列有一些细节差异这些差异恰恰是集成时容易翻车的地方。首先是词表大小Gemma用的是256000大小的词表比很多同规模模型大得多这意味着embedding层的参数量占比不低加载时对内存的瞬时压力更大。其次是RMSNorm的位置和实现方式Gemma在注意力层和前馈层之前都做了归一化顺序搞错的话输出会完全乱掉。还有一个容易被忽略的点是Gemma的注意力机制里用了RoPE位置编码但具体的旋转维度配置和基频参数跟其他模型不完全一样。如果你直接套用别的模型的加载代码权重能load进去但推理结果就是不对而且不会报错这种静默错误最难排查。我在第一次跑的时候生成的内容全是重复的乱码查了大半天才发现是位置编码的参数没对上。另外Gemma的权重发布格式有多种原始的是safetensors分片也有社区转换的PyTorch bin格式。不同格式的加载逻辑不一样分片文件的索引方式也有区别。这些细节在后面章节会具体展开。2. 环境搭建依赖版本这块必须较真2.1 PyTorch与CUDA的版本匹配逻辑这一步看起来基础但实际翻车率极高。PyTorch的版本、CUDA的版本、显卡驱动版本三者之间是链式依赖关系任何一环不匹配都会导致要么装不上要么装上了跑不起来。我的建议是先确定显卡驱动支持的CUDA最高版本然后倒推PyTorch版本而不是反过来。具体操作上先用nvidia-smi看驱动版本和它支持的CUDA版本。假设显示支持到CUDA 12.1那你就应该装对应CUDA 12.1构建的PyTorch。注意这里有个常见误解nvidia-smi显示的CUDA版本是驱动支持的最高版本不是你系统里装的CUDA版本实际运行时PyTorch自带CUDA运行时只要不超过驱动支持的上限就行。# 查看驱动和CUDA支持情况 nvidia-smi # 安装对应CUDA 12.1的PyTorch pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121装完之后一定要验证别急着往下走import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0)) print(torch.cuda.get_device_capability(0))is_available()返回True只是第一步还要看设备算力等级。Gemma用到了bfloat16精度这个精度需要算力8.0以上Ampere架构及以后才有硬件加速。如果你的卡是更老的架构bfloat16会回退到软件模拟速度会慢很多这时候应该改用float16。2.2 transformers与accelerate的版本陷阱Gemma的支持是在transformers的特定版本之后才加入的版本太老会直接报模型类型不认识。但版本也不是越新越好新版本可能引入API变更导致你参考的旧代码跑不通。我实测下来transformers 4.38以上的版本对Gemma支持比较完整accelerate建议0.27以上。pip install transformers4.38 accelerate0.27 safetensors sentencepiece这里特别提一下sentencepieceGemma的tokenizer依赖它缺了会在加载tokenizer时报错而且报错信息不太直观容易让人以为是权重文件的问题。另外safetensors这个库必须装现在主流权重都是这个格式它比传统的pickle格式加载更快也更安全。提示如果你之前装过其他模型的依赖建议用虚拟环境隔离。我遇到过transformers版本被其他项目锁死在旧版本导致Gemma加载失败的情况排查起来很浪费时间。2.3 显存预估与硬件底线在下载权重之前先算清楚你的硬件能不能扛住。Gemma 7B的原始权重是bfloat16精度参数量70亿光权重本身就要约14GB显存加上推理时的激活值、KV cache实际峰值会到16GB以上。2B版本原始精度约5GB推理峰值7GB左右。模型版本原始精度显存8bit量化4bit量化推荐硬件Gemma 2B约5GB约3GB约2GB8GB显存卡Gemma 7B约14GB约8GB约5GB16GB显存卡这张表是理论值实际会有浮动。如果你只有8GB显存想跑7B就必须走量化路线后面章节会讲具体怎么做。CPU推理也不是不行但7B模型在纯CPU上生成一段几百字的文本可能要几分钟实用性很低除非你只是做验证。3. 权重获取与加载分片文件的处理细节3.1 权重文件的组织方式Gemma的权重在发布时是按分片存储的一个7B模型通常拆成多个safetensors文件外加一个index.json索引文件。索引文件里记录了每个参数名对应哪个分片文件加载器靠它来定位。如果你手动下载一定要保证分片文件和索引文件在同一个目录下缺一个都会加载失败。目录结构大概长这样gemma-7b/ ├── config.json ├── generation_config.json ├── model-00001-of-00004.safetensors ├── model-00002-of-00004.safetensors ├── model-00003-of-00004.safetensors ├── model-00004-of-00004.safetensors ├── model.safetensors.index.json ├── tokenizer.json ├── tokenizer.model └── tokenizer_config.jsonconfig.json里定义了模型结构参数包括层数、隐藏维度、注意力头数、词表大小等。加载之前建议先打开看一眼确认torch_dtype字段和你的硬件匹配。有些转换版本会把这个字段写成float32那样加载后显存占用会翻倍需要手动改成bfloat16。3.2 用from_pretrained加载的正确姿势最省事的加载方式是用transformers的from_pretrained它会自动读索引文件、按需加载分片。但默认行为是把整个模型加载到内存再搬到显存中间会有个内存峰值。对于显存紧张的机器可以用device_map参数做分层加载。from transformers import AutoModelForCausalLM, AutoTokenizer import torch model_id ./gemma-7b tokenizer AutoTokenizer.from_pretrained(model_id) model AutoModelForCausalLM.from_pretrained( model_id, torch_dtypetorch.bfloat16, device_mapauto, low_cpu_mem_usageTrue )device_mapauto会让accelerate自动计算每层放哪显存不够时会往CPU内存里放一部分。low_cpu_mem_usageTrue能显著降低加载时的内存峰值原理是它不先创建完整模型再填权重而是边加载边分配。这两个参数组合起来是我在显存边缘情况下最常用的配置。加载完成后验证一下模型到底在哪print(model.hf_device_map)这个输出会告诉你每一层被分配到了哪个设备。如果看到大量层被分到cpu说明显存确实不够推理速度会受影响这时候就该考虑量化了。3.3 加载报错的典型排查路径加载阶段最常见的报错有三类我按出现频率排一下。第一类是KeyError提示某个参数名找不到。这通常是权重版本和模型代码版本不匹配比如用旧版transformers加载新版权重。解决办法是升级transformers或者检查config.json里的模型类型字段是否被正确识别。第二类是显存不足的CUDA out of memory。注意这个报错可能出现在加载过程中而不是推理时因为加载时会有瞬时峰值。对策是加low_cpu_mem_usageTrue或者先用量化配置加载。第三类是RuntimeError关于dtype不匹配。比如权重是bfloat16但你指定了float16某些算子会报错。统一dtype就能解决建议全程用bfloat16硬件支持的前提下。注意加载大模型时如果中途CtrlC中断可能会留下不完整的缓存文件下次加载会报奇怪的错误。遇到这种情况去缓存目录把对应的不完整文件删掉重来别在错误状态上反复尝试。4. 推理跑通从第一段生成文本开始4.1 最简推理脚本与参数含义模型加载好之后先跑一个最简的生成脚本验证链路通不通。这一步不要加任何花哨的参数就用默认配置确认能出正常文本再说。prompt 用一句话解释什么是机器学习 inputs tokenizer(prompt, return_tensorspt).to(model.device) with torch.no_grad(): outputs model.generate( **inputs, max_new_tokens100, do_sampleFalse ) print(tokenizer.decode(outputs[0], skip_special_tokensTrue))这里几个参数值得说清楚。max_new_tokens控制生成的最大长度注意它不包括输入部分的token数。do_sampleFalse走的是贪心解码输出确定但容易重复设成True会启用采样配合temperature和top_p控制多样性。第一次验证建议先用贪心排除采样参数带来的干扰。torch.no_grad()必须加否则PyTorch会构建计算图显存占用暴涨且毫无意义因为推理不需要梯度。4.2 生成质量异常的定位方法如果跑通了但输出是乱码、重复或者答非所问按下面的顺序排查。先看tokenizer是否匹配。用tokenizer.decode(tokenizer.encode(测试文本))做个往返测试如果解码回来和原文不一致说明tokenizer文件有问题或者加载错了版本。Gemma的tokenizer对空格和特殊字符的处理比较敏感往返测试能快速暴露问题。再看位置编码。前面提过Gemma的RoPE参数有特殊性如果生成内容呈现规律性重复比如一直重复同一个短语大概率是位置编码没生效或者参数错了。检查config.json里的rope_theta字段确认加载代码没有覆盖它。最后看精度。bfloat16在算力不足的卡上会静默降级输出可能变得不稳定。临时切到float32跑一次对比如果float32正常而bfloat16异常就是精度支持的问题。4.3 批量推理的写法与注意事项实际使用中往往是批量处理逐条推理效率太低。批量推理的关键是padding和attention mask的正确处理。prompts [解释梯度下降, 什么是过拟合, 简述注意力机制] tokenizer.padding_side left inputs tokenizer(prompts, return_tensorspt, paddingTrue).to(model.device) with torch.no_grad(): outputs model.generate( **inputs, max_new_tokens150, do_sampleFalse, pad_token_idtokenizer.pad_token_id ) for i, output in enumerate(outputs): text tokenizer.decode(output[i][inputs[input_ids].shape[1]:], skip_special_tokensTrue) print(f--- 样本{i} ---) print(text)这里有两个坑。第一padding_side必须设成left因为decoder-only模型是从左边开始生成右边padding会导致生成位置错乱。第二解码时要截掉输入部分只取新生成的部分否则输出里会混入prompt。pad_token_id也要显式指定Gemma的tokenizer默认可能没有pad token需要手动设成eos token的id。批量大小不是越大越好KV cache会随批量线性增长。7B模型在16GB卡上批量开到4基本就到顶了再大就OOM。建议从批量2开始试逐步往上加找到上限。5. 显存优化量化与KV cache管理5.1 bitsandbytes 8bit与4bit量化的取舍显存不够的时候量化是最直接的方案。bitsandbytes提供了8bit和4bit两种量化加载方式接入transformers只需要加几个参数。from transformers import BitsAndBytesConfig bnb_config BitsAndBytesConfig( load_in_4bitTrue, bnb_4bit_compute_dtypetorch.bfloat16, bnb_4bit_quant_typenf4, bnb_4bit_use_double_quantTrue ) model AutoModelForCausalLM.from_pretrained( model_id, quantization_configbnb_config, device_mapauto )nf4是4bit量化的数据类型针对正态分布的权重做了优化比普通的fp4精度损失小。use_double_quantTrue会对量化常数再做一次量化额外省一点显存代价是轻微的速度损失。compute_dtype指定计算时的精度设成bfloat16能在量化和精度之间取得平衡。8bit和4bit怎么选我的经验是显存够用8bit就别上4bit8bit的精度损失很小基本感知不到质量下降4bit在7B模型上会有可察觉的质量损失尤其是长文本生成时容易出现逻辑断裂。但如果显存是硬约束4bit是能跑和不能跑的区别那就没得选。5.2 KV cache的显存占用与控制很多人只算权重显存忽略了KV cache。KV cache的大小和序列长度、批量大小、层数、注意力头数都相关。粗略估算公式是2 × 层数 × 头数 × head_dim × 序列长度 × 批量 × 精度字节数。7B模型在序列长度2048、批量1的情况下KV cache大概占1GB左右批量4就接近4GB。控制KV cache的手段有几个。一是限制max_new_tokens生成长度直接决定cache增长。二是用past_key_values做增量推理避免重复计算。三是transformers较新版本支持的cache_implementationstatic能预分配固定大小的cache避免动态增长带来的碎片。outputs model.generate( **inputs, max_new_tokens200, use_cacheTrue, cache_implementationstatic )提示开了量化之后KV cache默认还是按原始精度存的这部分显存省不下来。如果显存实在紧张可以考虑用torch_dtypetorch.float16配合量化权重KV cache的占用能减半。5.3 CPU offload的代价与适用场景device_mapauto在显存不够时会自动把部分层放到CPU这就是CPU offload。它能让你在显存不足的机器上跑起来但代价是速度大幅下降因为每层计算都要在CPU和GPU之间搬数据。实测下来7B模型如果有一半层被offload到CPU生成速度会降到纯GPU的十分之一左右。这个速度对于交互式应用基本不可用但对于离线批处理还能接受。所以offload适合的场景是不追求延迟、只求能跑完的批处理任务。如果决定用offload建议手动指定device_map而不是用auto把embedding层和前面的层放GPU后面的层放CPU这样数据搬运的路径更短。具体配置需要根据你的显存大小试没有通用答案。6. 推理加速从能跑到跑得快6.1 torch.compile的实际效果与坑PyTorch 2.0之后引入的torch.compile能把模型编译成更高效的执行图对推理有加速效果。用法很简单model torch.compile(model, modereduce-overhead)但实际用下来有几个坑。第一首次编译时间很长7B模型可能要几分钟这段时间看起来像卡死了其实是正常的。第二编译后的模型对输入形状敏感如果每次输入长度变化很大会触发反复编译反而更慢。第三某些自定义算子可能不支持编译会回退到eager模式并给出警告。mode参数有几个选项reduce-overhead适合小批量低延迟场景max-autotune会花更多时间调优但峰值性能更好。我的建议是先用默认模式跑通确认加速有效再尝试其他模式。如果输入长度比较固定编译的收益会比较明显如果长度变化剧烈可能得不偿失。6.2 注意力实现的替换transformers支持替换注意力实现较新版本可以用Flash Attention 2或者PyTorch原生的scaled_dot_product_attention。后者不需要额外装包在支持的硬件上会自动启用高效实现。model AutoModelForCausalLM.from_pretrained( model_id, torch_dtypetorch.bfloat16, attn_implementationsdpa, device_mapauto )attn_implementation设成sdpa会走PyTorch内置的高效注意力设成flash_attention_2需要装flash-attn包。sdpa的兼容性更好加速效果在中等序列长度下已经不错。flash-attn在长序列下优势更明显但安装比较折腾对CUDA版本和编译环境有要求。实测在序列长度512到2048这个区间sdpa相比默认实现有20%到40%的加速具体取决于硬件。这个改动成本极低建议默认开启。6.3 投机解码的可行性投机解码speculative decoding用一个小的草稿模型先猜几个token再用大模型验证能显著加速生成。transformers里通过assistant_model参数支持这个特性。assistant AutoModelForCausalLM.from_pretrained( ./gemma-2b, torch_dtypetorch.bfloat16, device_mapauto ) outputs model.generate( **inputs, assistant_modelassistant, max_new_tokens200 )思路是用2B的Gemma给7B的Gemma做草稿。理论上草稿模型和大模型的分布越接近接受率越高加速越明显。但实际用下来两个模型都要占显存7B加2B的组合在16GB卡上会很紧张可能需要量化。而且加速比不稳定取决于任务类型代码生成这类确定性强的任务接受率高创意写作这类多样性强的任务接受率低。我的建议是显存充裕且任务类型固定时值得一试否则优先用前面两个更简单的加速手段。7. 实操中反复踩到的几个坑7.1 tokenizer的padding与特殊token问题Gemma的tokenizer默认没有pad token这在批量推理时会直接报错。解决办法是把pad token设成eos tokentokenizer.pad_token tokenizer.eos_token tokenizer.padding_side left但这样有个副作用padding位置填的是eos模型可能会把padding当成序列结束的信号影响生成。更稳妥的做法是在attention mask里正确标记padding位置让模型忽略它们。transformers在paddingTrue时会自动生成attention mask确保它被传进generate。另一个坑是Gemma的对话模板。Gemma有特定的对话格式用tokenizer.apply_chat_template能正确套用。如果手动拼接prompt而不套模板模型可能表现不佳因为它训练时见的就是特定格式。7.2 生成结果被截断或提前结束生成突然停住常见原因有三个。一是碰到了eos token这是正常结束。二是max_new_tokens设太小还没说完就被截断。三是生成了重复内容触发了某些停止条件。排查时先把max_new_tokens调大排除长度限制。然后检查输出末尾是不是eos如果是说明模型认为说完了可能是prompt引导有问题。如果输出是重复的调低temperature或者加repetition_penalty。outputs model.generate( **inputs, max_new_tokens300, repetition_penalty1.1, no_repeat_ngram_size3 )repetition_penalty大于1会惩罚已出现的tokenno_repeat_ngram_size禁止重复的n-gram。这两个参数能有效抑制重复但设太大会让输出变得生硬建议从小值开始调。7.3 多卡环境下的设备分配混乱多卡机器上device_mapauto有时会把模型拆到多张卡上但如果卡之间没有NVLink跨卡通信会成为瓶颈速度反而不如单卡。这种情况下建议手动指定device_map{: 0}强制全部放第一张卡前提是单卡显存够。如果确实需要多卡并行用device_mapbalanced让accelerate均匀分配或者手动写每层的映射。多卡推理的复杂度明显高于单卡除非模型大到单卡放不下否则不建议折腾。8. 一些让集成更顺手的经验权重下载这块分片文件加起来十几个GB网络不稳的时候容易断。建议用支持断点续传的工具下载下完核对一下文件大小和索引文件里记录的是否一致。我遇到过一次分片下载不完整加载时报的参数缺失错误查了很久才发现是文件本身的问题。缓存管理也值得注意。transformers默认把模型缓存在用户目录下几个模型下来能占几十GB。可以用HF_HOME环境变量把缓存指到大容量磁盘避免系统盘被撑爆。加载时如果本地已有缓存local_files_onlyTrue能避免联网检查加快加载速度。版本锁定是另一个容易被忽视的点。跑通一套配置之后把transformers、torch、accelerate的版本号记下来写进requirements。我吃过亏隔了几个月重装环境默认装了新版本结果API变了之前的代码跑不通又得重新调。最后说个心态上的事。本地跑大模型的配置过程确实琐碎报错信息也经常不直观但每一步踩过的坑都会变成经验。我现在的习惯是每解决一个问题就在脚本里加一行注释记下来下次遇到类似情况直接查自己的笔记比翻文档快得多。这套Gemma的PyTorch集成跑通之后换其他同架构的模型基本能复用大部分配置边际成本会越来越低。
阅读完成 · 觉得有帮助？
