大模型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)。
优化细节:在一些模型(如 Llama, Mixtral)中,这里会使用门控线性单元(Gated Linear Unit),这时可能会有 W_gate 和 W_up 两个上投影矩阵,所以 FFN 部分可能会有 3 个矩阵。但经典的 FFN 是 2 个。
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)。
在硬盘上:模型的权重参数以文件的形式(如 .safetensors 或 .bin)存储在硬盘上。
加载到内存 (RAM):当您启动程序时,这些权重文件首先被读入到计算机的内存 (RAM) 中。
分配到显存 (VRAM):为了利用 GPU 进行高速并行计算,上述所有的权重矩阵(WQ, WK, WV, WO, W_up, W_down, LM Head 等)最终都必须被加载到 GPU 的显存中。
显存是运行大模型最宝贵的资源。一个模型的显存占用主要由以下几部分构成: