【干货收藏】手搓一个Qwen3：从输入层到输出层的代码解析

摘要：2025年4月29号，阿里千问团队发表了Qwen3模型，这是 Qwen 系列大型语言模型的最新成员。他们的旗舰模型Qwen3-235B-A22B在代码、数学、通用能力等基准测试中，与 DeepSeek-R1、o1、o3-mini、Grok-3 和 Gemini

2025年4月29号，阿里千问团队发表了Qwen3模型，这是 Qwen 系列大型语言模型的最新成员。他们的旗舰模型 Qwen3-235B-A22B 在代码、数学、通用能力等基准测试中，与 DeepSeek-R1、o1、o3-mini、Grok-3 和 Gemini-2.5-Pro 等顶级模型相比，表现出极具竞争力的结果。此外，小型 MoE 模型 Qwen3-30B-A3B 的激活参数数量是 QwQ-32B 的 10%，表现更胜一筹，甚至像 Qwen3-4B 这样的小模型也能匹敌 Qwen2.5-72B-Instruct 的性能。

1、输入处理：输入的文本首先被转换为token ID序列，通过嵌入层将token转换为向量表示。

2、核心Transformer块：Qwen3由多个相同的Transformer块堆叠而成，每个块包含以下组件：3、自注意力机制：使用多头自注意力机制，能够捕捉输入序列中各个token之间的关系。4、归一化层：采用RMSNorm替代LayerNorm，在每个自注意力层和前馈网络之后进行归一化处理。5、前馈网络：使用GEGLU或SwiGLU激活函数，增强模型的表达能力。输出层：通过RMSNorm后，使用QWenLMHead生成下一个token的预测概率分布。

模型结构如下图所示：

根据上面的五个组成部分，我们分别来进行一个实现。

1、数据输入处理,通过加载tokenizer，将输入的input text转为token_ids

device = torch.device("gpu") # 这里根据实际情况选择，可以是cpu,也可以是gputokenizer_file_path = Path("qwen3") / "tokenizer-base.json"model_file = Path("qwen3") / "qwen3-0.6B-base.pth"tokenizer = Qwen3Tokenizer(tokenizer_file_path=tokenizer_file_path)model = Qwen3Model(QWEN_CONFIG_06_B)model.load_state_dict(torch.load(model_file))model.to(device)prompt = "什么是大语言模型？"input_token_ids_tensor = torch.tensor(tokenizer.encode(prompt),device=device).unsqueeze(0)

接下来我们进行第二步，Qwen3网络结构整体编码。

2、Qwen3网络结构编码

class Qwen3Model(nn.Module):def __init__(self, cfg):super.__init__# 模型的主要参数, 开源模型一般就是config.json文件里面的参数self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"],dtype=cfg["dtype"])# 大模型层数self.trf_blocks = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg["n_layers"])])self.final_norm = RMSNorm(cfg["emb_dim"])self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False,dtype=cfg["dtype"])# rope 旋转位置编码if cfg["head_dim"] is None:head_dim = cfg["emb_dim"] // cfg["n_heads"]else:head_dim = cfg["head_dim"]cos, sin = compute_rope_params(head_dim=head_dim,theta_base=cfg["rope_base"],context_length=cfg["context_length"])self.register_buffer("cos", cos, persistent=False)self.register_buffer("sin", sin, persistent=False)self.cfg = cfgself.current_pos = 0 # kvcache使用def forward(self, in_idx, cache=None):tok_embeds = self.tok_emb(in_idx)x = tok_embedsnum_tokens = x.shape[1]if cache is not None: # kvcachepos_start = self.current_pospos_end = pos_start + num_tokensself.current_pos = pos_endmask = torch.triu(torch.ones(pos_end, pos_end, device=x.device, dtype=torch.bool),diagonal=1)[pos_start:pos_end, :pos_end]else:pos_start = 0 # Not strictly necessary but helps torch.compilemask = torch.triu(torch.ones(num_tokens, num_tokens, device=x.device,dtype=torch.bool),diagonal=1)# Shape (1, 1, num_tokens, num_tokens) to broadcast across batch and headsmask = mask[None, None, :, :]for i, block in enumerate(self.trf_blocks):blk_cache = cache.get(i) if cache else Nonex, new_blk_cache = block(x, mask, self.cos, self.sin,start_pos=pos_start,cache=blk_cache)if cache is not None:cache.update(i, new_blk_cache)x = self.final_norm(x)logits = self.out_head(x.to(self.cfg["dtype"]))return logitsdef reset_kv_cache(self):self.current_pos = 0

这部分代码就是整个Qwen3 核心所在：

第一步：完成了了input text 到token ids的过程；

第二步(看结构图)： Token embedding layer ,通过nn.embedding将token ids 转为embedding

第三步(看图，红色块)： Transfomer Block，通过trf_blocks封装包含RMSNorm、GQA、FFN

第四步：Final RMSNorm , 对应代码中的 self.final_norm(x)

最后：Linear output layer, 对应 self.out_head

3、Transformer Block 编码

class TransformerBlock(nn.Module):def __init__(self, cfg):super.__init__self.att = GroupedQueryAttention( # GQA实现d_in=cfg["emb_dim"],num_heads=cfg["n_heads"],head_dim=cfg["head_dim"],num_kv_groups=cfg["n_kv_groups"],qk_norm=cfg["qk_norm"],dtype=cfg["dtype"])self.ff = FeedForward(cfg)self.norm1 = RMSNorm(cfg["emb_dim"], eps=1e-6)self.norm2 = RMSNorm(cfg["emb_dim"], eps=1e-6)def forward(self, x, mask, cos, sin, start_pos=0, cache=None):# Shortcut connection ,一种特殊的残差连接shortcut = xx = self.norm1(x)x, next_cache = self.att(x, mask, cos, sin, start_pos=start_pos,cache=cache) # Shape [batch_size, num_tokens, emb_size]x = x + shortcut # Shortcut connection for feed-forward blockshortcut = xx = self.norm2(x)x = self.ff(x)x = x + shortcut # Add the original input backreturn x, next_cache

这部分代码是Transformer Block核心(看途中红色或蓝色框)：

第一步： self.norm1(x) 参数归一化，后面紧跟着GQA

第二步：GQA,GQA是什么结构，可以参考下图

第三步：shortcut connnection

第四步： self.norm2(x)

第五步：self.ff(x)

第六步：shortcut connnection

4、GQA的编码

很多模型都会在这个基础去创新，比如最开始的MHA–>GQA–>MQA–>MLA，通过上图可以看到GQA核心是将KV进行分组，多个Q给到分组的KV(核心参数：num_kv_groups)，从而减少计算和内存开销。

class GroupedQueryAttention(nn.Module):def __init__(self,d_in, num_heads , num_kv_groups ,head_dim =None,qk_norm=False,dtype=None):super.__init__assert num_heads % num_kv_groups == 0self.num_heads = num_headsself.num_kv_groups = num_kv_groupsself.group_size = num_heads // num_kv_groupsif head_dim is None:assert d_in % num_heads ==0head_dim = d_in // num_headsself.head_dim = head_dimself.d_out = num_heads * head_dimself.W_query = nn.Linear(d_in, self.d_out,bias=False,dtype=dtype)self.W_key = nn.Linear(d_in,num_kv_groups * head_dim , bias=False,dtype=dtype)self.W_value = nn.Linear(d_in,num_kv_groups * head_dim,bias=False,dtype=dtype)self.out_proj = nn.Linear(self.d_out,d_in,bias=False,dtype=dtype)if qk_norm:self.q_norm = RMSNorm(head_dim,eps=1e-6)self.k_norm = RMSNorm(head_dim, eps=1e-6)else:self.q_norm = self.k_norm = Nonedef forward(self,x,mask,cos,sin,start_pos=0,cache=None):b, num_tokens,_ = x.shapequeries = self.W_query(x)keys = self.W_key(x)values = self.W_value(x)queries = queries.view(b,num_tokens,self.num_heads,self.head_dim).transpose(1,2)keys_new = keys.view(b,num_tokens,self.num_kv_groups,self.head_dim).transpose(1,2)values_new = values.view(b,num_tokens,self.num_kv_groups,self.head_dim).transpose(1,2)if self.q_norm:queries =self.q_norm(queries)if self.k_norm:keys_new = self.k_norm(keys)queries = apply_rope(queries,cos,sin,offset=start_pos)keys_new = apply_rope(keys_new,cos,sin,offset=start_pos)if cache is not None:prev_k , prev_v = cachekeys = torch.cat([prev_k,keys_new],dim=2)values = torch.cat([prev_v,values_new],dim=2)else:start_pos = 0keys ,values = keys_new,values_newnext_cache = (keys, values)keys = keys.repeat_interleave(self.group_size,dim=1)values = values.repeat_interleave(self.group_size,dim=1)attn_scores = queries @ keys.transpose(2, 3)attn_scores = attn_scores.masked_fill(mask, -torch.inf)attn_weights = torch.softmax(attn_scores / self.head_dim**0.5, dim=-1)context = (attn_weights @ values).transpose(1,2)context = context.reshape(b,num_tokens,self.d_out)return self.out_proj(context),next_cache

这里整个代码就是GQA的实现逻辑，核心还是QKV分组、维度对齐、计算的过程。接下来我们看看FFN

5、FFN实现

class FeedForward(nn.Module):def __init__(self,cfg):super.__init__self.fc1 = nn.Linear(cfg["emb_dim"],cfg["hidden_dim"],dtype=cfg["dtype"],bias=False)self.fc2 = nn.Linear(cfg["emb_dim"],cfg["hidden_dim"],dtype=cfg["dtype"],bias=False)self.fc3 = nn.Linear(cfg["emb_dim"],cfg["hidden_dim"],dtype=cfg["dtype"],bias=False)def forward(self,x):x_fc1 = self.fc1(x)x_fc2 = self.fc2(x)x = nn.functional.silu(x_fc1) * x_fc2return self.fc3(x)

FFN实现相对简单，一个多层感知机，通过三个线性层以及激活函数完成，

6、RMSNorm归一化

归一化是为了让参数在某一个范围内，以致模型训练过程比较稳定

class RMSNorm(nn.Module):def __init__(self,emb_dim, eps=1e-6, bias =False, qwen3_compatible=True):super.__init__self.eps = epsself.qwen3_compatible = qwen3_compatibleself.scale = nn.Parameter(torch.ones(emb_dim))self.shift = nn.Parameter(torch.zeros(emb_dim)) if bias else Nonedef forward(self,x):input_dtype = x.dtypeif self.qwen3_compatible:x = x.to(torch.float32)variance = x.pow(2).mean(dim=-1, keepdim=True)norm_x = x * torch.rsqrt(variance + self.eps)norm_x = norm_x * self.scaleif self.shift is not None:norm_x = norm_x + self.shiftreturn norm_x.to(input_dtype)

具体公式可参考：