大模型decoder中权重矩阵的理解 - 实践

一个 Decoder Block 核心由两个核心子层构成：自注意力层 (Self-Attention) 和 前馈神经网络层 (Feed-Forward Network)。

1. 自注意力层 (Masked Self-Attention Layer)

这是模型“看”输入序列并决定关注哪些部分的地方。它的参数主要体现在4 个权重矩阵上：

• 查询 (Query) 权重矩阵 (WQ)：将输入向量转换为一个“查询”向量。

• 键 (Key) 权重矩阵 (WK)：将输入向量转换为一个“键”向量。

• 值 (Value) 权重矩阵 (WV)：将输入向量转换为一个“值”向量。

• 输出 (Output) 权重矩阵 (WO)：在注意力计算完毕后，将结果整合并投影回原始维度。

优化细节三个独立的权重矩阵。就是：在实际实现中，为了计算效率，Q、K、V 这三个矩阵常常被合并成一个巨大的 W_qkv 矩阵一次性进行计算，然后再切分开。但从概念上讲，它们

2. 前馈神经网络层 (Feed-Forward Network, FFN)

这是模型进行“思考”和非线性变换的地方，它通常包含2 个 权重矩阵：

• 上投影 (Up-Projection) 权重矩阵 (W_up)：将注意力层的输出从模型的隐藏维度（例如 4096）扩展到一个更大的中间维度（例如 11008）。

• 下投影 (Down-Projection) 权重矩阵 (W_down)：将扩展后的向量重新投影回模型的原始隐藏维度（例如 4096）。

优化细节 2 个。就是：在一些模型（如 Llama, Mixtral）中，这里会使用门控线性单元（Gated Linear Unit），这时可能会有 W_gate 和 W_up 两个上投影矩阵，于是 FFN 部分可能会有 3 个矩阵。但经典的 FFN

3. 归一化层 (Layer Normalization)

Decoder Block 中通常还有两个 LayerNorm 层，一个在自注意力层之前，一个在前馈网络层之前。它们没有大的权重矩阵，但有可学习的参数：

• Gamma (γ)：一个缩放向量。

• Beta (β)：一个偏移向量。
  它们的参数量远小于权重矩阵。

总结一个 Decoder Block 的参数：

组件	主要权重矩阵/参数	数量	主要功能
自注意力	W_q, W_k, W_v, W_o	4 个矩阵	关联序列中的不同位置信息
前馈网络	W_up, W_down	2 个矩阵	进行非线性计算，增强模型表示能力
归一化层	Gamma (γ), Beta (β)	2 组向量	稳定训练过程

一个完整的 GPT 模型是由很多个这样的 Decoder Block 堆叠起来的（例如 Llama-7B 有 32 个），再加上一些其他参数。

最终的输出层 (LM Head)

在所有 Decoder Block 的最顶层，还有一个非常重点的权重矩阵：

• 语言模型头 (LM Head)：它负责将最后一个 Decoder Block 的输出向量，投影到整个词汇表的大小（例如 50257）。这个矩阵的输出经过 Softmax 后，就得到了下一个词的概率分布。

优化细节 (权重绑定 Weight Tying)：为了节省参数，很多模型的输入词嵌入矩阵 (Input Embedding Matrix) 和这个最终的 LM Head 矩阵是共享权重的。

参数主要在哪个上面被分配了空间？

您的第二个疑问非常关键，答案是GPU 的显存 (VRAM)。

1. 在硬盘上：模型的权重参数以文件的形式（如 .safetensors 或 .bin）存储在硬盘上。

2. 加载到内存 (RAM)：当您启动程序时，这些权重文件先被读入到计算机的内存 (RAM) 中。

3. 分配到显存 (VRAM)：为了利用 GPU 进行高速并行计算，上述所有的权重矩阵（WQ, WK, WV, WO, W_up, W_down, LM Head 等）最终都必须被加载到 GPU 的显存中。

显存是运行大模型最宝贵的资源。一个模型的显存占用主要由以下几部分构成：