本文较长,建议点赞收藏。更多AI大模型应用开发学习视频及资料,在智泊AI。
阿里云推出的 Qwen 3 (通义千问3) 是目前性能最强大的开源模型之一,在多种任务上都表现出色,例如推理、编程、数学和多语言理解。
在这篇博客文章中,我们将从头开始构建一个 0.8 亿参数、包含两个专家层的 Qwen 3 MoE 模型。
模型亮点
- 旗舰版本: Qwen3-235B-A22B 在 MMLU-Pro、LiveCodeBench 和 AIME 等重要基准测试中名列前茅,其性能甚至能与顶尖的专有模型相媲美。1
- 架构创新: Qwen 3 采用了 Mixture-of-Experts (MoE) 架构。这种架构的优势在于,在处理每个查询时,它只会激活部分(而非全部)参数,从而在保证高性能的同时,实现了更高的效率。
- 强大功能:
-
- 支持高达 128K 的上下文长度。
- 支持 119 种语言。
- 引入了独特的“思考”与“非思考”双模式,以平衡深度推理和更快的推理速度。4
前提条件
这篇博客的教程假设读者具备神经网络和Transformer架构的基础知识,并且熟悉Python编程,但无需掌握面向对象编程(OOP)。教程中提供了相关的视频链接作为学习资源。
Qwen 3 MoE(混合专家)架构的核心思想是:与其使用一个庞大且通用的网络来处理所有任务,不如使用一个由多个小型、专业网络(专家) 组成的团队,并由一个 “路由器” 来分配任务。
- 专家(Experts): 这些是较小的、专门化的神经网络,通常是简单的前馈网络(FFNs)或多层感知机(MLPs) 。每个专家擅长处理特定类型的信息或模式。
- 路由器(Router): 这是一个小型网络,充当“管理者”的角色。它的任务是接收输入数据(例如一个词的嵌入向量),然后决定将这个任务分配给哪一个或哪几个最适合的专家。
通过一个句子“the cat sat”来理解MoE的工作流程:
- 分词(Tokenization): 首先,句子被分解成单独的词元(tokens):
"The","cat","sat"。 - 词元嵌入(Token Embedding): 每个词元被转换为一个向量,这个向量包含了词的语义信息。例如,
"cat"被表示为一个高维的嵌入向量。 - 路由器选择专家: 当
"cat"的向量进入MoE层时,路由器会对其进行分析。假设有4个专家(E1, E2, E3, E4),路由器可能会根据"cat"的特征,认为E2(擅长处理名词)和E4(擅长处理动物概念)是最佳选择。 - 分配任务: 路由器会给被选中的专家分配分数或“权重”(例如,E2获得70%的权重,E4获得30%的权重)。然后,
"cat"的向量只会被发送给E2和E4,而E1和E3则无需参与计算,从而节省了大量的计算资源。 - 专家处理: E2和E4分别处理
"cat"的向量,并生成各自的输出(Output_E2和Output_E4)。 - 结果加权合并: 最后,两个专家的输出会根据路由器的权重进行加权合并,得到最终的输出:
Final_Output = (0.7 * Output_E2) + (0.3 * Output_E4)。这个最终输出会继续传递到模型的下一个层。
一个词元在Qwen 3 MoE模型中的完整旅程如下:
- 输入文本(Input Text)进入分词器(Tokenizer) 。
- 分词器生成数字词元ID(Token IDs) 。
- 词嵌入层(Embedding Layer)将ID转换为具有意义的向量(Embeddings) ,并添加位置信息(Positional Info) ,这通常通过RoPE(旋转位置编码)实现。
- 这些向量流经多个Transformer模块。每个模块包含:
-
- 自注意力机制(Self-Attention): 词元之间相互“关注”,其性能通过RoPE得到增强。
- MoE层: 路由器将词元发送给特定的专家。
- 归一化(Normalization): 使用RMSNorm等技术来稳定训练。
- 残差连接(Residual connections): 帮助信息在网络中流动。
- 最后一个模块的输出进入最终层(Final Layer) 。
- 最终层生成词汇表中每个词元的Logits(分数) 。
- Logits被转换为概率,并预测出下一个词元。
环境设置与文件下载
本节介绍如何为构建 Qwen 3 MoE 模型做准备。首先需要安装必要的 Python 库,然后下载模型的核心文件。
# 下载所需的模块
pip install sentencepiece tiktoken torch matplotlib huggingface_hub tokenizers safetensors
本指南将使用一个较小的 Qwen 3 MoE 版本,该版本包含两个专家(Experts),每个专家有0.8亿参数。你可以通过两种方式下载所需的文件:
# 导入tqdm用于进度条,snapshot_download用于下载模型文件
from tqdm import tqdm
from huggingface_hub import snapshot_download
# 定义Hugging Face仓库ID和本地保存目录
repo_id = "huihui-ai/Huihui-MoE-0.8B-2E"
local_dir = "Huihui-MoE-0.8B-2E"
# 从Hugging Face下载模型快照
# 忽略.bin文件,只下载config、tokenizer和safetensors权重
snapshot_download(
repo_id=repo_id,
local_dir=local_dir,
ignore_patterns=["*.bin"], # 跳过大型的.bin文件,只获取safetensors
tqdm_class=tqdm # 使用标准的tqdm显示进度条
)
模型文件与参数
Qwen 3 使用的是 字节对编码 (Byte Pair Encoding, BPE) 算法。BPE 是一种子词分词算法,它通过迭代地将语料库中最频繁的相邻字符对合并成新的词元来构建词汇表。这使得模型能够有效地处理未知词汇并保持词汇表大小适中。
- 用途:将原始文本(如
"the only thing I know is that I know")转换为模型能理解的数字序列([785, 1172, 3166, 358, 1414, 374, 429, 358, 1414]),并能将这些数字序列解码回文本。 - 特点:词汇表大小为 151669,包含通过 BPE 算法合并形成的子词单元。
首先,我们需要将输入的文本转换成模型能够理解的词元(tokens) 。Qwen 3 模型使用特定的聊天模板来构建对话,这有助于模型区分用户的输入和自己的响应。
- 聊天模板:Qwen 3 的模板使用了
<|im_start|>和<|im_end|>等特殊词元来定义对话的边界。例如,一个用户提问的格式通常是<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n。 - 词元化过程:代码首先获取这些特殊词元的 ID,然后将用户的提示 (
"The only thing I know is that I know") 编码成词元 ID 列表。最后,将所有这些 ID 按模板结构拼接成一个完整的序列。 - 结果:一个包含所有词元 ID 的 PyTorch
tensor,代表了完整的对话上下文。
词元化之后,我们需要将每个词元 ID 转换成一个词嵌入向量。词嵌入是一个稠密的、高维度的向量,它捕捉了词元的语义和上下文信息。
- 嵌入层:我们通过
nn.Embedding模块创建一个嵌入层。这个层的尺寸由模型的词汇表大小 (vocab_size) 和隐藏层维度 (dim) 决定。 - 加载预训练权重:为了利用 Qwen 3 模型的知识,我们直接将下载好的预训练权重文件 (
model_weights["model.embed_tokens.weight"]) 加载到我们新创建的嵌入层中。 - 生成嵌入:将词元 ID 的张量输入到这个嵌入层,就会得到一个形状为
[17, 1024]的张量,其中 17 是词元数量,1024 是词嵌入的维度。
在将词嵌入向量输入到 Transformer 层之前,对其进行归一化是一个关键步骤,它能确保训练的稳定性和效率。Qwen 3 使用 RMSNorm (Root Mean Square Layer Normalization) 来完成这个任务。
- RMSNorm 函数:
rms_norm函数通过计算输入张量的均方根 (RMS) 来进行归一化。它将输入张量乘以其均方根的倒数,从而使归一化后的张量具有单位方差。 - 作用:RMSNorm 简化了传统的 LayerNorm,省略了减去均值的步骤,只进行重新缩放,这在保持性能的同时,提供了更高的计算效率。
- 应用:我们使用模型权重中第一层 (
model.layers.0.input_layernorm.weight) 的归一化权重,对生成的词嵌入向量进行归一化。这为接下来的注意力机制做好了准备。
分组查询注意力(Grouped-Query Attention, GQA)
Qwen 3 模型采用 GQA 机制来优化计算效率。传统的自注意力机制中,每个查询头(Query head)都有一个对应的键头(Key head)和值头(Value head)。而 GQA 允许多个查询头共享一小组键头和值头。
- 实现方式:
-
- 模型有 16 个查询头 (
n_heads = 16),因此q_proj权重被重塑为[16, 128, 1024]。 - 模型有 8 个键头和值头 (
n_kv_heads = 8),所以k_proj和v_proj权重被重塑为[8, 128, 1024]。
- 模型有 16 个查询头 (
- 计算:我们首先从这些重塑后的权重中,提取第一个查询头和第一个键/值头的权重,然后将归一化后的词嵌入向量与这些权重相乘,得到每个词元对应的 Q、K、V 向量。
# Get the pre-computed rotations for our sequence of 17 tokens
freqs_cis_for_tokens = freqs_cis[:len(tokens)]
# --- Apply RoPE to Query vectors ---
# Reshape [17, 128] to [17, 64, 2] and view as complex numbers [17, 64]
q_per_token_as_complex_numbers = torch.view_as_complex(q_per_token.float().view(q_per_token.shape[0], -1, 2))
# Apply rotation via complex multiplication
q_per_token_rotated_complex = q_per_token_as_complex_numbers * freqs_cis_for_tokens
# Convert back to real numbers and reshape to [17, 128]
q_per_token_rotated = torch.view_as_real(q_per_token_rotated_complex).view(q_per_token.shape)
# --- Apply RoPE to Key vectors (same process) ---
k_per_token_as_complex_numbers = torch.view_as_complex(k_per_token.float().view(k_per_token.shape[0], -1, 2))
k_per_token_rotated_complex = k_per_token_as_complex_numbers * freqs_cis_for_tokens
k_per_token_rotated = torch.view_as_real(k_per_token_rotated_complex).view(k_per_token.shape)
print("Shape of rotated Query vectors:", q_per_token_rotated.shape)
#### OUTPUT ####
# Shape of rotated Query vectors: torch.Size([17, 128])
#### OUTPUT ####
在计算注意力之前,需要为 Q 和 K 向量注入位置信息,以使模型能够理解词元在序列中的顺序。Qwen 3 使用 RoPE (Rotary Position Embedding) 来实现这一功能。
- RoPE 原理:RoPE 通过旋转 Q 和 K 向量的每一对维度来编码位置信息。每个 2D 维度对(例如,维度0和1,维度2和3)的旋转角度会根据其位置和维度而变化。
- 实现步骤:
-
- 预先计算一个包含旋转角度的查找表 (
freqs_cis)。这个表包含了所有可能位置的旋转复数。 - 将 Q 和 K 向量重塑为复数形式。
- 通过复数乘法,将预计算的旋转角度应用到 Q 和 K 向量上。
- 将旋转后的复数向量转换回实数向量。
- 预先计算一个包含旋转角度的查找表 (
现在,我们可以计算每个词元对其他所有词元的重要性。
- 点积:将旋转后的 Q 向量与 K 向量的转置相乘,得到一个
[17, 17]的矩阵,其中每个元素表示一个词元对另一个词元的注意力分数。 - 缩放:将分数除以
head_dim的平方根 (128**0.5),以防止内积过大,导致 softmax 函数梯度消失。 - 因果掩码(Causal Masking) :由于我们构建的是一个自回归模型,词元不应该“看到”未来的词元。因此,我们创建一个上三角矩阵,并将对角线以上的所有元素设置为负无穷,从而确保模型在预测当前词元时只关注其前面的词元。
- Softmax:将掩码后的分数应用 Softmax 函数,将其转换为 0 到 1 之间的概率,这些概率被称为注意力权重。
实现多头注意力(Multi-Head Attention)
多头注意力机制通过并行运行多个独立的注意力计算来捕捉输入序列中不同位置和方面的信息。
- 并行计算 :代码通过一个
for循环遍历所有的16个注意力头。在每个循环中,它提取当前头的Q、K、V权重,应用RoPE位置编码,计算注意力分数和权重,并最终生成一个大小为[17, 128]的输出张量。 - GQA :由于采用了 分组查询注意力 (GQA),
k_proj和v_proj权重只被分为8个共享的头。代码通过head // (n_heads // n_kv_heads)(即head // 2)来确保每两个查询头共享一个键/值头。 - 拼接与投影 :所有16个头的输出被收集在一个列表中,然后通过
torch.cat沿着最后一个维度拼接成一个更大的张量[17, 2048]。这个张量再通过一个输出投影层 (o_proj) 投影回模型的隐藏层维度[17, 1024]。 - 残差连接 :投影后的结果 (
embedding_delta) 被加回到注意力层的输入 (token_embeddings_unnormalized),形成了第一个残差连接。这有助于训练非常深的网络。
# --- Step 1: The MoE Router ---
# The router is a simple linear layer that determines which expert to send each token to.
# It projects our [17, 1024] tensor to a [17, num_experts] tensor of scores (logits).
gate = model_weights["model.layers.0.mlp.gate.weight"]
router_logits = torch.matmul(embedding_after_attention_normalized, gate.T)
# We apply softmax to the logits to get probabilities, and then find the expert with the
# highest probability for each token.
routing_weights = torch.nn.functional.softmax(router_logits.float(), dim=1).to(torch.bfloat16)
routing_expert_indices = torch.argmax(routing_weights, dim=1)
print("Router logits shape:", router_logits.shape)
print("Expert chosen for each of the 17 tokens:", routing_expert_indices)
# --- Step 2: The Expert Layers ---
# Each expert is a SwiGLU-style Feed-Forward Network.
expert0_w1 = model_weights["model.layers.0.mlp.experts.0.gate_proj.weight"]
expert0_w2 = model_weights["model.layers.0.mlp.experts.0.down_proj.weight"]
expert0_w3 = model_weights["model.layers.0.mlp.experts.0.up_proj.weight"]
expert1_w1 = model_weights["model.layers.0.mlp.experts.1.gate_proj.weight"]
expert1_w2 = model_weights["model.layers.0.mlp.experts.1.down_proj.weight"]
expert1_w3 = model_weights["model.layers.0.mlp.experts.1.up_proj.weight"]
# --- Step 3: Process tokens with their chosen expert ---
final_expert_output = torch.zeros_like(embedding_after_attention_normalized)
#### OUTPUT ####
# Router logits shape: torch.Size([17, 2])
# Expert chosen for each of the 17 tokens: tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
#### OUTPUT ####
MoE是Qwen 3架构的核心,它在注意力层之后被应用。
- 预归一化:首先,对注意力层输出进行 RMSNorm 归一化,为MoE层做准备。
- 路由器(Router) :一个名为
gate的线性层充当路由器。它将归一化后的嵌入向量([17, 1024])投影成一个包含专家分数的张量([17, 2])。然后,通过softmax获得每个词元分配给每个专家的概率,并使用argmax选出概率最高的专家。 - 专家(Experts) :每个专家都是一个 SwiGLU 风格的前馈网络。代码加载了两个专家的权重,并为每个词元循环处理。
- 加权求和:每个词元只会被其被选中的专家处理。专家的输出会乘以其对应的路由器概率,最终被整合到
final_expert_output张量中。 - 残差连接:MoE模块的输出被加回到其输入 (
embedding_after_attention),形成了第二个残差连接,完成了整个Transformer层。
# The LM Head weights are the same as the embedding weights (weight tying)
lm_head_weights = model_weights["model.embed_tokens.weight"]
# We only care about the last token's output to predict the next token
last_token_embedding = final_embedding_normalized[-1]
# Calculate the logits by multiplying with the LM Head
logits = torch.matmul(last_token_embedding, lm_head_weights.T)
print("Shape of the final logits:", logits.shape)
#### OUTPUT ####
# Shape of the final logits: torch.Size([151936])
#### OUTPUT ####
最后,所有的步骤被整合到一个主循环中,循环遍历模型的28个层。
- 层循环:
for layer in range(n_layers)循环负责处理模型的每一层。在每个循环中,它都执行完整的注意力子层和MoE子层。 - 数据流:前一层的输出 (
final_embedding) 会作为下一层的输入,并在每个子层中通过残差连接不断更新。 - 结果:经过所有28个层的处理后,我们得到一个形状为
[17, 1024]的最终嵌入向量,这个向量包含了所有词元的完整上下文信息,并准备好进入最终层进行下一个词元的预测。
# Find the token ID with the highest score
next_token_id = torch.argmax(logits, dim=-1)
print(f"Predicted Token ID: {next_token_id.item()}")
# Decode the ID back to a string to see the predicted word
predicted_word = tokenizer.decode([next_token_id.item()])
print(f"\nPredicted Word: '{predicted_word}'")
#### OUTPUT ####
# Predicted Token ID: 12454
# Predicted Word: 'nothing'
#### OUTPUT ####
学习资源推荐
如果你想更深入地学习大模型,以下是一些非常有价值的学习资源,这些资源将帮助你从不同角度学习大模型,提升你的实践能力。
本文较长,建议点赞收藏。更多AI大模型应用开发学习视频及资料,在智泊AI。