1. 先搞清楚:一次推理到底在跑什么
1.1 前向传播的骨架
当你决定自己手写一个LLM推理引擎的时候,第一步往往不是写代码而是画图。你得把模型的forward pass拆成一张数据流图:输入token序列进来,先查embedding表得到向量,然后过N层transformer block,每一层里先后做attention和MLP,中间穿插RMSNorm和残差连接,最后过final norm和lm_head得到logits。这个过程看起来简单,但真正动手之后你会发现,代码只会决定计算怎么做,而权重文件决定了计算能不能做出来。
我把“权重”理解成一组已经训练好、固定下来的数字。推理引擎要做的事,本质上是拿着这些数字,按照架构定义好的顺序,做一系列矩阵乘法和向量运算。所以权重既是程序的输入数据,也是程序最重要的资源。你写再漂亮的kernel,如果权重没加载对,或者精度转换有问题,结果就是一堆乱码。更麻烦的是,权重文件在不同格式、不同模型之间差异极大,不提前设计好加载层,后面每接入一个新模型就要痛苦一次。
这一章我集中讲权重。包括权重在模型里的分布规律、主流文件格式的差异、从磁盘到显存/内存的加载流程、推理时权重是怎么参与计算的,以及我踩过的各种和权重有关的坑。如果你也在写自己的推理引擎,这一篇应该能帮你省下不少时间。
1.2 权重的分布:参数藏在哪些层里
以LLaMA系模型为例,一个7B模型差不多有70亿个参数,这里面绝大部分权重都集中在两类地方:attention子层和MLP子层。理解权重分布不是学术兴趣,它直接决定内存规划策略和量化优先级。
先说attention。每一层transformer里通常有四个投影矩阵,名字在各个框架里略有差异,但含义都一样:
- q_proj:把隐藏状态投影成query向量
- k_proj:投影成key向量
- v_proj:投影成value向量
- o_proj:把attention输出投影回hidden size
以标准7B模型为例,hidden size通常取4096,所以q_proj、k_proj、v_proj、o_proj每个都是4096乘4096的大矩阵。单看这四个,参数量就是4乘4096再乘4096,大约6700万。一共32层,光attention的投影权重就贡献了大概21亿参数。知道这个数字有什么意义?当你做显存规划的时候,就能按层去估算,也能确认“某一层加载完应该有多少字节”,如果对不上说明加载有问题。
再看MLP。LLaMA的MLP是三个矩阵:gate_proj、up_proj、down_proj。gate和up把hidden size从4096投射到中间维度11008,down再投射回4096。这三个矩阵的参数量大概是4096乘11008乘3,约等于1.35亿每层,32层加起来就是43亿参数。你看,权重的分布其实非常集中,绝大部分都压在MLP上。这也是为什么很多量化方案会优先对MLP做激进压缩,因为它的体量最大,压缩收益最明显。
除了这些,每层还有RMSNorm的权重向量,维度是hidden size,一个只有4096个数。最后是token embedding表,7B模型的词表一般是32000,embedding维度4096,这就是32000乘4096,大约1.31亿个参数。我在这里提醒一个高频误区:很多模型的embedding和lm_head是共享的,也就是tied weights。后面专门讲,几万字不在话下,因为这在写加载器的时候是个容易踩坑的点。
2. 权重文件的真实面貌:格式、布局与组织
2.1 主流序列化格式与选择
你从模型社区下载一个模型,拿到的往往不是单个文件,而是一个模型目录。里面通常有这几个东西:config.json记录架构超参数,比如层数、hidden size、注意力头数;generation_config.json记录生成参数;权重文件可能是pytorch_model.bin或者model.safetensors;还有就是tokenizer文件,分词器需要词表和合并规则。
权重文件格式主要有两大类:PyTorch的.bin,本质是pickle序列化的state_dict;以及safetensors。我个人强烈建议新写的引擎直接支持safetensors,原因很简单:它把tensor的元数据和原始数据分区存放,用mmap可以零拷贝直接读取,不必像pickle那样反序列化整个对象。同时它不支持任意代码执行,安全得多。老模型的.bin文件里,一个7B模型的state_dict全在一坨pickle里,加载时要把整个文件读进内存再解析,峰值内存很容易多出好几个GB。safetensors的文件头是一段JSON,记录了每个tensor的名字、dtype、shape和字节偏移,找数据就是一次seek加一次read的事。
如果你要写一个能加载GGUF格式的引擎,情况会复杂一点。GGUF是llama.cpp生态的格式,它把超参数、tokenizer词表、tensor数据和量化信息全部封装在一个文件里,用专用的键值对结构描述。它的tensor数据区和safetensors类似,按名称记录offset和shape,但是文件头解析逻辑更复杂,因为它要支持多种量化类型,比如Q4_K、Q6_K、Q8_0这些。我的建议是:除非你要直接跑GGUF量化模型,否则不要把它的加载逻辑和safetensors的加载逻辑混在一起。两个模块分开写,对外暴露同一个接口,后面维护会舒服很多。
还有一种裸权重文件,比如某些原始发布里的pytorch_model.bin目录,里面每个文件对应一个tensor,用np.memmap就能读。这种格式在几个开源模型上还很常见,但要么是为了极简,要么是历史遗留,实际工程里最好先转换到统一格式再喂给引擎。
2.2 张量命名与维度映射
不同模型家族对同一概念的命名并不一致。LLaMA的叫法是q_proj、k_proj、v_proj、o_proj、gate_proj、up_proj、down_proj,GPT-2里叫c_attn、c_proj,Mistral基本沿用了LLaMA的命名,Falcon又不一样。所以一个通用推理引擎的核心要有一张张量名字映射表。从工程角度说,最好在模型启动阶段就根据config.json里的model_type把外部名字映射到引擎内部的统一命名空间。比如llama的model.layers.0.self_attn.q_proj.weight和mistral的对应名字,应该统一解析成layer[0].attn.qkv.weight这样的内部结构。
我提一个容易被忽略的点:很多开源权重文件里的tensor名字带model.、transformer.这样的前缀,或者干脆不带前缀。加载的时候必须做归一化处理。如果你不管三七二十一直接按名字找,遇到某些模型会加载不出来。我写过一个tensor name normalizer,把常见前缀、层的数字位置、子模块名字分别拆出来再重组,实测下来能把huggingface上绝大多数公开模型映射到同一套内部结构。
还有个维度问题。PyTorch的Linear层权重形状一般是out_features在前、in_features在后,但有些模型保存时转置过,或者某些推理框架用了in、out的布局。你在做矩阵乘法、量化、算子调度之前,必须先统一约定好自己引擎的张量布局。最稳妥的策略是加载后第一时间把原始shape解析出来,记录一个meta信息,别依赖“文件名里说是什么就是什么”。
2.3 元数据:权重跟架构的“婚约”
权重文件只有数字,并没有说它应该怎么用。真正决定计算图的是config.json。我建议在引擎实现里把加载权重和构建计算图分开。计算图先根据config.json把每层要创建的权重Tensor全部声明出来,包括名字、shape、dtype和初始值,加载器再去权重文件里逐个匹配。这个设计的好处是,遇到权重文件缺失或者shape不匹配,你能立刻知道是哪一层、缺哪个量,而不是等跑到某个算子时报一堆摸不着头脑的错误。
我在加载一个7B模型时,正常情况下每层应该有8个tensor:attn的4个投影、MLP的3个投影、1个RMSNorm权重。再加上embedding和final norm,一共就是层数乘8加上额外几个。如果加载完发现数量对不上,优先怀疑权重文件是不是用低精度保存,或者模型有部分层是共享参数的。这些东西都要靠系统性的元数据对比来发现。写一个“期望张量清单”和“实际加载清单”的比对工具,比你在调试器里一行行看要快得多。
3. 从磁盘到显存:权重的装载流程
3.1 mmap与零拷贝
说到加载权重,最容易犯的错误就是把整个权重文件一次性读进内存,再转成目标dtype,再拷到显存。对一个14GB的fp16模型来说,这一条路径下来,内存峰值会轻松超过30GB,而实际真正需要的内存只有14GB。正确做法是先用mmap把文件映射到虚拟内存空间,然后按tensor粒度去读取。safetensors本身就是为这个流程设计的:解析文件头拿到每个tensor的offset和length,通过memoryview或者指针直接访问对应区域。比如用frombuffer基于这部分内存直接创建tensor,是不需要拷贝的。如果你用自己的加载器,也可以直接用pread或者mmap偏移切出bytes。
不过mmap也不是万能药。文件被放在网络盘,或者系统内存本来就很紧张的时候,mmap触发的page fault会把进程卡住。我的经验是加载权重时保留一个流式拷贝的开关,默认对本地磁盘走mmap,对网络路径或特殊环境走普通读文件再进内存。这是工程里很小的一个设计,但能救命的场景不少。
另一个要注意的点是字节序。权重文件通常按小端序存储,x86、ARM的小端都没问题,但如果你把引擎跑在大端机器上,比如某些特定芯片环境,就要在加载层做字节序转换。这个很多教程不会讲,但真遇到了非常难受。踩过一次之后,我把字节序转换做成了加载器的可选阶段,默认不启用,但遇到可疑数据时会自动检测并警告。
3.2 dtype转换与内存规划
权重文件的原始dtype一般有三种:fp32、fp16、bf16。加载进引擎后,你要决定推理时用什么精度计算。fp32最省事但太费内存,7B模型就是28GB,普通消费级显卡直接GG。fp16内存占用是14GB,计算快,但数值范围有限,某些模型在前向推理时也容易出inf,尤其是attention里的softmax之前,QK的转置乘积结果可能非常大。bf16的数值范围和fp32一样大,内存同样省一半,非常适合推理,前提是硬件支持bf16计算。如果你的引擎跑在纯CPU环境,或者老一点的GPU,可能fp16反而是更稳的选择。
除开精度,还有内存规划的问题。加载过程中,如果你边读边转dtype,就可能同时在内存里存在原始文件、转换后tensor、以及GPU侧tensor的同一份数据。想让峰值内存可控,可以采用一条小流水线:先读一小部分权重到临时buffer,转好dtype填进预分配好的目标tensor,然后丢掉这一部分的原始数据。这样最多只需要buffer大小的额外内存,而不是整个文件的大小。
做显存规划时还有个细节很多人忽略:权重的显存不等于模型总显存占用。KV cache虽然名字里带cache,但它是推理过程中动态增长的,和权重无关。很多人因为经常把KV cache的占用也归类到模型权重里,于是计算显存需求时把两者混在一起,最后导致OOM。我的建议是拆开统计:权重是一块固定的大头,KV cache按token数和层数动态估算,两者分开打印。后面排查显存问题时,一眼就能看出是谁在占地方。
3.3 分片加载与模型并行
大模型经常是分片保存的,早期是pytorch_model-00001.bin、00002.bin这样的多个文件,现在safetensors也常用safetensors.index.json做分片索引。写加载器的时候,最好一开始就支持一个逻辑权重分散在多个物理文件里的情况。实现上并不复杂:解析index文件得到每个tensor的名字到文件的映射,读数据时先定位文件,再做偏移读取。这个能力在模型超过单卡显存、或者要做多机推理时几乎是硬需求。
真正麻烦的是并行加载时的内存峰值。比如你用8个线程同时读8个分片,每个分片动辄2GB,一个不注意就内存爆了。我用的方案是有界生产者-消费者队列:每批只派发2到3个分片,解析完一批,拷贝到目标权重区域并完成精度转换,再派发下一批。队列深度做成可配置项,默认值设在单卡显存和单文件大小的比例附近。这样既能利用多线程的IO带宽,又不会让临时内存无限膨胀。
4. 权重在推理中是怎么“跑”的
4.1 矩阵乘法与计算路径
一次token生成,本质上是把输入向量挨个通过权重矩阵做变换。embedding查表拿到行向量后,第一件事就是把输入向量和q_proj、k_proj、v_proj这三个矩阵分别相乘。在实现里,这往往是GEMM(通用矩阵乘)而不是GEMV(矩阵乘向量)。为什么?因为实际推理是批量的,一个prompt序列有几百个token,就算只跑一次生成,也要同时处理多个位置,形成一个小矩阵乘。
权重矩阵在显存里的排布方式对性能影响很大。大多数推理框架会要求权重按特定分块方式排布,比如CUDA的cutlass库喜欢把矩阵拆成16乘16或者8乘8的小块,按block-tile的方式组织,这样SM的shared memory利用率最高。如果你只是先做一个不求甚解的引擎,可以直接用现成的cuBLAS、rocBLAS或者OpenBLAS,权重保持普通行主序就行。但我强烈建议加一层布局切换的抽象:加载完成后按需转置或重排权重,算子侧根据layout选择对应的kernel。这个抽象在后期接入量化weight时尤其有用,因为反量化kernel最怕权重排版乱变。
4.2 量化权重的反量化路径
量化是权重处理里最大的坑。以最常见的group quantization为例,权重不是直接以int4或int8存储,而是被分成一个个group,比如每128个元素一组,每组保存一个scale和一个zero_point。推理时,你需要先根据scales把整数权重还原成浮点,再去做矩阵乘。
这里有个优化点:不要在每个算子内部实时反量化,而是加载阶段反量化成fp16存一份。这样当然省事,但会损失量化节省内存的核心优势。更常见的方案是“按需反量化”:第一次使用某个量化tensor时,反量化到对应精度的临时缓存里,后续如果同一层被反复调用,直接复用缓存。不过对于LLM这种一次生成要重复调用所有层的场景,缓存效用有限,更关键的还是把反量化做得足够高效,比如在CUDA上用向量化加载加查表,快速把int4转成fp16。
如果你是自己设计引擎,我建议第一版先别上量化,用fp16或者bf16把整条链路跑通,再慢慢加量化支持。因为量化的时序问题会让你分不清到底是权重本身坏了,还是反量化算法写错了。这类问题极其浪费时间。
4.3 权重共享与tied embeddings
前面提到,很多模型的token embedding和lm_head是同一个权重。这对写引擎的人意味着两件事:第一,加载时你可能只有一个embedding文件,lm_head是空的或者是引用关系;第二,当你做量化时,不能对同一个权重做两份不同的处理,否则最终推理会不一致。
我见过不少工具开发者,因为没考虑tied embeddings,结果输出第一个token总是错的。为什么?因为他们在加载embedding之后又单独加载了lm_head,但lm_head的文件不存在,于是初始化为随机值。随机矩阵乘logits,输出当然就乱了。排查这类问题最好的办法是在加载完成后做一次模型拓扑层面的检查:embedding的shape和lm_head的shape一致时,强制只用一份数据,不要重复初始化。更稳的做法是把tied机制作为加载器的一等公民,在元数据里显式标记哪些是共享tensor。
5. 常见问题与排查技巧实录
5.1 加载失败:命名、shape与缺失
我在写引擎的过程里最常遇到三类加载失败。
第一类是命名不匹配。解决思路前面说过,做一个名字归一化层。遇到不匹配时,第一时间打印缺失的tensor和多余tensor的完整名字,把对照表打出来,而不是只报一个key not found。对照表会让你立刻发现是层序号偏移、前缀缺失还是命名风格不同。
第二类是shape不匹配。比如某个模型在attention里用了GQA,也就是Grouped Query Attention,那么k_proj的形状可能不是通常的hidden size乘hidden size,而是hidden size乘num_kv_heads再乘head_dim。这种如果你在加载阶段沿用旧模型的形状去校验,就会报错。应对方法是对每个tensor都保存一个期望形状列表,加载时做宽松匹配和严格校验两种模式的切换。开发期用严格模式抓问题,生产期用宽松模式兼容变体。
第三类是文件缺失或文件头不完整。从网上下载大文件经常出现aborted后缀的残缺文件,所以我在加载器里加了文件大小和哈希校验。对safetensors来说,文件头长度字段如果解析出异常值,直接报corrupt file,不要继续解析,否则会读到错乱数据。其实还有一个隐蔽的问题:某些文件头里的JSON带BOM头,或者包含注释,解析器如果不够宽容也会挂。为这个坑绕了两天后,我换成了容错JSON解析并手动跳过BOM。
5.2 显存 / 内存不足的排查
OOM是家常便饭。我建议在加载阶段就把模型权重占用、额外buffer占用、KV cache占用三个数值全部打出来,而不是等到算子跑起来报CUDA OOM才去猜。权重占用是固定的,buffer在加载完会释放,KV cache是变化的。如果发现OOM发生在生成阶段而权重占用没问题,基本可以确定是KV cache的评估不准,或者并发序列太多。经常有人把context长度设得很大,导致KV cache的预留量远超实际,这类问题一打印就能看出来。
另一个排查方向是显存碎片。有些框架在加载多个分片文件后释放了临时内存,但因为分配粒度不规律,显存被割成很多小块,后续KV cache申请大块连续显存时失败。解决方法是开一个显存池,至少保证KV cache的分配在池内进行,或者干脆在加载后做一次显存碎片整理。与其等到OOM再处理,不如在分配环节就做一个统一入口。
5.3 加载后精度表现奇怪
如果你加载完权重后发现模型输出全是乱码或者重复同一句话,先别怀疑模型本身,绝大多数情况下是权重精度转换出的问题。一个经典案例是:把fp32权重直接截断成fp16,但fp16的有效数字位数有限,如果模型里某些权重的绝对数值特别大,比如超过65504,转成fp16就会变成inf,进而污染整个后续计算。
处理这种问题的方案是,在加载dtype转换时使用safe cast:对超出fp16表示范围的值做clip,而不是让它直接溢出成inf或NaN。有些框架会默认开启这个clamp,有些不会,你得自己确认。另外,bf16和fp16之间互转也会引入微小误差,如果你发现两次加载结果有细微差别,很可能是位级精度问题,不影响绝大多数场景,但做算子级单元测试时确实很烦人。我的经验是固定一套转换路径,别一会儿走fp16一会儿走bf16,不然会把误差来源搞得很乱。
5.4 调试工具与日志策略
最后说调试工具的实战经验。我在推理引擎里专门维护了一个权重加载报告,每次启动时输出:
- 期望tensor数量和实际tensor数量
- 每个tensor的名字、shape、dtype、来自哪个文件哪个offset
- dtype转换说明,原始是什么,线上是什么,用什么策略
- 是否出现tied weight或shared weight
- 加载耗时和峰值内存
这份报告看起来啰嗦,但在排查“为什么同一个模型在不同机器上表现不同”的时候,能帮你省下大把时间。比如同一个权重文件,在一台机器上被转成bf16跑,在另一台机器上被转成fp16跑,输出概率分布当然有差异,但如果没这份报告,你会以为是引擎写错了。日志要分级,平时默认只打摘要,遇到问题再打开full trace。full trace会记录每个tensor的加载细节,虽然慢一点,但关键时刻能救命。
6. 踩坑总结与实操建议
6.1 三套表示的分离设计
写推理引擎这段时间,我最大的感受是:模型架构现在的复杂度并不高,注意力、MLP、归一化就那几件事,真正让人头疼的反而是权重这一层。权重文件的格式、命名、分片、精度、内存布局、量化方案,每一件都需要在动手写算子之前先想好。
我强烈建议你在引擎里维持三套表示的分离。第一套是外部格式的表示,比如safetensors或者GGUF里tensor长什么样;第二套是内部张量的表示,一个统一的Tensor类,带shape、dtype、layout、device;第三套是计算图层面的表示,描述某层需要哪些逻辑权重。加载器要做的事情就是把外部表示转成内部张量,再按计算图的期望去装配。千万不要把三种职责放在一个类里,否则后面加新格式或者新模型时,你会改到怀疑人生。
我还加了第四层,专门处理元数据。每个tensor加载完都会记录它的来源文件、源token在文件里的位置、原始dtype、当前dtype、是否经过layout重排。这套元数据在调试分布并行和量化问题时价值极大。有一次我发现两个分片加载顺序不同会导致最终结果有细微差异,就是这个元数据帮我锁定了是设备端浮点加法顺序导致的,不是权重本身的错。
6.2 转换工具与开发节奏
如果你每天要加载几十次模型,每次都去解析safetensors再转dtype,纯属浪费时间。我写了一个独立的convert工具,把各种来源的模型权重统一转换成自己引擎的私有压缩包格式,格式很简单:一个目录,里面放tensor索引文件加原始二进制。转换成自有格式之后,引擎启动时解析成本极低,还方便做增量校验和缓存。
这种私有格式还能记录一些额外元数据,比如provenance,也就是这批权重是从哪个原始仓、哪个commit、哪个转换脚本生成的。模型文件在路上被改过、脚本版本对不上这类问题,一查provenance就清楚了。做多模型对比实验时,这个信息直接决定你要不要重新拉权重。
如果你从零开始写引擎,我建议按这样的顺序推进:先用fp32加载一个最小的GPT-2这样的小模型,验证整条链路;再换成fp16或bf16的大模型;然后再接safetensors的mmap加载;最后才考虑GGUF量化和其他格式。把权重这一层拆得越干净,后面的算子优化、KV cache优化、采样策略设计都会轻松不少。
最后再分享一个小经验:权重的拷贝和转换一定要做成可中断、可恢复的流程,尤其是加载几十GB模型的时候。真遇到加载到一半程序崩了,如果每次都要从头再来,心态很容易崩。我后来在加载器里加了checkpoint机制,每加载完几层就记录一下进度,崩溃后可以从断点继续。这个功能写起来不难,但体验提升非常明显。希望这篇关于权重的心得能帮到正在折腾自己LLM推理引擎的朋友。