Page 220 - 《软件学报》2026年第1期
P. 220
高彦杰 等: 大语言模型预训练系统关键技术综述 217
Q ( K T
)
( ) QW E i KW
Q k v V i i V
· F i VW .
overlinehead i = Attention QW ,E i KW ,F i VW = ¯ P· F i VW = Softmax √
i i i i i
d k
( )
V
使用 ¯ P· F i VW 为第 i 个头部计算上下文嵌入. 需要注意的是, 上述操作仅需要 O(nk) 的时间和空间复杂度.
i
因此, 如果我们能够选择一个较小的投影维度 k, 使得 k 远小于 n, 就能够显著减少内存的消耗.
近似与稀疏化和分桶本质上是对于 QK 矩阵进行近似计算, 只需考虑局部元素; 低秩分解与降维则通过降维
来减小输出中间结果矩阵的维度. 另外, 一些工作, 例如 Long-short Transformer [71] , 采用了上述 3 种方法的混合优
化策略. 表 5 对上文介绍的 4 类高效注意力机制进行对比, 分析了其优势、不足和适用场景.
表 5 高效注意力机制对比
高效注意力类型 优势 不足 适用场景
节省访存开销. 对准确度几乎没有
访存优化 影响. 该方法与其他方案正交, 可 仅能节省访存开销. 在计算密集型算 注意力访存为瓶颈的场景
子场景下, 优化效果可能降低
以结合使用而互不干扰
该方法存在一定的信息损失. 不适用
近似与稀疏化 计算效率和访存效率均有提升 于注意力分数矩阵较为稠密的场景. 注意力分数矩阵数据稀疏. 适用于
仅需考虑局部注意力的场景
近似方法的实现需要软硬件协同设计
该方法仅适用于需要考虑局部注意
该方法对数据变化较为敏感. 该方法
分桶 计算效率和访存效率均有提升 力上下文的场景. 数据分桶较为均
可能会导致分桶倾斜
匀, 有助于负载的均衡分配
有一定信息损失. 需要优化配置降维 在注意力分数矩阵为低秩的场景
低秩分解与降维 计算效率和访存效率均有提升
维度 下, 效果更加显著
3.3.2 混合精度训练
在大型语言模型的训练中, 显存和加速器核心是有限的资源. 使用低精度数据类型可以显著降低计算和内存
开销. 混合精度训练引入了半精度浮点数 (例如 Float16、BFloat16 等) 来训练大语言模型, 无需修改超参数, 尽可
能不损失模型精度的情况下, 提升训练性能和效率. 例如, 将 Float32 替换为 Float16 数据类型, 这会将相应张量的
内存需求减半, 并且在 NVIDIA GPU 等最新架构上, 可以利用张量核心等特殊计算单元进一步加快计算速度.
在混合精度训练中, 通常使用半精度格式存储权重张量、激活张量和梯度张量. 然而, 由于这种格式的数值范
围比单精度更窄, 可能导致信息丢失或收敛性问题. 为了解决这个问题, Micikevicius 等人 [72] 提出了 3 种防止关键
信息丢失的技术. 首先, 建议维护一个累积的单精度权重张量副本, 每个优化器步骤后得到的梯度 (在前向和反向
传播中, 此副本四舍五入为半精度). 其次, 使用损失缩放以保留较小幅度的梯度值. 最后, 使用半精度算术累加并
转换为单精度输出, 在存储到内存之前将其转换为半精度. 这些技术的结合有助于克服混合精度训练中可能出现
的数值范围狭窄引起的问题. 目前, 混合精度训练方法也被广泛应用于大型语言模型的训练, 例如在 Llama 2 的微
调方案中, 用户仍然可以选择混合精度的配置.
如图 9 所示, 在混合精度训练中, 主模型权重首先被转换为 16 比特的 Half 类型, 然后与 Float16 (FP16) 类型
的激活一起进行前向传播, 生成激活. 在反向传播计算中, 计算激活梯度时使用 Float16 类型的输入权重和激活梯
度, 生成 16 比特的激活梯度; 计算权重梯度时需要 16 比特的激活和激活梯度, 计算出 16 比特的权重梯度, 最终更
新模型权重. 这样, 大部分核心计算和显存消耗较大的张量能够保持低精度, 从而加速训练并节省显存.
FP16
权重 FP16
FP16 前向传播 激活
激活
Float to
half
FP16 FP16
激活梯度 权重
FP16 反向 FP16 激活梯度
权重梯度 传播-权重 FP16
激活
FP32 FP32
主权重 (FP32) 权重更新 更新后的主权重
图 9 混合精度训练 [72]

