FP16 Online Softmax 误差分析与混合精度优化

FP16 Online Softmax 误差传播与混合精度

在这里我们先简单介绍一下Online Softmax的意义：
Online Softmax把传统Softmax的two-pass优化成一次扫描，两者时间复杂度都是\(O(n)\)，并且在只做Softmax的情况下，由于Online Softmax比传统Softmax要多做些计算操作，导致Online Softmax比传统Softmax的性能还低，但从算子融合的角度来看就不一样了，如果矩阵的维度比较大的话，一行需要多个block进行覆盖，因为传统Softmax需要全局最大值，所以这就需要我们进行规约操作了，这就代表gemm算完一块block之后直接算Softmax，节省了读取内存的时间，但在最后计算完毕后，要进行一个全局规约，规约浪费的时间可能会比我们之前节省的读取时间还要长，所以在这里就体现出了Online Softmax的优势，它不需要进行全局规约。

理论分析

接下来，看Online Softmax 的递推公式：

\[\begin{aligned} m_i &= \max(m_{i-1}, x_i) \\ d_i &= d_{i-1} \cdot e^{m_{i-1}-m_i} + e^{x_i - m_i} \end{aligned} \]

其中 \(d_0 = 0, m_0 = -\infty\)。
这里引入精度模型
FP16 舍入单位：\(u = 2^{-11} \approx 4.88 \times 10^{-4}\)。
对任意实数 \(z\)，FP16 舍入满足：

\[\text{fl}(z) = z(1+\delta), \quad |\delta| \leq u \]

这里考虑一个简单情况，假设 \(m_i,x_i\)用 FP32 精确存储，仅 \(d_i\) 用 FP16 存储和计算。
则计算值 \(\hat{d}_i\) 满足：

\[\hat{d}_i = \text{fl}\Big( \text{fl}\big(\hat{d}_{i-1} \cdot \text{fl}(e^{m_{i-1}-m_i})\big) + \text{fl}(e^{x_i-m_i}) \Big) \]

这里我们用一阶扰动对其进行分析：
令 \(\hat{d}_i = d_i + \epsilon_i\)，记 \(\gamma_i = e^{m_{i-1}-m_i}\)，\(\alpha_i = e^{x_i-m_i}\)。

根据上面的式子，展开并保留一阶项，发现能把\(d_i\)消掉，于是得到下式：

\[\epsilon_i \approx \gamma_i \epsilon_{i-1} + \underbrace{d_{i-1}\gamma_i \delta_{\text{total}} + \alpha_i \delta_{\text{exp+add}}}_{\eta_i} \]

其中 \(|\delta_{\text{total}}| \lesssim 3u\)（mul+exp+add 的累积），\(|\eta_i| \lesssim 3u \cdot d_i\)。

把递推项展开，由三角不等式得到上界：

\[|\epsilon_n| \leq \sum_{i=1}^n |\eta_i| \prod_{j=i+1}^n \gamma_j \]

通过观察我们发现，\(\gamma_j = e^{m_{j-1}-m_j} \in (0,1]\)。

• 当 \(x_j \gt m_{j-1}\) 时，\(\gamma_j <1\)，历史误差被压缩
• 当 \(x_j \leq m_{j-1}\) 时，\(\gamma_j = 1\)，误差无损传递

最坏情况是：当输入序列为单调递减，此时对所有 \(j\) 有 \(\gamma_j = 1\)，误差线性累积满足：

\[\frac{|\epsilon_n|}{d_n} \lesssim 3 \cdot n \cdot u \]

通过该式得知，相对误差的上界对n很敏感，当我们输出一个比较大的n时，相对误差会变得比较大。这时候我们就会想办法去压缩这个误差，我们想，如果我们去算FP32下的误差，那不就是把FP16的结果中的\(u\)改成FP32情况下的\(u\)，这就可以在一定的误差下“容下”更大的n。
于是，我们试着把输出的数据用FP16精度，输入数据和中间计算用FP32精度，这其实就是flash attention中的场景，我们来分析误差：
这个误差分析其实很显而易见，跟FP16的情景差不多，只是变成了FP32下的\(u\)，有下面式子：

\[\hat{d}_i = \text{fl}_{u32}\Big( \text{fl}_{u32}\big(\hat{d}_{i-1} \cdot \text{fl}_{u32}(e^{m_{i-1}-m_i})\big) + \text{fl}_{u32}(e^{x_i-m_i}) \Big) \]

按照上面结果，同时因为输出也要转成FP16，所以这里引入了一个舍入误差\(u_{16}\)，我们能轻易得出结果：

\[\frac{|\epsilon_n|}{d_n} \lesssim 3 \cdot n \cdot u_{32} + u_{16} \]

为什么要进行混合精度，答案已经显而易见。接下来简单谈一下真实场景的性能，其实对这种混合精度可能会有性能和精度权衡的担心，我们就拿Attention机制来看，在进行softmax之前，首先会进行gemm。一般会把这两个算子进行融合，在计算gemm时，通常会输入FP16下的矩阵，因为FP16使得数据内存减半，提高加载速度，而且大矩阵相乘一般都用高效的FP16 Tensor Core，之后在进行累加的时候会把之前的FP16数据转换成FP32进行相加，防止误差累加，因为这两个算子通过融合成为了一个融合算子，没有中间矩阵的显存读取浪费，FP16的优势就体现不出来了，虽然FP16转FP32时会产生一些损耗，但损耗很小，且在online softmax中FP32和FP16两种情况下计算速度差距不大，而且能换来很大的精度提升。再一步说如果online softmax也采用FP16,那势必会影响到gemm的精度选择，如果gemm还用FP32累加的话，反而到online softmax还要转换精度，从而损耗时间，如果gemm不用FP32累加，那造成的误差传播又进一步增大，所以在该场景下性能和精度的权衡结果显而易见了。但是在选取方案的时候，还是要多根据现实场景灵活变化，经验很有效，但是不能死脑筋。