GLM-4模型中的注意力机制张量形状问题解析

在THUDM团队开发的GLM-4大型语言模型中，研究人员发现了一个与注意力机制实现相关的张量形状不匹配问题。这个问题出现在模型的非SDPA（Scaled Dot-Product Attention）模式下的前向传播过程中。

问题背景

GLM-4模型采用了基于Transformer的架构，其中自注意力机制是其核心组件之一。在实现注意力计算时，模型需要将查询(query)、键(key)和值(value)三个张量进行矩阵运算。具体来说，当不使用优化的SDPA实现时，模型会执行标准的注意力计算流程。

问题定位

在原始的modeling_chatglm.py代码中，第264行附近存在一个张量形状定义问题。当构建注意力输出张量时，代码错误地将value_layer和query_layer的维度顺序进行了错误的排列：

output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))

这种维度排列会导致后续的矩阵运算无法正确执行，因为张量的形状不匹配矩阵乘法的基本要求。

解决方案

正确的张量形状应该调整为：

output_size = (value_layer.size(0), value_layer.size(1), query_layer.size(1), value_layer.size(3))

这一修改确保了：

1. 批量维度(value_layer.size(0))得到正确保留
2. 注意力头数(value_layer.size(1))保持正确
3. 查询序列长度(query_layer.size(1))与键值序列长度正确对应
4. 每个头的特征维度(value_layer.size(3))保持不变

技术影响

这个形状不匹配问题会导致以下后果：

1. 在非SDPA模式下运行时可能引发形状不匹配错误
2. 注意力计算结果不正确，影响模型的语言理解能力
3. 可能导致梯度计算错误，影响模型训练效果

最佳实践建议

对于基于Transformer架构的模型实现，开发人员应当：

1. 仔细检查所有张量操作的输入输出形状
2. 为关键张量操作添加形状断言(assert)
3. 在不同运行模式下(如SDPA/非SDPA)进行全面测试
4. 使用张量形状可视化工具辅助调试

这个问题虽然看似简单，但凸显了在实现复杂神经网络架构时对张量形状管理的重要性。正确的形状处理是确保模型数学运算正确性的基础。