概述在从零开始写Qwen3一模型结构分析中已经识别了Qwen3的基本结构了解了使用哪些关键组件现在开始实现将模型分为这几个模块attn注意力计算不包含短接和QKV嵌入无参数self_attn自注意力调用attnnorm主要实现RMSNormropefeedback实现SwiGLUtransformer_block一个完整的Transformer Decoder层qwen3完整的模型本章的目标就是快速搭建一个模型能够推理即可之后再自己手写组件从零开始写Qwen3目录模型组件Attn直接使用现成的fromtorch.nn.functionalimportscaled_dot_product_attentionreturnscaled_dot_product_attention(q,k,v,is_causalis_causal,enable_gqaTrue,scaleself.n_head_embed**-0.5,)需要注意qkv以及输出o的维度是 BxHxSxDH是多头在enable_gqa时Q的头可以是KV的整倍数O的头数和Q相同is_causal表示是否开启遮罩。a. 这里非常奇怪的一点理论上在decode阶段开启KVCacheQ长度只有1这个参数有没有无所谓但如果传了True会导致结果错误所以我们在decode阶段将其手动设置为Falsescale就是自注意力中的那个缩放D \sqrt{D}Dnorm有现成的def__init__(...):self.weightnn.Parameter(torch.ones(norm_dim))defforward(...):returnnn.functional.rms_norm(x,self.weight.shape,self.weight,self.epseps是为了防止分母为0添加的一个数这个参数一般模型config.json会提供Qwen3-0.6B是1e-6rope这个没有现成的需要手写原理RoPE是模仿向量内积的那个夹角cos θ \cos\thetacosθ弄的让两个向量旋转它们位置对应的角度然后内积就可以体现它们之间相对位置了对于二维向量可以在平面上执行旋转但高维向量无法保证都在同一个平面上其实并不需要旋转高维向量x ⋅ y ∑ i 1 d x i y i ∑ k 1 d / 2 ( x 2 k − 1 y 2 k − 1 x 2 k y 2 k ) \mathbf{x} \cdot \mathbf{y} \sum_{i1}^{d} x_i y_i \sum_{k1}^{d/2} \left( x_{2k-1} y_{2k-1} x_{2k} y_{2k} \right)x⋅yi1∑dxiyik1∑d/2(x2k−1y2k−1x2ky2k)两个高维向量的内积其实可以看做多个二维向量内积之和大模型隐藏层维度一般是偶数如果真是奇数可以补充一个0上去rope的做法就是把特征向量当成多个二维向量对每个二维向量进行旋转对一个二维向量的旋转就是乘一个矩阵R ( θ ) [ cos θ − sin θ sin θ cos θ ] R(\theta) \begin{bmatrix} \cos\theta -\sin\theta\\ \sin\theta \cos\theta \end{bmatrix}\\R(θ)[cosθsinθ−sinθcosθ]实际上每个小向量可以旋转不同的角度比如低维度的可以旋转大一点的角度高维度旋转小一点的角度这样在相同cos θ \cos \thetacosθ时高维度必须距离更远才能达到和低维度相同的系数低维度可以捕获短距离、局部的关系高维度则捕获长距离、全局的关系一般是这样生成基础旋转角的θ β − 2 i / d \theta\beta^{-2i/d}θβ−2i/di是维度组向下除以2取整d是维度beta是10000乃至更大这个参数也在config.json中有记录对每个在位置n的词元嵌入旋转n θ n\thetanθ实现注意到旋转之后有这个x ′ cos θ x − sin θ y y ′ sin θ x cos θ y x ′ ⃗ [ x ′ ; y ′ ] cos θ x ⃗ sin θ σ ( x ⃗ ) x\cos \theta x -\sin\theta y\\ y\sin \theta x \cos \theta y\\ \vec{x}[x;y]\cos \theta \vec{x} \sin \theta \sigma(\vec{x})x′cosθx−sinθyy′sinθxcosθyx′[x′;y′]cosθxsinθσ(x)这里σ \sigmaσ是奇偶交换并把交换后的偶数部分取负实际上σ \sigmaσ除了奇偶交换原始版本还有其他方式比如GPT-Neox的交换前一半和后一半Qwen3也用这个在llama.cpp中将交换奇偶的版本叫做normal交换前后一半的叫做neox将上面的公式扩展后就有了rope的实现x ′ ⃗ cos θ ⃗ ∘ x ⃗ sin θ ⃗ ∘ σ ( x ⃗ ) \vec{x}\cos \vec{\theta} \circ \vec{x} \sin \vec{\theta}\circ \sigma(\vec{x})x′cosθ∘xsinθ∘σ(x)这个rope在单次前向计算中对所有层不管是K还是V这个cos θ ⃗ , sin θ ⃗ \cos\vec{\theta},\sin\vec{\theta}cosθ,sinθ都是相同的所以全局计算一次即可如果可以保证生成长度不会超过最大长度那甚至可以预先计算最大长度的这两个向量然后后面直接截取其中一部分使用即可简单计算defbuild_cos_sin_embed(self,dtype,position_ids:torch.Tensor)-tuple[torch.Tensor,torch.Tensor]:inv_freq1.0/(self.base_freq**(torch.arange(0,self.head_dim,2,deviceposition_ids.device).float()/self.head_dim)).unsqueeze(0)freqstorch.einsum(bj,bk-bjk,position_ids,inv_freq)embtorch.cat([freqs,freqs],dim-1)return(emb.cos().to(dtype),emb.sin().to(dtype))这里做了一个乘法其实就是n θ n\thetanθ为每个位置生成一个旋转向量einsum是爱因斯坦求和用于张量计算的bj,bk-bjk 的意思是输入两个矩阵下标分别是bj和bk然后乘出来一个三维数组下标是 bjk也就是相同 b 的数乘到一起C b j k A b j B b k C_{bjk}A_{bj}B_{bk}CbjkAbjBbk如果是 jb,bk-jk那就是常规的矩阵乘法了C j k ∑ j A b j B j k C_{jk}\sum_j A_{bj}B_{jk}Cjk∑jAbjBjk通常后面这个求和符号可以省略selfAttn这个需要自己手写也简单QKV进行投影QK需要进行后norm然后把QKV通过viewtranpose转为 BHSD 的格式对QK执行rope然后计算attn得到oo转置到BSHD然后reshape到BSD后投影输出所有的投影直接使用Linear即可其他feedback其实没什么特别的组件Linear和Swish激活函数就没了transformer_block 就是上面组件拼接即可qwen3在从零开始写Qwen3一模型结构分析中展示了forward就是一个Embedding嵌入执行多层TransformerBlock最后NormLinear就行运行推理加载参数HuggingFace下载的模型文件夹下有safetensors文件Qwen3-0.6B较小只有一个文件这就是模型参数里面包含所有张量和张量的名字通过张量的名字来加载到参数中实际上这些张量的名字就是transformers中成员变量的名字这些成员要么嵌套包含其他组件要么包含一个叫做.weight的成员类型是nn.Parameter后者就是要加载的参数实际上所有的张量名字都是以.weight结尾的比如一个张量的名字是model.layers.8.self_attn.k_norm.weight它就是编号为8从0开始的解码层的自注意力组件的k_norm的参数执行推理有个完整的模型就可以推理了整个模型的输入是嵌入id列表可以直接使用from tokenizers import Tokenizer这个库这个库和transformers的AutoTokenizer效果是一样的因为都是加载了模型目录下的tokenizer.jsontokenizerTokenizer.from_file(model_path/tokenizer.json)print(prompt)inputstensor.tensor([tokenizer.encode(prompt)])输出参数是BxSxD的logits只需要最后一个S的特征向量它的长度就是字典长度每个值代表一个词元可以简单argmax得到最可能的下一个词元id然后把这个词元id拼接到 inputs预测下一个词元不断如此直到达到最大长度或者遇到 eos提示词可以直接写字符串这是大模型的补全模式但模型目录下还有一个tokenizer_config.json文件其中有一个chat_template成员这是一个Jinja语法的模板字符串它可以填充模板实现多轮对话的格式也是对话模式可以直接加载上withopen(model_path/tokenizer_config.json)asf:datajson.load(f)templatejinja2.Template(data[chat_template])prompttemplate.render(messages[{role:user,content:介绍一下你自己}])tokenizerTokenizer.from_file(model_path/tokenizer.json)这个模板实际上就是我们常用的ChatMessage格式。模板中还支持工具调用写tool_calls就行不知道Qwen3-0.6B的模型的工具调用能力如何有机会写个Agent看看这个模板基本是这样的截取前半部分{%-iftools%}{{-|im_start|system\n}}{%-ifmessages[0].rolesystem%}{{-messages[0].content\n\n}}{%-endif%}{{-# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within tools/tools XML tags:\ntools}}{%-fortoolintools%}{{-\n}}{{-tool|tojson}}{%-endfor%}{{-\n/tools\n\nFor each function call, return a json object with function name and arguments within tool_call/tool_call XML tags:\ntool_call\n{\name\: function-name, \arguments\: args-json-object}\n/tool_call|im_end|\n}}{%-else%}{%-ifmessages[0].rolesystem%}{{-|im_start|system\nmessages[0].content|im_end|\n}}{%-endif%}{%-endif%}这是Jinja语法上面的内容大概是这样如果有tools这个变量传入才处理如果第一个messages也就是可以传数组的role是system把系统提示词放到总体提示词开头接着写入工具提示词把工具描述放进去直接json但这个json通常都是有严格格式要求的否则也就是没有传入工具简单把系统提示词传入注意到每轮对话都是以|im_start|role开头|im_end|结尾eos就是这个|im_end|现在没有做KVCache运行速度很慢而且越来越慢这是我们下一节要解决的问题