MOSS：高效、精确的 FP8 LLM 训练，采用微尺度和自动尺度变换

1. MXFP8 的量化逻辑

MXFP8 的核心计算逻辑是：每 32 个连续数值作为一组，找到这组里绝对值最大的数，据此算出一个只表示“2 的幂次”的共享缩放因子，然后用这个因子去除组内每个数，把结果存成 FP8。

2. 信噪比是什么

信噪比（Signal-to-Noise Ratio，SNR）是衡量有用信号强度相对于噪声强度的指标。在论文里，它用来评价量化后的张量与原张量有多接近：SNR 越高，量化噪声越小，量化保真度越好。

基本定义

SNR 是信号功率与噪声功率之比：

$$\mathrm{SNR}=\frac{P_{\mathrm{signal}}}{P_{\mathrm{noise}}}$$

通常用分贝（dB）表示：

$$\mathrm{SNR}_{\mathrm{dB}}=10\log_{10}\frac{P_{\mathrm{signal}}}{P_{\mathrm{noise}}}$$

• SNR = 0 dB：信号功率和噪声功率相等。
• SNR 越高：噪声相对信号越小，质量越好。
• SNR 为负：噪声比信号还强。

论文中的具体形式

论文第 3.1 节式 (4) 把量化 SNR 定义为：

$$\mathrm{SNR}_{\mathrm{linear}} = \frac{\mathbb{E}\left[\|\vec{X}\|_2^2\right]}{\mathbb{E}\left[\|D\vec{Q}_X-\vec{X}\|_2^2\right]}$$

取对数得到 dB 值：

$$\begin{aligned} \mathrm{SNR}_{\mathrm{dB}} &= 10\log_{10}\mathrm{SNR}_{\mathrm{linear}} \\ &= 10\log_{10}\frac{\mathbb{E}\left[\|\vec{X}\|_2^2\right]}{\mathbb{E}\left[\|D\vec{Q}_X-\vec{X}\|_2^2\right]} \end{aligned}$$

论文式 (4) 把线性功率比和它的 dB 表示写在同一行上，中间用等号连接。这两个量其实不是一回事：前者是个比值，后者是对数。按定义应当分成两步写，先写比值，再取对数。若严格按“线性比值等于 dB 值”来读，式子是不成立的。

其中：

• $\vec{X}$ 是原始高精度张量；
• $D\vec{Q}_X$ 是反量化后的张量；
• $D\vec{Q}_X-\vec{X}$ 就是量化误差，也就是“噪声”；
• 分子是原始信号功率，分母是量化噪声功率。

直观理解：量化后的值偏离原值越小，分母越小，SNR 越高，说明量化越保真。

论文为什么用 SNR

FP8 只有 8 位，量化的核心问题是“量化后损失了多少信息”。SNR 把这种损失量化成一个可比较的数，于是可以公平比较不同量化方案：

• per-tensor 量化；
• per-group 量化；
• MOSS 的两级微缩放。

论文 Table 7 就是用 SNR（dB）比较这三种方案。结果显示 MOSS 的 SNR 比 per-group 高约 3.0–3.4 dB，比 per-tensor 高约 9.2–9.4 dB，说明 MOSS 量化误差更小。

一个直观类比

把原始张量想成一首歌，量化误差想成背景杂音。SNR 就是“歌声音量 / 杂音音量”。杂音越小，SNR 越高，听感越接近原唱。论文里的“原唱”是 FP32/BF16 张量，“翻唱”是 FP8 量化再反量化的结果，SNR 衡量翻唱有多接近原唱。

注意 dB 是对数尺度

SNR 提高 3 dB 大约意味着噪声功率减半，提高 10 dB 意味着噪声功率降到约十分之一。所以论文中“高 3 dB”不是微小差距，而是量化噪声显著降低。

3. 什么是范数？

范数（norm）是数学中用来衡量一个向量“大小”或“长度”的函数。你可以把它理解成“向量的尺子”：给一个向量，范数返回一个非负实数，表示这个向量有多长、多大。

最直观的范数：欧几里得范数

对向量 $\vec{x}=(x_1,x_2,\dots,x_n)$，最常见的范数是欧几里得范数（也叫 $L_2$ 范数）：

$$\|\vec{x}\|_2=\sqrt{x_1^2+x_2^2+\cdots+x_n^2}$$

这就是中学里“向量长度”的公式。例如 $\vec{x}=(3,4)$，则 $\|\vec{x}\|_2=\sqrt{9+16}=5$。

范数的一般定义

一个函数 $\|\cdot\|$ 要称为范数，必须满足三条性质：

• 非负性：$\|\vec{x}\|\ge 0$，且 $\|\vec{x}\|=0$ 当且仅当 $\vec{x}=0$。
• 齐次性：$\|a\vec{x}\|=|a|\|\vec{x}\|$，缩放向量，范数按比例缩放。
• 三角不等式：$\|\vec{x}+\vec{y}\|\le\|\vec{x}\|+\|\vec{y}\|$，两边之和大于等于第三边。

常见的几种范数

• $L_1$ 范数：$\|\vec{x}\|_1=\sum_i|x_i|$，绝对值之和。
• $L_2$ 范数：$\|\vec{x}\|_2=\sqrt{\sum_i x_i^2}$，欧几里得长度。
• $L_\infty$ 范数：$\|\vec{x}\|_\infty=\max_i|x_i|$，最大绝对值。
• $L_p$ 范数：$\|\vec{x}\|_p=(\sum_i|x_i|^p)^{1/p}$，上面几种是 $p=1,2,\infty$ 的特例。

论文里的 RMS 范数

论文式 (4) 下面写：

$$\|\vec{X}\|^2=\sqrt{\sum_{i=1}^{k}x_i^2}$$

这里需要指出：这个写法在数学上有点不严谨。按通常定义，$\|\vec{X}\|_2^2=\sum_i x_i^2$，而 $\|\vec{X}\|_2=\sqrt{\sum_i x_i^2}$。论文写的是“RMS 范数”，RMS 是均方根（Root Mean Square）：

$$\mathrm{RMS}(\vec{X})=\sqrt{\frac{1}{k}\sum_{i=1}^{k}x_i^2}$$

它等于 $L_2$ 范数除以 $\sqrt{k}$。所以论文想表达的是“用均方根来衡量张量大小”，即先平方、求平均、再开方。分子 $\mathbb{E}[\|\vec{X}\|^2]$ 是信号的平均功率，分母是量化误差的平均功率，两者之比就是 SNR。严格说式子里 $\|\cdot\|^2$ 的记法和 RMS 定义有混用，但物理含义是清楚的：衡量张量的均方幅度。

4. 为什么要进行两级微缩放

先看现有 per-group 量化的问题

论文 1 节和 3.1 节反复强调，当前 FP8 训练框架（如 COAT）对激活做 per-group 量化，每个 128 值组一个 FP32 缩放因子。这样做数值上很准，但有个致命效率问题：

缩放因子是沿着 GEMM 的内维 K 分布的。Tensor Core 做完一段乘加后，要在 CUDA Core 上把 partial sum 乘以对应的组缩放因子再累加。而 CUDA Core 的 FP32 峰值吞吐只有 FP8 Tensor Core 的 1.6%，反量化一次 partial sum 的代价约等于 60 次 Tensor Core MAC。结果 GEMM 主循环被慢速反量化拖垮。

Figure 1 显示：per-group 的 COAT 吞吐远低于 per-tensor 的 Transformer Engine。所以per-group 太贵。

那退回 per-tensor 行不行？

per-tensor 只有一个 FP32 因子，GEMM 主循环可以全在 Tensor Core 上跑，反量化放到 epilogue，效率高。但精度不够：LLM 激活有离群值，整个张量共用一个因子，小值被压得很惨，SNR 低（Table 7 显示 per-tensor 比 per-group 低约 6 dB）。

于是矛盾是：

• per-tensor：快，但精度差；
• per-group：精度好，但主循环里反量化太贵。

两级微缩放的思路：把“贵的高精度缩放”放粗粒度，把“便宜的细粒度缩放”放细粒度

MOSS 的出发点是：能不能既保留细粒度对局部动态范围的适应，又不让 GEMM 主循环里做昂贵的 FP32 反量化？

答案是把缩放拆成两层：

• Level-1 全局因子 $s$：FP32，粗粒度（整个 $k_1\approx 10K$ 范围共享），精度高但数量少。它只在 epilogue 里用一次，不在主循环里反复乘，所以贵一点没关系。
• Level-2 局部因子 $ss_i$：E8M0（2 的幂），细粒度（每 32 值一个），但格式极便宜，E8M0 乘法可以融进 Tensor Core 的 MMA 里（相当于指数偏移），不需要 CUDA Core 做 FP32 乘法。

这样，细粒度缩放由便宜的 E8M0 承担，高精度由粗粒度的 FP32 承担，GEMM 主循环里只有 Tensor Core 操作，反量化全部推迟到 epilogue（Figure 3b）。这就是“两级”的出发点。

具体怎么做：全局因子怎么来的，局部因子怎么配

第 1 步：按 32 值子块算浮点缩放 $s_i$

对第 $i$ 个 32 值子块 $\vec{X}_i$：

$$s_i = \frac{\max(|\vec{X}_i|)}{\Delta_{\max}^{\mathrm{FP8}}}$$

这是“这个子块要用多大缩放才能刚好不溢出”的精确值，是浮点数。如果直接用 $s_i$ 做量化，精度最好，但每个子块一个 FP32 因子，数量是 $k_1/32$ 个，存储和 GEMM 内乘法都贵，这又退回了 per-group 的老问题。

第 2 步：取所有 $s_i$ 的最大值作为全局因子 $s$

$$s = \max_i(|s_i|)$$

为什么取最大值？ 因为要保证整个 $k_1$ 范围内所有子块都不溢出。对任意子块：

$$\begin{aligned} \left|Q_{\vec{X}_i}\right| &= \left|\frac{\vec{X}_i}{s\cdot ss_i}\right| \\ &\le \frac{\max(|\vec{X}_i|)}{s\cdot ss_i} \end{aligned}$$

而 $s\ge s_i=\max(|\vec{X}_i|)/\Delta_{\max}^{\mathrm{FP8}}$，且 $ss_i\le 1$，所以 $|Q_{\vec{X}_i}|\le\Delta_{\max}^{\mathrm{FP8}}$，安全。

取最大值是保守的：它按“最坏的那个子块”定全局尺度，保证不溢出。代价是其他子块的有效缩放偏大、值偏小，但这些偏差由 Level-2 的 $ss_i$ 来补。

第 3 步：用全局因子归一化，算每个子块的 E8M0 局部因子

$$ss_i = \left\lceil \frac{s_i}{s} \right\rceil_{\mathrm{E8M0}} = 2^{\lceil \log_2(s_i/s)\rceil}$$

因为 $s_i\le s$，所以 $s_i/s\in(0,1]$，$ss_i$ 是 $(0,1]$ 内的 2 的幂。它的作用是：把全局因子“定大了”的那部分按子块补回来。某个子块的真实需求 $s_i$ 比全局 $s$ 小多少，$ss_i$ 就把它按 2 的幂缩回去多少。取 2 的幂是为了 E8M0 能存、能在 Tensor Core 里廉价实现。

第 4 步：有效缩放 = $s\cdot ss_i$，反量化在 epilogue

每个子块量化时用：

$$\begin{aligned} Q_{\vec{X}_i} &= \left\lceil \frac{\vec{X}_i}{s\cdot ss_i} \right\rceil_{\mathrm{FP8}} \\ DQ_{\vec{X}_i} &= Q_{\vec{X}_i}\cdot s\cdot ss_i \end{aligned}$$

GEMM 主循环里，$Q_x$ 先乘 $ss_x$（E8M0，便宜），再和 $Q_w$ 做 Tensor Core MMA；得到的 FP32 partial sum 在 epilogue 里乘 $s_W$ 和 $s_x$（FP32，只做一次）。这样：

• 细粒度：每 32 值一个 $ss_i$，保留了对局部动态范围的适应；
• 低开销：细粒度因子是 E8M0，贵重的 FP32 乘法只在 epilogue 做一次；
• 不溢出：全局因子取最大值保证安全。

为什么这样比 per-group 好、比 per-tensor 好

• 比 per-tensor 好：per-tensor 只有一个因子，无法适应 32 值块间的局部差异；MOSS 的 $ss_i$ 提供了 32 值粒度的局部适应。论文式 (7) 证明 $\mathrm{SNR}_{\mathrm{MOSS}}>\mathrm{SNR}_{\mathrm{per-group}}$，Table 7 实测 MOSS 比 per-group 高 3.0–3.4 dB，比 per-tensor 高 9.2–9.4 dB。
• 比 per-group 好：per-group 的 FP32 因子沿 K 维分布，要在主循环里做 CUDA Core 反量化；MOSS 把 FP32 乘法放到 epilogue，主循环只有 Tensor Core 的 E8M0 微缩放，效率接近 per-tensor（Figure 3b、Table 6）。
• 直觉：per-group 是“每个小组都用一把精确但昂贵的尺子，每量一次都要停下来换尺”；per-tensor 是“全程一把粗尺子，快但不准”；MOSS 是“一把精确的主尺（FP32 全局）定基准，每个小组配一个便宜的 2 的幂微调刻度（E8M0），微调在测量过程中自动完成，不用停”。

一句话总结：全局因子取所有子块 $s_i$ 的最大值，是为了用 FP32 一次性保证不溢出；局部 E8M0 因子把全局因子“定大”的部分按 2 的幂补回来，实现细粒度精度；两者相乘，既保住 per-group 的精度，又拿到接近 per-tensor 的效率。

5. 自动缩放

自动缩放（automatic scaling）是 MOSS 为权重设计的一种缩放因子确定方法。它的核心原理是：利用 Adam/AdamW 优化器的“更新有界性”，也就是权重每一步的变化量被学习率 $\eta$ 限制住，从而提前预测权重最大绝对值的增长上界，直接算出缩放因子，而不必每步去 HBM 里读权重、做 max-reduction。

下面按“问题—原理—公式—做法—效果”拆开讲。

要解决的问题：just-in-time scaling 太贵

FP8 量化需要缩放因子 $s=\max(|W|)/\Delta_{\max}^{\mathrm{FP8}}$，保证权重落在 FP8 动态范围内。传统做法是 just-in-time（JIT）缩放：每一步都读整个 FP32 权重张量，算最大值，再写回 FP8。论文 Table 1 显示，对大张量这一步要 0.54 ms；而且涉及多次 HBM 读写，抵消了 FP8 的性能收益。

自动缩放的目标：不读权重、不做 max-reduction，也能给出安全且够准的缩放因子。

原理：Adam 类优化器的更新有界性

论文 2.2 节和 3.2 节的核心观察是：Adam/AdamW 的更新量

$$\Delta_t = \eta \cdot \frac{\hat{\mathbf{m}}_t}{\sqrt{\hat{\mathbf{v}}_t}+\epsilon}$$

有一个关键性质：更新幅度与梯度的整体缩放无关。把梯度 $g_t$ 乘一个因子 $c$，一阶矩 $m_t$ 乘 $c$，二阶矩 $v_t$ 乘 $c^2$，代入后 $c$ 抵消。所以即使 FP8 梯度动态范围很大，Adam 的更新仍然稳定。

论文 Theorem 2 进一步给出上界：

$$|\Delta_t| \le \begin{cases} \eta \cdot \dfrac{1-\beta_1^t}{\sqrt{1-\beta_2^t}}, & \text{若 } 1-\beta_1^t > \sqrt{1-\beta_2^t}\\[6pt] \eta, & \text{否则} \end{cases}$$

证明思路：由 Jensen 不等式（凸函数 $f(x)=x^2$）有 $(m_t)^2\le v_t$，所以 $|m_t/\sqrt{v_t}|\le 1$。于是 $|\Delta_t|$ 主要由 $\eta$ 和偏差校正项决定。实际 LLM 训练中通常 $1-\beta_1^t<\sqrt{1-\beta_2^t}$（例如 $\beta_1=0.9,\beta_2=0.95$），所以每一步的有效更新约被学习率 $\eta$ 界住：$|\Delta_t|\lesssim \eta$。

直觉：不管梯度多大多小，Adam 每步最多把权重挪动约 $\eta$ 那么多。这相当于在参数周围隐式定义了一个“信任域”。

公式：预测权重最大值的上界，直接算缩放因子

既然每步最多挪 $\eta$，那么从初始权重 $W_0$ 出发，第 $t$ 步时权重的最大绝对值满足：

$$\max(|W_t|) \le \max(|W_0|) + \eta \cdot t$$

于是第 $t$ 步的缩放因子可以直接预测：

$$s_t = \frac{\max(|W_0|) + \eta \cdot t}{\Delta_{\max}^{\mathrm{FP8}}} = s_0 + \frac{\eta \cdot t}{\Delta_{\max}^{\mathrm{FP8}}}$$

其中 $s_0=\max(|W_0|)/\Delta_{\max}^{\mathrm{FP8}}$ 只在初始化时做一次 max-reduction 得到。之后每一步只需用学习率 $\eta$ 和步数 $t$ 套公式，常数时间、零 HBM 读取。这就是论文式 (10)。

为什么取上界而不是精确值？ 取上界保证权重不会溢出 FP8 范围；代价是缩放因子偏大、权重值偏小，但 MOSS 通过定期重缩放来修正。

具体做法：区间内预测，区间末重缩放

完全不用 max-reduction 会累积误差，所以 MOSS 采用固定区间重缩放：

1. 初始化：做一次 max-reduction，得到 $s_0$。
2. 区间内：每步用 $s_t=s_0+\eta t/\Delta_{\max}^{\mathrm{FP8}}$ 预测缩放因子，不做任何 max-reduction。
3. 区间末（默认每 500 步）：重新读一次当前权重，做 max-reduction，更新 $s_0$ 和基准，进入下一个区间。

论文 Figure 4 显示，区间长度 500 时，自动缩放的轨迹始终略高于 JIT 缩放，说明它偏保守、保证不溢出，且两条曲线很接近，精度损失很小。

为什么能省开销：Table 1 和 Table 10

• Table 1：不同大小的权重张量，JIT 缩放要 0.08–0.54 ms，自动缩放恒定 0.02 ms，几乎可忽略。
• Table 10：LLaMA-2-7B 微调端到端，JIT 每步缩放开销 3.8 ms、总步时 68.5 ms、吞吐 38,642 token/s；自动缩放开销 0.2 ms、总步时 63.1 ms、吞吐 41,998 token/s，端到端提速 8.7%。
• Table 11：自动缩放与 JIT 的下游精度差异在 $\pm0.2\%$ 内，属于统计噪声（$\pm0.3\%$），说明精度没有退化。

区间长度的权衡：Table 9

方法	区间	缩放开销 (ms/step)	有效吞吐	NumGLUE 精度
BF16	—	0	1.0×	58.7%
JIT	1	3.8	1.0×	59.3%
MOSS	100	0.03	1.043×	59.4%
MOSS	500（默认）	0.02	1.044×	59.4%
MOSS	2000	0.01	1.044×	58.1%

区间太短（JIT 每步）开销大；区间太长（2000）跟踪不上尺度漂移，精度略降；100–500 是最佳区间，开销可忽略、精度与 BF16 持平甚至略高。

和 JIT、delayed scaling 的区别

• JIT scaling：每步读当前权重算 max，最准但最贵（多次 HBM 读写）。
• Delayed scaling：用历史窗口的最大值估计当前缩放，假设统计一致，遇到离群值可能失稳。
• 自动缩放（MOSS）：利用 Adam 更新有界性预测缩放因子上界，区间内零 max-reduction，区间末重校准。它既不是“看历史”，也不是“看当前”，而是“按优化器性质推算上界”。

一句话总结

自动缩放的原理是：Adam/AdamW 每步对权重的更新幅度被学习率 $\eta$ 界住（$|\Delta_t|\lesssim\eta$），所以权重最大绝对值随时间至多线性增长 $\max(|W_t|)\le\max(|W_0|)+\eta t$；据此直接预测缩放因子 $s_t=s_0+\eta t/\Delta_{\max}^{\mathrm{FP8}}$，只在初始化做一次 max-reduction，之后区间内零开销、区间末重校准，从而在保持精度的同时消除 JIT 缩放的内存访问和 max-reduction 成本。