立志写出让纯夏都能看懂的教程讲解

我们继续自顶向下学习 nano-vllm 的架构，本章我们主要解析其自带支持的 qwen3-8b 的模型结构。本章我们主要解析其自带支持的 qwen3-8b 的模型结构。我们将首先介绍 qwen3 的模型结构，然后从代码入手快速学习了解整个前向推理过程，至于具体的基础算子实现我们放在后面再讲。

模型架构

先简单对 Qwen3 模型有个概念，整体其实很好理解，就是一个 Decoder-Only 的模型，每个 Decoder 层都有：

  • RMS Norm + Attention：这里 Attention 是 Group Query Attention，顾名思义，就是多个 Q 头会对应相同的 K, V 头； 比较特殊的是，在具体计算 QKV 时，会对 Q, K 做 per-head 的 RMS Norm 归一化；
  • RMS Norm + MLP：这里做了个相较于普通的直接两层 Linear，使用了门控网络来让模型自己学习动态选择特征的能力，并且用 SiLU 做激活函数，整体思路比较 intuition。当 gate 值接近 0 时，SiLU 接近于 0，代表该特征被抑制，反之，当 gate 很大，SiLU 接近于 gate，该特征值正常通过。

qwen3.png

我们在 qwen3.py 的代码里，可以很直接的找到上面这个模型结构对应的代码细节。接下来，我们将继续自顶向下的顺序，对代码进行讲解，并在其中适当地掺杂原理解析。

Qwen3ForCausalLM & Qwen3Model

Qwen3ForCausalLM 就是 ModelRunner 使用的接口类了，封装了我们上面图画的整个 Qwen3Model 模型结构，包括词嵌入层、多个解码器层和最终的 RMSNorm 层。注意下，LM Head 和 Qwen3Model 是分开的，Qwen3Model 这个类可以看作是一个通用的特征提取器。

Qwen3Model 这个类实际上就是我们上图模型的拼装：

  • 首先， VocabParallelEmbedding 将 token ids 转换为 word embedding，这是一个 [vocab_size, hidden_size] 大小的矩阵。为了支持多卡并行加速，所以我们很自然能够想到将这个 embedding 矩阵按行切割，平均划分个多卡进行张量并行，每卡只负责一部分内容，此时维度 [vocab_size / tp_size, hidden_size]。【你可能要问，那多卡并行的时候，token_id 怎么知道要从哪个卡里获得对应的 embedding 呢？这个问题我们留到下章讲解。】；
  • 然后，Qwen3DecoerLayer 堆叠的 layer 将处理 embedding，和上图一样，每个 decoderLayer 都包含一个 self-attention 模块和一个 MLP 模块；
  • 最后，RMSNorm 对最终的 hidden_states 进行归一化。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class Qwen3Model(nn.Module):

    def __init__(
        self,
        config: Qwen3Config,
    ) -> None:
        super().__init__()
        # word embedding 层，列并行实现，词表被划分到不同的 GPU 上，每个 GPU 只处理自己负责的词表部分
        self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size)
        # decoderLayer 堆叠，每个 decoderLayer 包含一个 self-attention 模块和一个 MLP 模块
        self.layers = nn.ModuleList([Qwen3DecoderLayer(config) for _ in range(config.num_hidden_layers)])
        # 最后一个 RMSNorm 层，应用在所有 decoderLayer 之后，对最终的 hidden_states 进行归一化
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
    ) -> torch.Tensor:
        hidden_states = self.embed_tokens(input_ids)
        residual = None
        for layer in self.layers:
            hidden_states, residual = layer(positions, hidden_states, residual)
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states

class Qwen3ForCausalLM(nn.Module):
    packed_modules_mapping = {
        "q_proj": ("qkv_proj", "q"),
        "k_proj": ("qkv_proj", "k"),
        "v_proj": ("qkv_proj", "v"),
        "gate_proj": ("gate_up_proj", 0),
        "up_proj": ("gate_up_proj", 1),
    }

    def __init__(
        self,
        config: Qwen3Config
    ) -> None:
        super().__init__()
        # Qwen3Model 包含了整个 Transformer 模型的结构，包括词嵌入层、多个解码器层和最终的 RMSNorm 层
        self.model = Qwen3Model(config)
        # LM 头，维度为 (vocab_size, hidden_size)，
        # 用于将模型的输出 hidden_states 投影到词表大小的维度上，生成 logits
        self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
        if config.tie_word_embeddings:
            self.lm_head.weight.data = self.model.embed_tokens.weight.data

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
    ) -> torch.Tensor:
        return self.model(input_ids, positions)

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        return self.lm_head(hidden_states)

Qwen3Attention

继续拆解 Qwen3Model 模块，根据上文描述，它是由 Qwen3Attention 和 Qwen3MLP 组合实现的。

Q,K,V 计算本身很简单，就是一个 Linear 层映射，这里为了充分利用多卡张量并行，我们在 QKVParallelLinear 的实现里将 Q,K,V 的矩阵映射计算合并在了一起计算，合并后的 Q K V 权重形状为 $(num\_heads \times head\_dim + num\_kv\_heads \times head\_dim \times 2, hidden\_size)$ (Pytorch 按照 [Output Size, Input Size] 存储，F.linear(x, W) 计算的时候权重矩阵会转置 x @ W.T ) 的线性层，这个矩阵在计算的时候，会按照转置后列平均切分给多卡并行计算，合并后输出就是 $[Q, K, V]$ 矩阵；

在 Multi-Head Attention 计算完后，我们通常会有一个 Linear 层，将多头的维度重新压缩回 hidden_size 的大小，这一层 Linear 的维度是 $(hidden\_size, total\_num\_heads \times head\_dim)$。同样的道理，为了充分利用多卡张量并行，我们在 attention 计算完后的 Linear 层也是类似的处理：按照转置后行平均切分给多卡做并行计算，各个卡计算完成后 all-reduce 通信求和获得完整的输出。

Qwen3 模型比较特殊的一点在于，其计算完 Q, K 后引入了 RMSNorm 进行归一化。

Group Query Attention

在代码里，我们会看到 Multi-Head Attention 的计算中，Q 头和 K, V 头数量并不一样，这实际上是 Group Query Attention：多个 Q 会共享相同的 K, V，相比原始的 Multi-Head 这么做能在保持性能的同时有效降低 KV-Cache 的压力。实际上我觉得和巻积里的 Depth-wise Convolution 思路基本是差不多，对 CV 模型比较熟悉的同学肯定很眼熟。

GQA.png

RoPE 位置编码

还有一个值得讲的是 RoPE 位置编码的使用，这里我们暂且按住不表，就暂且简单总结为：

有兴趣的朋友可以先直接参考苏老师在科学空间写的博客：

https://www.spaces.ac.cn/archives/8265

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
class Qwen3Attention(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        max_position: int = 4096 * 32,
        head_dim: int | None = None,
        rms_norm_eps: float = 1e-06,
        qkv_bias: bool = False,
        rope_theta: float = 10000,
        rope_scaling: tuple | None = None,
    ) -> None:
        super().__init__()
        tp_size = dist.get_world_size()
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        # 计算每个 GPU 的 Q 头 (多 GPU 并行)
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = num_kv_heads
        assert self.total_num_kv_heads % tp_size == 0
        # 计算每个 GPU 的 KV 头 (多 GPU 并行)
        self.num_kv_heads = self.total_num_kv_heads // tp_size
        self.head_dim = head_dim or hidden_size // self.total_num_heads
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim ** -0.5
        self.qkv_bias = qkv_bias

        # QKV 算子融合
        # 合并后的 QKV 权重形状为 (hidden_size, num_heads * head_dim + num_kv_heads * head_dim * 2)
        # 列并行线性层，输出被分成三部分：Q、K 和 V
        self.qkv_proj = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=qkv_bias,
        )

        # 输出投影
        # 输入维度为 total_num_heads * head_dim，输出维度为 hidden_size
        # 行并行线性层，输入被分成 num_heads * head_dim 的块，每个块独立进行线性变换
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
        )

        # RoPE 位置编码
        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position,
            base=rope_theta,
            rope_scaling=rope_scaling,
        )
        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            self.num_kv_heads,
        )

        # Qwen3 独有特性: 如果不使用 QKV 偏置，则对 Q 和 K 进行独立的 RMSNorm
        if not self.qkv_bias:
            self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
            self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        # 输入 hidden_states 形状为 (batch_size, seq_len, hidden_size)
        # 一次性计算 Q、K、V，
        # 输出形状分别为 Q(batch_size, seq_len, num_heads * head_dim)，
        # 和 K,V(batch_size, seq_len, num_kv_heads * head_dim)
        qkv = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q = q.view(-1, self.num_heads, self.head_dim)
        k = k.view(-1, self.num_kv_heads, self.head_dim)
        v = v.view(-1, self.num_kv_heads, self.head_dim)

        # Qwen3 独有特性: 如果不使用 QKV 偏置，则对 Q 和 K 进行独立的 RMSNorm
        if not self.qkv_bias:
            q = self.q_norm(q)
            k = self.k_norm(k)
        
        # 应用 RoPE 位置编码
        q, k = self.rotary_emb(positions, q, k)
        o = self.attn(q, k, v)
        # 输出投影，并行计算，内部包含了 all-reduce 操作
        output = self.o_proj(o.flatten(1, -1))
        return output

Qwen3MLP

作为 Qwen3 的 FFN，这里基本就是一个门控网络的思路。和前面 Attention 类似，为了利用张量并行，我们将 gate 和 up 两个 Linear 映射的权重网络合并在一起进行计算。同时，对于模型 $SiLU(gate) * up$ 的操作，我们也正好融合为了一个算子 SiluAndMul 一起进行。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Qwen3MLP(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
    ) -> None:
        super().__init__()

        # 算子融合：将 gate_proj 和 up_proj 合并为一个 MergedColumnParallelLinear
        # gate_proj 和 up_proj 的输出维度分别为 intermediate_size 和 intermediate_size，合并后输出维度为 intermediate_size * 2
        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size,
            [intermediate_size] * 2,
            bias=False,
        )

        # 输出投影
        # 用于将 intermediate_size 维的激活输出投影回 hidden_size 维
        self.down_proj = RowParallelLinear(
            intermediate_size,
            hidden_size,
            bias=False,
        )
        assert hidden_act == "silu"
        # 算子融合：将 SiLU(x) * y 融合为一个 SiluAndMul，减少内存占用和计算开销
        self.act_fn = SiluAndMul()

    def forward(self, x):
        # gate_up_proj 的输出被分成两部分：gate 和 up
        gate_up = self.gate_up_proj(x)
        # 激活函数融合：同时对 gate 和 up 应用 SiLU 激活，并计算 gate * up
        x = self.act_fn(gate_up)
        # 输出投影，降维回 hidden_size
        x = self.down_proj(x)
        return x

说些什么吧！

beaudar