LLM之GLM模型

1 简介

GLM模型是清华大学提出的一个LLM, 这个模型在中文数据集上的效果表现比较不错, 网络上也有一些针对这个模型的开发。
Paper: 《GLM: General Language Model Pretraining with Autoregressive Blank Infilling》
Repo:

2 算法原理

阅读GLM论文后,我总结了GLM的3个Feature:

  • 独特的训练方式 / 优化目标: 这理解也是论文的核心卖点, 讲了一个“general” language model的故事。而所谓的general, 是指一个预训练好的模型,可以同事在NLU、条件生成和无条件生成任务中均表现很好。

  • 一些模型结构的微调: 1)对layernorm和残差连接的顺序做了调整(将layernorm放在Block前可以提高大模型的精度及训练的稳定性 Ref: 参考文献1); 2)将Block内MLP中两个FC间的激活函数替换成了GELU(而非其它模型使用的ReLU)。GLM的Block内是Self-Attention + MLP组成, Self-Attention通常有两个矩阵参数, 一个是QKV矩阵用于计算qkv向量, 一个参数矩阵用于做特征变换(就是一个FC做矩阵乘); MLP通常有两个参数矩阵做特征变换(2次矩阵乘), 先升维,再降维, 有点借鉴ResNet bottleneck的思想。

  • 2D位置编码: 服务于Feature 1, 使用Rotary Position Encoding。

2.1 模型结构简图

GLM-Block也是基本的Attention-Block + GLU-Block.
glm-block

2.2 模型训练

数据处理:将输入tokens随机mask掉一些text span, 然后将mask掉的这些token用一个特殊token表示, 从而形成输入到GLM模型中的corrupted token ids. 这个过程用论文中的这张图可以很好的表达:
VNMgMc
优化函数: 使用的是最大似然估计, 最大化生成概率的对数和。其中是所有m个被mask掉的text span的排列序号的集合, 一共可以有中排列方案; $p_{\theta}(s_{z_i}|x_{corrupt},s_{z})ix_{corrupt}i-1 maxE_{z\sim Z_{m}} [\Sigma_{i=1}^{m}\log P_{\theta}(s_ {z_{i}}|x_{corrupt},s_{z<i})$]

GLM模型使用了P-tuning V2这一fine-tuing schema, 这个方法也是清华Jie Tang老师团队的工作。

2.3 Rotary位置编码

GLM工作使用的位置编码是国内追一科技提出的旋转位置编码(Rotary Position Embedding), 这一工作还是比较受领域内人员的认可, 像LLaMA也是采用了这种位置编码方案。

RoPE借助attention的计算, 实现了将绝对位置编码转换成相对位置编码。具体来说, 如果使用表示绝对位置编码, 我们期望找到一个合适的映射关系, 使其满足:

经过一系列推导, 得到的变换可以用下面的计算表达:
ixh37D
这是一个分块对角矩阵, 而每个分块又是一个2D旋转矩阵, 所以这种位置编码方式叫做旋转位置编码。另外, 由于这个矩阵是一个稀疏矩阵, 处于计算效率的考量, 可以用下面的计算来实现。
YdNCY1

不同于原论文的旋转位置编码方法, GLM做了一点细微的改动: 每个token有两个position_id, 第一个id表示在输入中的index,第二个id表示在text span内部的index, 这两个id可以参考图2(c)中的例子。所以对一个128维的特征, GLM将它split成2个64维特征, 分别使用这两个id去计算旋转位置编码。

3 张量并行方案

使用的是NV的Megatron方案, Megatron并行的方案可以参考文献1. 如果要高度概括一下, 可以总结成以下三点:

  • 对Self-Attention, 在head_num的维度上切分进行张量并行, 即每张卡都对部分头进行注意力计算;
  • MLP的两层FC权重, 一个横切 + 一个纵切, 这样的切法可以减少一次通信次数;
  • 每个Block内需要进行两次AllReduce操作, 一次是在attention_dense后, 一次是在MLP的最后一层FC后。
    nPWcRC

如果约定一下符号, Self-Attention的切分可以表示成:

  • [b, s, h] x [h, 3h / p] = [b, s, 3h / p]: h是n个头的特征维数, 每张卡上负责n / p个头的注意力计算, 所以特征数是 n / p * h / n = h / p;
  • [b, s, h / n] x [b, s, h / n].transpose = [b, s, s]: 每张卡都会为每个头计算attention score
  • [b, s, s] x [b, s, h / n] = [b, s, h / n]: 每个头都计算出自己的value, 然后每张卡内的value值concatenate起来, 得到[b, s, h / p], (对应上图的Y1或Y2).
  • [b, s, h /p] x [b, h / p, h] = [b, s, h], (对应上图的Z1或Z2)
  • 将各张卡上的结果做一次all_reduce_sum

4 优化手段

  • 针对特定硬件的算子优化
  • 内存复用
  • 算子融合
  • 数据类型, e.g. INT8 / INT4

5 参考资料

1 《Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism》
2 博采众长的旋转位置编码

Comments

Unable to load Disqus, please make sure your network can access.