GIFT 阅读笔记:量化用的坐标系,比量化器本身更重要

笔记日期： 2026-07-16 笔记作者： Zhongzhu Zhou 论文标题： GIFT: Geometry-Informed Low-precision Gradient Communication for LLM Pretraining 作者： Jieying Wang, Shuyuan Fan, Mingkai Zheng, Zhao Zhang（罗格斯大学 Rutgers University） arXiv： 2607.07494 状态： arXiv 预印本,2026 年 7 月

一句话总结

训练大语言模型的时候,如果用多张 GPU 并行训练,每一步都要把各个 GPU 上算出来的梯度同步(平均)一次——这个同步操作本身就很贵。一个常见的省钱办法是把梯度用低精度格式(比如 FP8)来通信,而不是用标准的 FP32,这样传输的数据量直接减少四分之三。做法很直接:拿到一个层的梯度张量,选一个缩放系数让它的数值落进 FP8 那个很窄的可表示范围里,然后量化。这篇论文的核心发现是:这个”直接”的做法会悄悄丢掉精度,而且原因跟 FP8 这个格式本身没关系——梯度往往是各向异性的,也就是说,它在某些方向上”拉得很长”,在另一些方向上”缩得很短”,像一个雪茄形状的点云,而不是一个球。用一个标量缩放系数,没办法同时兼顾”长的方向”和”短的方向”:长方向缩放得刚刚好,短方向就缩放得太小(相当于精度浪费在了没什么变化的方向上);反过来如果按短方向去校准,长方向又会溢出截断。结果就是量化误差在不同方向上分布得很不均匀,而恰好对优化最重要的那些方向,反而可能是被扭曲得最厉害的。GIFT 的解法既不是发明新的量化格式,也不是改优化器,而是换一个坐标系:量化之前,先把梯度变换到一个看起来近似球形(各向同性)的坐标系里,在那个坐标系里量化、通信,通信完了之后再变换回来,交给优化器。这个变换用的数学工具,跟 K-FAC 一类自然梯度优化方法里用来做曲率近似的东西(Fisher 信息矩阵 / Kronecker 分解曲率近似)是同一套,但关键区别是:GIFT 没有把优化器变成自然梯度优化器——它只是借用这套几何结构来构造更好的通信坐标系,变换完之后梯度会原样变回它本来的坐标系,再交给优化器做参数更新,优化器本身完全没被碰过。为了让这个想法在 LLM 预训练这种规模下真正可用,论文做了三次层层递进的简化——只保留变换的”输入侧”、用一个秩为 32 的低秩近似代替完整矩阵、只对被识别为最容易被量化损伤的一小部分层应用——最终得到的系统只多花 3%-9% 的额外内存,却保住了完整几何变换的大部分精度收益,比对每一层都上完整变换要便宜得多。在 Llama-300M 和 Llama-600M 的预训练实验里,GIFT 相比直接的欧氏 FP8 通信,下游任务质量保持得更好(14 个下游任务里跟 FP32 比赢了 7 个,而欧氏基线只赢 4-5 个),同时依然保留了 FP8 通信时间加速的相当一部分,在 64 张 GH200 超级芯片上给 600M 模型带来了实测 7.6% 的端到端预训练时间缩减。

关键要点

• 论文把 FP8 梯度通信误差重新定义为部分是坐标系问题,不完全是数值格式问题:在原始的欧氏参数坐标系里量化一个各向异性的梯度张量,必然会让不同方向被扭曲的程度不一样,因为一个per-tensor(甚至 per-block)的标量缩放系数,没办法把方差本来就不一样的各个方向的动态范围拉到同一水平。
• GIFT 的变换建立在 K-FAC(Kronecker 分解曲率近似)逐层 Fisher 近似之上,$F_W \approx A \otimes G$,其中 $A = \mathbb{E}[aa^\top]$ 是输入激活的二阶矩,$G = \mathbb{E}[\delta\delta^\top]$ 是输出梯度的二阶矩——这跟 K-FAC 在做优化器预条件时用的是同一套结构,但这里被重新用作临时的通信坐标系,完全不碰优化器的更新规则本身。
• 完整的双侧白化变换是 $\widetilde W_g = L_G^{-1} W_g L_A^{-\top}$,用到 Cholesky 分解 $A = L_A L_A^\top$、$G = L_G L_G^\top$;这个变换把 K-FAC 度量下那个各向异性的”椭球”映射成一个圆的欧氏球,这样一个 FP8 缩放系数就能公平地对待每个方向了。
• 论文的一个消融实验(Table I)显示:输出侧变换几乎没有任何贡献,跟不做变换的欧氏基线几乎一样;而只用输入侧变换就几乎拿到了完整双侧 K-FAC 变换的全部保真度收益——这是一个真正有用的经验性简化,直接把变换的开销砍掉了一半。
• 对输入侧因子 $A$ 做秩为 32 的低秩近似,精度几乎跟用完整矩阵一样好(Table II),而一个朴素的对角近似几乎不比不做变换的基线强多少——这说明收益真正来自 $A$ 的非对角(跨维度)相关结构,而不只是给每个坐标单独乘一个缩放系数。
• GIFT 只对一个通过离线画像筛选出来的、数量很小的”数值脆弱层”子集应用几何感知变换(论文 600M 设置里是 13 个 MLP fc2 层,按 FP8 边界命中率排序选出来的),其余层照旧走普通欧氏快速路径——这种选择性部署设计把总开销控制得很小(300M 模型多花 3.33% 内存,600M 模型多花 8.98%),而不是像给每一层都上完整 K-FAC 状态那样开销随层数线性增长。
• 在下游任务评测上(Table III,14 个任务),GIFT 在 300M 和 600M 两个规模上都是 14 个任务里赢 FP32 的 7 个,明显强于欧氏基线的 4/14 和 5/14——甚至在 300M 规模上还赢过更重的完整 K-FAC变体(6/14),这直接支撑了论文”选择性部署”这个设计选择,而不是”处处都上完整变换”。
• 尽管下游任务上有明显差距,GIFT 和欧氏基线的验证损失曲线却几乎完全重合(Figure 6)——这是一个挺有意思、也有点反直觉的发现,论文用它来论证:验证损失是通信保真度效应的一个不完整代理指标,下游评测在这个场景下才是更能说明问题的信号。
• 实测的系统层面收益是:600M 的 Llama 在 64 张 NVIDIA GH200 超级芯片上,GIFT 带来了 7.6% 的端到端预训练时间缩减,而同规模下纯欧氏 FP8 基线本身能拿到 10.79% 的缩减——GIFT 明确地牺牲了一部分原本纯 FP8 能拿到的速度,来换取更好的下游保真度,这个取舍论文说得很坦诚,没有藏着掖着。
• GIFT 相对 FP32 减少了 75.0% 的梯度通信量(这跟任何 FP8 方案是一样的,毕竟省下的字节数是格式本身决定的),论文明确区分了这个”载荷体积缩减”和更小的”端到端时间缩减”——后者还包含了前向/反向计算、优化器更新、量化/反量化、几何变换本身、以及同步开销。
• 论文明确把这个坐标系视角的核心想法框定为对未来 FP4 通信是前向兼容的——一旦软硬件支持成熟,因为几何变换的开销大体上跟上层用的是哪种低精度格式是相互独立的。

前置知识:需要先搞懂的东西

这篇论文横跨三个大多数人分开接触的领域:(1)为什么分布式 LLM 预训练里梯度通信本身就是个瓶颈,(2)低精度(FP8)通信是怎么工作的、误差从哪来,(3)K-FAC / 自然梯度这条优化研究脉络,它提供了 GIFT 借用的数学工具。这一节把这三块都搭起来,再进入论文自己的贡献。

为什么梯度通信是个瓶颈

现代 LLM 是在很多张 GPU 上同时训练的,通常会组合使用数据并行(DP)、张量并行(TP)、流水线并行(PP)——合起来叫”3D 并行”。这几种并行方式各自引入了不同的通信模式:

• 数据并行:每个 worker 都保存一份完整的模型副本,处理训练 batch 里不同的一片数据。反向传播算完之后,每个 worker $r$ 都有了自己的本地梯度 $g^{(r)}$,这些梯度必须在优化器更新之前,在所有 $R$ 个 worker 之间做平均,这样每个副本才会应用同一个更新,保持同步:

$$\bar g = \frac{1}{R}\sum_{r=1}^{R} g^{(r)}. \tag{1}$$

这个求平均的操作是用一个叫 all-reduce 的集合通信原语来实现的(或者数学上等价的 reduce-scatter 加 all-gather,常见于 ZeRO 这类内存分片的场景)。这个集合通信的开销,会同时随着模型规模(要平均的参数变多了)和 worker 数量(参与规约的节点变多了)增长,这正是为什么它在大规模场景下会变成主要开销:论文引用的前人工作报告说,一个 8.3B 参数的 GPT 模型在 128 张 A100 上跑,通信要吃掉总预训练时间的 40%。

• 张量并行:单个层被拆到不同的设备上(比如权重矩阵的不同列在不同 GPU 上),所以一个层内部的前向/反向计算过程中就得交换部分结果。
• 流水线并行:模型被切成一串顺序的阶段,分布在不同设备上,所以相邻的阶段要交换边界处的激活(前向)和边界处的梯度(反向)。

GIFT 专门聚焦在数据并行的梯度 all-reduce 这条路径上——也就是在数据并行副本之间做梯度平均的那部分通信——因为这条路径跟低精度压缩最直接兼容,论文借此把”用什么坐标系通信”这个变量单独隔离出来研究,同时保持优化器、模型架构、以及其余的并行策略都不变。

低精度通信为什么有用,又在哪里出问题

缩小通信瓶颈最直接的办法就是缩小每个梯度值传输时占用的比特数。如果用 FP8(8 位浮点)而不是 FP32(32 位浮点)来通信梯度,传输的字节数直接变成四分之一——这是一个跟模型、集合通信算法都无关的、几乎是”白捡”的带宽缩减。这也是为什么 FP8 梯度通信(以及更窄的 FP4)成了一个活跃的系统研究方向:比如 FP8-LM 和 COAT 对模型权重、激活、梯度、优化器状态都做缩放加量化;SDP4Bit 在分片数据并行场景下用基于傅里叶变换的 4-bit 通信。

但是把数值格式变窄,在另一个意义上并不是免费的:FP8 的动态范围比 FP32 小得多,精度也粗糙得多,所以把一个数值圆整到 FP8 那只有 256 个可表示级别里的某一个,必然会丢掉一些信息。现有的低精度方案基本都用同一套办法处理这个问题:选一个缩放系数 $s$(通常是张量里、或者张量某个块里的最大绝对值),把张量除以 $s$ 让数值落进 FP8 的可表示范围,然后再量化。这个办法只有在张量数值在各个维度上大致均匀的时候才管用——如果梯度某个维度上的数值在 $10^{-4}$ 量级,另一个维度在 $10^{-1}$ 量级,一个按较大维度校准的缩放系数,会把已经很小的那个维度的数值进一步压到 FP8 格式最粗糙、精度最低的区域,对那个方向的扭曲会远大于其他方向。论文的核心诊断是:真实的 LLM 梯度恰恰就是这种不均匀、各向异性的对象,而这种各向异性——而不是 FP8 格式本身——是一个之前没怎么被仔细审视过的通信误差主要来源。

FP8 到底是什么,以及它的”窄”为什么会放大各向异性

值得具体说说”FP8”在数值上到底意味着什么,因为论文的整个论证都建立在各向异性如何跟这个格式的窄度相互作用上。一个浮点数存储为(符号位、指数位、尾数位)。FP32 分配 8 个指数位、23 个尾数位,既有很大的动态范围,范围内的精度也很细。LLM 训练里常用的 FP8 有两种变体:E4M3(4 个指数位、3 个尾数位——偏向精度而非范围,最大值约 $\pm448$)和 E5M2(5 个指数位、2 个尾数位——偏向范围而非精度,最大值约 $\pm57344$,但每个数量级之间只有 4 个可表示的尾数级别)。不管哪种,FP8 每个倍频程只有 $2^3=8$ 或 $2^2=4$ 个尾数级别,相比 FP32 的 $2^{23}$——单个数量级内的分辨率降低了大约 5-6 个数量级。标准做法(FP8-LM、COAT、以及本文的欧氏基线都是这么做)是per-tensor 或 per-block 缩放:算出 $s = \mathrm{amax}(X)/\mathrm{FP8\_MAX}$,把张量除以 $s$ 让最大值落在格式最大可表示值附近,量化,通信完之后再乘回 $s$。这个缩放步骤正确地处理了”量级”问题(保证整个张量不会溢出或下溢 FP8 的范围),但对”形状”问题完全没有帮助:张量里每个元素,不管它对应 K-FAC 椭球意义下的哪个”方向”,都被用同样的相对精度量化,因为缩放系数 $s$ 和格式固定的尾数宽度是均匀作用在所有张量元素上的。这正是 GIFT 的坐标变换要填补的空白:单纯缩放只能拉平”整体量级”,只有换一个基才能拉平”跨方向的相对重要性”。

K-FAC 和 Fisher 信息矩阵:借来的工具

要修复一个各向异性问题,你得先有办法度量这种各向异性——某种”这个方向和那个方向应该区别对待多少”的概念。这正是优化理论里 Fisher 信息矩阵 $F(\theta)$ 提供的东西:它是一个局部度量,说明一个小的参数扰动 $\Delta\theta$ 到底会让模型的输出分布变化多少:

$$\|\Delta\theta\|_F^2 = \Delta\theta^\top F(\theta)\, \Delta\theta. \tag{2}$$

直觉上讲,如果 $F(\theta)$ 在某个方向上很”拉伸”(特征值大),在另一个方向上很”压缩”(特征值小),那说明模型的行为对扰动拉伸方向极其敏感,对扰动压缩方向几乎不敏感——同样绝对大小的扰动,取决于你往哪个方向移动,意义完全不一样。这正是论文需要的”各向异性”概念,只不过这里是应用在参数上,而不是直接应用在梯度上(不过这两者其实密切相关,后面会看到)。

问题是,现代 LLM 的 $F(\theta)$ 是一个(参数数量)² 项的矩阵——直接构造、求逆、甚至存储都不可能。**K-FAC(Kronecker 分解曲率近似)**是一套打磨了很多年的近似方案,通过两个简化让这件事变得可行:

1. 按层块对角化:不用一个巨大的矩阵把网络里每个参数跟其他每个参数耦合起来,K-FAC 把 Fisher 近似为块对角的,每层一个块: $$F(\theta) \approx \text{blockdiag}(F_{W_1}, F_{W_2}, \ldots, F_{W_L}). \tag{3}$$ 这丢掉了跨层的曲率信息(扰动第 3 层的权重和扰动第 7 层的权重之间怎么相互作用),但保留了层内曲率——实践中对前馈式的层来说,这通常是占绝大部分的项。
2. 层内的 Kronecker 分解:对一个权重矩阵为 $W \in \mathbb{R}^{d_\text{out}\times d_\text{in}}$ 的单个线性层,设 $a$ 是这个层的输入激活,$\delta$ 是损失对这个层预激活输出的梯度。K-FAC 把这一层的 Fisher 块近似成两个小得多的矩阵的 Kronecker 积: $$F_W \approx A \otimes G, \qquad A = \mathbb{E}[aa^\top],\quad G = \mathbb{E}[\delta\delta^\top]. \tag{4}$$ 这里 $A \in \mathbb{R}^{d_\text{in}\times d_\text{in}}$ 捕捉输入侧的二阶统计(输入到这层的激活,在 batch 维度上怎么协同变化),$G \in \mathbb{R}^{d_\text{out}\times d_\text{out}}$ 捕捉输出梯度侧的二阶统计(从这层反向传播出去的梯度怎么协同变化)。K-FAC 不需要一个 $(d_\text{in}\cdot d_\text{out})^2$ 项的巨大矩阵,只需要估计和存储两个小得多的矩阵,大小分别是 $d_\text{in}^2$ 和 $d_\text{out}^2$。

为什么这跟”通信”有关系,而不只是跟”优化”有关系?因为 K-FAC 最初是被设计成一个预条件器——你会用 $F_W^{-1}$ 在优化器更新之前重新缩放梯度,把梯度下降变成一个近似的自然梯度方法。GIFT 明确没有这么做:“这篇论文没有把 K-FAC 当作优化器使用,也没有应用自然梯度更新……只是把它当作一个可行的局部度量,用来定义通信坐标。” 论文重新使用了同一个数学对象(块对角的 Kronecker 分解曲率近似),但用在完全不同的地方:作为通信步骤前后一个临时的换基,变换之后梯度会原样映射回它原来的坐标系,在此之前,不变的优化器根本不会碰到它。这个”同一套数学,不同的用途”值得记住,因为这正是 GIFT 能声称”不改变优化器、模型、训练配方、通信集合操作或低精度格式”,却依然能有意义地改变通信保真度的原因。

核心想法:在白化坐标系里量化,而不是在原始坐标系里

前置知识都铺好了,下面一步步推导论文的核心机制。

第一步:K-FAC 的”椭球”,以及它为什么是各向异性的

把 Kronecker 近似(式 4)代入局部度量的定义(式 2),对一个权重扰动 $\Delta W$(用标准的按列堆叠的向量化约定 $\mathrm{vec}(\cdot)$,把矩阵展平成一个长列向量):

$$\|\Delta W\|_{F_W}^2 = \mathrm{vec}(\Delta W)^\top (A \otimes G)\, \mathrm{vec}(\Delta W). \tag{5}$$

用 Kronecker 积/迹的恒等式 $\mathrm{vec}(X)^\top (A \otimes G)\, \mathrm{vec}(X) = \mathrm{tr}(A X^\top G X)$,可以化简成一个好用得多的矩阵迹表达式,不用真的去构造那个巨大的 Kronecker 积:

$$\|\Delta W\|_{F_W}^2 = \mathrm{tr}\!\left(A\, \Delta W^\top G\, \Delta W\right). \tag{6}$$

满足 $\{\Delta W : \|\Delta W\|_{F_W}^2 \le \epsilon\}$ 的扰动集合,在参数空间里画出的是一个椭球——不是球,因为 $A$ 和 $G$ 在不同方向上一般有不同的特征值。具体来说:如果 $A$ 最大的特征值是最小特征值的 100 倍,那么 K-FAC 度量认为沿着 $A$“刚性”特征方向的扰动,比同样大小、沿着”柔软”特征方向的扰动重要 100 倍。这个椭球正是论文 Figure 1 那张卡通式二维散点图(真实梯度样本在欧氏坐标下看起来是”雪茄形”)在图示的东西:“扰动到底有多重要”这个真实几何是拉伸的、不是圆的,而一个统一的(per-tensor 标量)量化缩放系数,恰恰是应对拉伸几何的错误工具。

这张图正是式(9)推导的直接可视化对应物:左图是原始 $\Delta W$ 空间,K-FAC 椭球 $\{\Delta W : \|\Delta W\|_{F_W}^2 \le \epsilon\}$ 真的是拉伸的(有些方向即便欧氏大小相同,K-FAC 加权的”重要性”也远大于其他方向);右图是应用 $U = L_G^\top \Delta W\, L_A$ 之后的 $U$ 空间,同一个椭球(通过式 9 的代数)变成了一个圆球。一个统一的、坐标轴对齐的 FP8 量化网格,跟左图那种几何形状很不匹配——它必须同时兼顾长轴(有截断的风险)和短轴(把大部分可表示级别浪费在一个几乎不怎么变化的方向上)——但跟右图那种几何形状是匹配的,因为每个轴的散布程度都差不多。

第二步:白化——把椭球变成球

处理椭球型度量的标准招数是白化:找一个线性变换,把椭球映射成一个圆球,这样在新坐标系里,一个统一的缩放系数又变得合适了。为此,对 $A$ 和 $G$ 做 Cholesky 分解(任何矩阵平方根都行,Cholesky 是标准的实用选择,因为它数值上稳定,而且只需要下三角那一半):

$$A = L_A L_A^\top, \qquad G = L_G L_G^\top. \tag{7}$$

代入式(6):

$$\|\Delta W\|_{F_W}^2 = \mathrm{tr}\!\left(L_A L_A^\top\, \Delta W^\top L_G L_G^\top\, \Delta W\right). \tag{8}$$

利用迹的循环性质($\mathrm{tr}(XYZ) = \mathrm{tr}(ZXY)$)把两个 $L_A^\top$ 和两个 $L_G^\top$ 重新组合到一起:

$$\|\Delta W\|_{F_W}^2 = \mathrm{tr}\!\left(L_A^\top\, \Delta W^\top L_G L_G^\top \Delta W\, L_A\right) = \left\| L_G^\top \Delta W\, L_A \right\|_F^2. \tag{9}$$

最后一步用到了恒等式 $\mathrm{tr}(X^\top X) = \|X\|_F^2$(Frobenius 范数的平方就是 $X^\top X$ 的迹),代入 $X = L_G^\top \Delta W\, L_A$。这一步是整个推导的关键:它说明,原始坐标系里那个复杂的各向异性椭球度量,跟先把 $\Delta W$ 用 $U = L_G^\top \Delta W\, L_A$ 变换一下之后的、朴素的各向同性欧氏(Frobenius)范数完全相等。换句话说,定义

$$U = L_G^\top\, \Delta W\, L_A, \tag{10}$$

那么原来的椭球约束 $\{\Delta W : \|\Delta W\|_{F_W}^2 \le \epsilon\}$ 就精确地映射成了圆的欧氏球 $\{U : \|U\|_F^2 \le \epsilon\}$。这就是为什么这个变换叫”白化”:它把一个相关的、跟方向有关的(各向异性)对象,映射成了一个不相关的、跟方向无关的(各向同性)对象——跟经典信号处理里对有色噪声信号做白化,是同一个核心想法。

第三步:把白化用到通信上,而不是优化上

上面这个洞察是关于”对权重的扰动”的,但 GIFT 需要把它用到”正在被通信的梯度”上。论文把梯度 $W_g$ 的完整几何感知通信坐标定义为

$$\widetilde{W}_g = L_G^{-1}\, W_g\, L_A^{-\top}. \tag{11}$$

为什么这里用的是逆因子,而不是式(10)里同样的因子?因为式(10)是在给一个”要拿去跟度量比较的扰动”做白化——你要乘上几何因子,把一个欧氏球映射成椭球形的扰动。这里我们做的是反过来的事:我们已经有一个梯度,它”生活”在一个坐标系里,而这个坐标系里真正有意义的度量是各向异性的 K-FAC 椭球,我们想把它重新表达到一个坐标系里,让朴素的欧氏度量匹配的正是原来的 K-FAC 度量——也就是说,我们想在量化之前撤销这种各向异性。这就需要逆因子 $L_G^{-1}$ 和 $L_A^{-\top}$,作用在对应的两侧。等价地,用向量化的形式写:

$$\mathrm{vec}(\widetilde W_g) = \left(L_A^{-1} \otimes L_G^{-1}\right) \mathrm{vec}(W_g), \tag{12}$$

这个变换在量化之前,把梯度里由 $A$ 捕捉的主要输入侧缩放(通过 $L_A^{-1}$)和由 $G$ 捕捉的输出侧缩放(通过 $L_G^{-1}$)都去掉了。在这个变换后的坐标系里,梯度分布更接近各向同性,所以一个坐标轴对齐的 FP8 量化网格——本质上对每个坐标都一视同仁——跟这个梯度打交道时,会比在原始的各向异性欧氏坐标里均匀得多。

白化之后,GIFT 应用跟任何欧氏空间方案一样的低精度通信流程——缩放、量化到 FP8、all-reduce、反量化、求平均:

$$\widehat{\widetilde W}_g = \mathrm{CommFP8}\!\left(\widetilde W_g\right), \tag{13}$$

然后,在优化器碰到它之前,把同步、反量化之后的结果映射回原始的欧氏梯度坐标系:

$$\widehat{W}_g = L_G\, \widehat{\widetilde W}_g\, L_A^\top. \tag{14}$$

这最后一步”映射回去”正是让 GIFT 严格地只是一个通信时刻的技巧、而不是优化器改动的原因:$\widehat{W}_g$ 精确落回了优化器本来期望的空间,所以同一个 FP32/BF16 的 Muon(或者 Adam,或者任何其他)优化器更新可以原样应用。模型架构、分布式并行策略、优化器更新规则一个都没被碰过——变的只有通信时刻临时用的坐标系。

为什么”不是 K-FAC 优化器”是一个真正的设计选择,而不只是一句免责声明

值得停下来想想,为什么论文这么明确地强调 GIFT 不是一个 K-FAC 式的自然梯度方法,因为这是一个有真实取舍的设计决定,不只是一句免责声明。一个真正的 K-FAC 优化器会在每一步都把预条件器 $F^{-1}$ 应用到梯度上,把变换后的梯度喂给优化器的更新规则——这会实实在在地改变优化轨迹本身,好坏都有可能(K-FAC 预条件的训练每一步可能收敛更快,但每一步也更贵,而且块对角 Kronecker 结构近似的质量,直接决定了最终轨迹是不是真的更好)。GIFT 完全绕开了这个问题:因为变换在优化器看到任何东西之前就已经被撤销了(式 14),优化动力学相对欧氏基线是可证明不变的——唯一会变的是,平均之后的梯度经过 FP8 往返之后能有多忠实。这是一个真正干净的关注点分离:意味着实验里观察到的 GIFT 和欧氏基线之间的任何精度差异,都可以明确归因于通信保真度,而不是纠缠在某种不同优化轨迹的效应里。这个设计可能失效的边界情况是:如果 K-FAC 因子太陈旧,或者估计得太差,以至于”白化”变换实际上在某个方向上增大了方差而不是减小它——我们在批判性分析那一节会回到这一点,因为论文的因子刷新周期(每 50 步)是一个真实的、没有被充分检验的风险点。

从理论到实践:三次简化

完整的双侧变换(式 11)在概念上很干净,但对每个符合条件的层、每一步都应用,对真正的 LLM 预训练来说太贵了:给每一层的 $A \in \mathbb{R}^{d_\text{in}\times d_\text{in}}$ 和 $G \in \mathbb{R}^{d_\text{out}\times d_\text{out}}$ 都做计算和 Cholesky 分解,每个刷新周期都要重来一遍,增加了真实的计算和内存开销。论文第五节系统地通过三次简化把这个想法一路精简下来,每一步都先用一个针对性的消融实验来验证,再决定采纳——这是一种方法论上很讲究的模式,值得逐步讲清楚,因为这是”我们猜这样会更便宜”和”我们测量过这个具体简化几乎不损失精度”之间的区别。

这张路线图在读接下来三个小节的时候值得一直放在心里:每个方框代表一次刻意的范围缩减,每条箭头代表论文在正式采纳这次简化之前,先跑的一个针对性实验(分别对应 Table I、Table II 和 Figure 3 的画像研究),而不是先简化再指望精度撑得住。最终状态——路线图底部——就是下面算法 1 要形式化的那个选择性、只用输入侧、秩 32 的设计。

简化一:只保留输入侧

问题:双侧变换里的输入侧因子 $L_A$ 和输出侧因子 $L_G$,到底是哪一个真正带来了实际的保真度收益?

实验:一个受控的单步 FP8 往返测试。从一个 FP32 梯度 $g$ 出发,把它映射到候选的通信坐标系里,在那个坐标系里做一次 FP32→FP8→FP32 的往返,把结果映射回欧氏梯度空间,然后用四个指标跟原始的 $g$ 比较重建出来的梯度 $\hat g$:

$$\mathrm{RelL2} = \frac{\|\hat g - g\|_2}{\|g\|_2}, \qquad \mathrm{Cos} = \frac{\langle \hat g, g\rangle}{\|\hat g\|_2 \|g\|_2}, \qquad \mathrm{MaxErr} = \|\hat g - g\|_\infty, \qquad \mathrm{MSE} = \frac{\|\hat g - g\|_2^2}{n}. \tag{15}$$

RelL2、MaxErr、MSE 都是”越小越好”的误差量级指标,只是敏感的地方不一样(RelL2 是归一化的平均情形度量,MaxErr 抓住单个误差最大的元素,MSE 对大误差做平方放大);Cos 度量的是 FP8 往返有没有保住梯度的方向——这可能是对优化来说最重要的性质,毕竟梯度下降主要关心的是往哪个方向走。

结果(论文 Table I):只用输出侧的变换,跟不做变换的欧氏基线在四个指标上几乎没有区别(RelL2 $5.283\times10^{-2}$ 对基线的 $5.282\times10^{-2}$)——几乎没有任何贡献。相反,只用输入侧的变换几乎在每个指标上都跟完整双侧 K-FAC 变换一样好(RelL2 输入侧是 $1.768\times10^{-2}$,完整 K-FAC 是 $1.771\times10^{-2}$——差别在第四位有效数字上)。这是一个真正有用的经验发现,而且并不是先验上显然的:它说明,对这个特定的、用 Muon 优化的 Llama 设置来说,基本上所有对 FP8 保真度有影响的各向异性都来自每一层的输入激活侧,而不是输出梯度侧。设计上的结论是:完全去掉输出侧变换,只保留

$$\widetilde W_g = W_g\, T_A, \tag{16}$$

其中 $T_A$ 是从 $L_A^{-1}$ 导出的输入侧变换(对应式 11 里 $L_A^{-\top}$ 那一项的简化)。仅这一步就把每一层需要计算和存储的几何机制大致减半了,因为 $G$ 和它的 Cholesky 因子 $L_G$ 完全不再需要了。

为什么可能是这样,边界在哪里? 一个可能的机制性解释:对于 MLP 层(论文画像研究的目标,见下面简化三)来说,输入激活 $a$ 是被上游整个网络塑造出来的——归一化统计量、之前的非线性、残差累积——在 batch 维度上可能发展出相当强的相关、拉长的结构。而输出梯度 $\delta$ 是一个已经通过链式法则被很多下游层平滑过的反向传播信号,经验上到达任何一层的时候可能已经比较接近各向同性了。这是论文没有明确检验的一个假设(论文报告了结果,但没有深挖为什么输入侧占主导),值得指出:完全有可能这种不对称性是 Muon 优化器特有的、是这个深度/宽度配置特有的、或者是这个特定预训练配方特有的,换到比如注意力层的梯度,或者换一个优化器的梯度统计量,不一定能原样迁移。

简化二:秩 32 的低秩近似

问题:一旦确定只用输入侧因子 $A$,它到底需要表示得多精确?对每个被选中的层都用完整的 $d_\text{in}\times d_\text{in}$ 矩阵 $A$(以及它的 Cholesky 因子)依然存储和计算都很贵。

设定:写 $A \approx \mathbb{E}[xx^\top]$(一个半正定的二阶矩矩阵,$A \succeq 0$),考虑它的低秩特征分解近似

$$A \approx U_r \Lambda_r U_r^\top, \tag{17}$$

只保留前 $r \ll d_\text{in}$ 个特征向量/特征值。用这个精简表示来构造输入侧变换,而不是用完整的 $A$,大幅削减了存储(只需要 $r$ 个长度为 $d_\text{in}$ 的向量,而不是一整个 $d_\text{in}\times d_\text{in}$ 矩阵)和应用变换的计算开销。

结果(论文 Table II):对 $A$ 做对角近似(也就是假设 $A$ 完全没有跨维度相关性——只有每个维度自己的方差)几乎不比不做变换的欧氏基线强(RelL2 $5.262\times10^{-2}$ 对基线的 $5.282\times10^{-2}$——基本没有收益)。这是一个重要的负面结果:它直接证明了几何感知的收益来自 $A$ 的非对角、跨维度相关结构,而不仅仅是给每个坐标单独乘上它自己的方差。一个朴素的”只按每个维度自己的量级归一化”方案——乍一看可能像是个显而易见的廉价近似——在这里基本上帮不上什么忙。相反,对完整矩阵做低秩近似几乎能拿回全部收益:秩 8 得到 RelL2 $1.958\times10^{-2}$,秩 16 是 $1.731\times10^{-2}$,秩 32 是 $1.721\times10^{-2}$——基本跟完整 $A$ 的结果 $1.768\times10^{-2}$ 一样好,存储和应用却便宜得多。论文最终选定秩 32”作为我们实现的运行点,而不是一个跟模型无关的常数”——诚实地承认这个具体数字是针对这个具体设置调出来的超参数,不是普适规律。

为什么会有这个模式? 直觉是,激活统计量里”有意思”的各向异性结构,通常集中在少数几个主导方向上(比如由归一化尺度、残差流随深度增长的量级、或者广义 LLM 量化文献里有很多记录的少数几个高度活跃的”离群”特征维度共同造成的、系统性偏大的方向),剩下一长串方向本来就已经比较接近各向同性了。一个只针对主导方向的低秩修正,就能捕捉到大部分有用的白化效果,而对角近似——只能单独缩放每个轴,永远不能混合或旋转维度之间的关系——完全没办法修正维度之间的相关性,这正是论文的对角近似结果失败的地方。

简化三:只在脆弱层上选择性部署

问题:即便已经做了只用输入侧和秩 32 这两步简化,给模型里每一层都应用几何感知分支依然会带来不小的额外开销。是不是所有层都真的需要它?

画像流程:在一个 32-GPU、600M 参数的运行上,跑欧氏基线预训练的前 100 步,对每个 MLP 层(fc1 和 fc2 子层都算),把它的梯度张量展平,算出会被使用的 FP8 量化缩放系数,把展平后的梯度编码成 FP8,然后统计编码后的数值有多少次撞到了 FP8 格式的上边界或下边界(也就是饱和/截断,而不是圆整到内部某个值)。把这些边界命中率在 100 步的画像窗口里分别对 fc1 和 fc2 层求平均,把每层上下边界命中率加起来,用这个总和作为数值脆弱度分数——一个直接的、代价很小的、跟模型无关的代理指标,衡量”在欧氏基线下,这一层的梯度现在到底被 FP8 量化扭曲得有多严重”。

发现(Figure 3):脆弱度高度集中。在 600M 模型的画像里,排名最靠前的 13 个脆弱层全都是 fc2 层(MLP 块里的第二个线性层,也就是把扩展后的隐藏维度重新投影回去的那个),而且在这组之后分数有一个明显的下降——排名第 14 的层(排名里第一个 fc1 层)分数明显更低。数字 13 明确不是被当作一个跟模型无关的常数;它只是”这个模型和这个训练配方下,脆弱度排序流程选出来的运行点”,同样的画像流程”可以针对不同的架构重新自动跑一遍来选择层集合”。

这条曲线的形状正是”硬性截断而不是连续、分级的部署”这个决定的依据:脆弱度不是平滑地分散在所有层上,而是集中在一个明确的组里,之后有一个明显的”拐点”式下降。这个拐点正是论文选择规则(“取分数明显下降之前的那些层”)具体操作化出来的层集合。

设计:只对这个通过画像筛选出来的脆弱层集合 $S$ 启用几何感知(输入侧、秩 32)分支,其余层照旧走朴素的欧氏快速路径。这正是 GIFT 实用性的核心:论文没有做一个笨重、均匀应用到所有地方的几何感知系统(比如每层都上完整 K-FAC 状态,论文的结果显示这样做既更贵、下游任务上还略逊于选择性设计——见 Table III 里 300M 的”完整 K-FAC”行只赢了 6/14,低于 GIFT 的 7/14),而是一个混合体:模型的大部分都保持跟欧氏基线一样便宜,额外的机制只集中在画像数据说真的有用的地方。

为什么偏偏是 fc2 层最脆弱? 一个可能的机制性原因,虽然论文没有明确说出来:fc2 子层的输入激活是 MLP 非线性函数(比如作用在扩展隐藏表示上的 GELU/SiLU)的输出,在更广义的量化文献里,这类激活是出了名的会发展出重尾、高度不均匀的统计量——少数几个”热门”特征通道可能携带远大于其他通道的量级。如果这就是机制,就直接解释了为什么 fc2 的输入侧几何(简化一里提到的 $A$)特别受益于白化,也和简化一发现输入侧(而不是输出侧)最重要这一点对上了。

这跟 LLM 量化文献里一个记录得比较充分的现象——有时被称为”激活离群值”或者”巨型激活”——是相通的:在很多 transformer 架构里,少数几个特定的隐藏维度(常常跟训练过程中出现的某些类似注意力汇聚点或者偏置行为有关)会发展出比典型维度大得多的激活量级,而且特别集中在某些 MLP 子层和网络的某些深度上。如果 GIFT 这种集中在 fc2 上的脆弱度轮廓反映的是同一个底层现象,那就说明这种具体的层选择模式(集中在 fc2,而不是均匀分散)不是这个特定 Llama/Muon/OpenWebText 配方的偶然产物,而是 transformer MLP 块在 FP8 量化下更结构性的一个特点——这是一个论文自己的画像数据支持、但没有明确拿去跟更广义的离群激活文献对照检验的假设。

拼在一起:完整的选择性 GIFT 算法

论文的算法 1 把三次简化结合起来,形式化成完整的、逐步、逐层的通信规则,还加了一个之前没提到的细节:误差反馈。下面是逐行讲解:

算法 1:带输入侧白化和误差反馈的选择性 GIFT 通信

需要:固定的被选中层集合 S(来自离线画像,简化三)
需要:对每个 l 属于 S,输入侧分解 A^(l) ≈ L_A^(l) (L_A^(l))^T,
     每 K = 50 个训练步刷新一次
需要:误差缓冲区 {R^(l)}_{l in S},初始化为 0
需要:worker 数量 N,FP8 量化器 Q(·; s),缩放规则 Scale(·)

 1: 对每个训练步:
 2:     对每个有权重梯度 W_g^(l) 的层 l:
 3:         如果 l 不在 S 中:                          # 朴素欧氏快速路径
 4:             s^(l) <- Scale(W_g^(l))
 5:             s^(l) <- AllReduceMax(s^(l))            # 在各 worker 间同步缩放系数
 6:             Q^(l) <- Q(W_g^(l); s^(l))               # 量化到 FP8
 7:             Q_hat^(l) <- AllReduceSum(Q^(l))          # all-reduce 量化后的值
 8:             W_hat_g^(l) <- Dequantize(Q_hat^(l); s^(l)) / N
 9:         否则:                                        # 几何感知分支
10:             U^(l) <- W_g^(l) . (L_A^(l))^{-T}          # 输入侧白化变换
11:             U_tilde^(l) <- U^(l) + R^(l)              # 加上上一步留下的误差反馈
12:             s_loc^(l) <- Scale(U_tilde^(l))            # 白化坐标里的本地缩放系数
13:             Q_tilde^(l) <- Q(U_tilde^(l); s_loc^(l))   # 量化白化后的梯度
14:             U_tilde_deq^(l) <- Dequantize(Q_tilde^(l); s_loc^(l))
15:             R^(l) <- U_tilde^(l) - U_tilde_deq^(l)     # 更新误差反馈缓冲区
16:             s_comm^(l) <- s_loc^(l)
17:             如果这一步是缩放系数同步步:
18:                 s_comm^(l) <- AllReduceMax(s_loc^(l))  # 偶尔也同步一下缩放系数
19:             结束
20:             Q_hat^(l) <- AllReduceSum(Q_tilde^(l))
21:             U_hat^(l) <- Dequantize(Q_hat^(l); s_comm^(l)) / N
22:             W_hat_g^(l) <- U_hat^(l) . L_A^(l),T       # 映射回欧氏坐标系
23:         结束
24:     结束
25:     优化器用 {W_hat_g^(l)} 更新参数
26: 结束

逐行走一遍两个分支:

• 第 3-8 行(未被选中的层):这就是任何欧氏 FP8 基线方案都会跑的标准流程——算本地缩放系数、在各 worker 之间同步这个缩放系数(这样大家都用同一个参考缩放系数量化)、量化、all-reduce 量化后的值(求和)、反量化、除以 worker 数量 $N$ 得到平均后的梯度。完全不涉及几何。
• 第 10 行(白化):对一个被选中的层,梯度先被右乘 $(L_A^{(l)})^{-\top}$——这正是式(11)/式(16)完整变换的输入侧简化,用这一层自己的 K-FAC 因子逐层应用。
• 第 11-15 行(误差反馈):GIFT 没有直接量化白化后的梯度 $U^{(l)}$,而是先加上上一步留下的误差反馈残差 $R^{(l)}$,量化这个和,然后把新的残差算成量化前的值和反量化后的值之差。这是一个更广义的梯度压缩文献里已经很成熟的技术(论文引用的 DGC 和 1-bit Adam 都用到过):不去丢掉单次量化圆整掉的东西,而是把这个圆整误差在本地累积起来,在下一步重新注入回去,这样系统性的量化偏差就不会随着很多步悄悄累积,而是最终会被折算回来,而不是永久丢失。在白化坐标系里应用误差反馈(而不是当作欧氏空间里一个不相关的事后修正)让这个技巧跟通信公式的其余部分保持一致。
• 第 16-19 行(缩放系数同步):每个 worker 用自己(局部不一样)的梯度算出来的本地缩放系数 $s_\text{loc}^{(l)}$,通常直接给那个 worker 自己的量化步骤用,但在指定的”缩放系数同步步”上,论文会先在各 worker 间 all-reduce 出最大的本地缩放系数,让大家都用一个共同的参考缩放系数量化——这是每个 worker 用自己本地缩放系数(省掉一次同步往返,但精度是本地的)和跨 worker 用统一缩放系数(避免不同 worker 圆整行为微妙不一致,但需要额外同步)之间的标准取舍。论文提到,在变换后的坐标系里用本地缩放系数”在我们的消融实验里比同步一个全局缩放系数保真度更好”——也就是说,默认模式偏向每个 worker 用本地缩放系数,只是偶尔才做全局同步。
• 第 20-22 行(all-reduce 和映射回去):量化后的白化值跟朴素路径一样被 all-reduce 和反量化,然后——关键的一步——通过右乘 $L_A^{(l),\top}$ 映射回欧氏坐标系,这样下游消费 $\widehat W_g^{(l)}$ 的任何代码(优化器)完全不需要知道几何感知分支存在过。

论文强调的两个实现细节:第一,几何感知分支真的只在小的固定子集 $S$ 上活跃,所以模型通信代码路径的绝大部分跟朴素基线是逐字节完全一样的;第二,误差反馈是具体在变换后的坐标系里发生的,论文的论证是这样能保持”跟通信公式一致,而不是把误差反馈当作一个独立的事后修正”——这是一个设计一致性的论证,而不是一个测量出来的消融实验,值得注意这是论文断言了一个设计原则、却没有跟”在欧氏空间里、映射回去之后再做误差反馈”做并排对比的地方。

这张图是算法 1 用伪代码描述的东西在系统层面的直观呈现:一个所有层都要经过的共享”通信核心”(量化 → all-reduce → 反量化),几何感知分支被画成一条可选的岔路——进入核心之前的白化变换,以及核心之后的映射回去变换——只有被选中的脆弱层才会走这条岔路。这张图直观地说明了为什么 GIFT 的开销是有界的:那条昂贵的岔路在架构上被限制在图里一小部分节点上,而不是均匀涂抹在每条边上。

实验结果

论文从四个角度评测 GIFT:系统层面的扩展性收益、验证损失保持情况、下游任务质量保持情况、以及内存开销。

实验设置

在 OpenWebText 上做两种 LLaMA 式预训练配置:一个 约 300M 参数的模型,序列长度 4096;一个 约 600M 参数的模型,序列长度 2048(不同序列长度是为了适应 GPU 内存限制,不是一个独立的实验变量)。除非特别说明,端到端实验都用 32 张 GPU、全局 batch size 512、micro batch size 4、Muon 优化器、学习率 $5\times10^{-4}$,余弦衰减到最小值 $5\times10^{-5}$。所有实验都跑在 TACC 的 Vista 超级计算机上,用 NVIDIA GH200 Grace Hopper 超级芯片节点(每张 GPU 96GB HBM3,CPU-GPU 紧密集成设计)。

一个值得特别指出的、方法论上比较讲究的选择:论文没有简单地拿 GIFT 跟单一的”欧氏 FP8”基线比——它先在 600M 规模上多跑了一些实验,来确定到底哪个基线变体最强,这样主要的对比就不会不小心跟一个弱的基线比。它发现 BF16 梯度通信在 7/14 下游任务上跟 FP32 打平,但绝对任务数值整体上稍低一点(所以 FP32 依然是参考基准);per-block 欧氏 FP8(block size 512)只在 4/14 上跟 FP32 打平,不如按层的欧氏 FP8 的 5/14——所以论文全程用更强的按层欧氏 FP8作为主基线,而不是一个弱的稻草人变体。

系统层面收益:随 GPU 数量扩展(Figure 5)

Figure 5 画的是相对 FP32 基线的单步时间改善(正值 = 比 FP32 快,负值 = 比 FP32 慢),随 GPU 数量增加而变化的曲线,300M 和 600M 两个模型都有,比较的是欧氏 FP8 基线和 GIFT。三个观察:

1. 低精度通信的优势随规模增长——相对 FP32 的加速比在 GPU 数量更多的时候更明显,600M 模型尤其如此,因为随着模型和集群规模增长,通信占总单步时间的比例越来越大(要同步的参数更多,参与同步的 worker 也更多)。
2. GIFT 一直比朴素欧氏基线慢,因为它引入了额外的几何感知计算(白化变换、映射回去变换,以及每 50 步一次的 Cholesky 因子维护)——这是一个诚实的、预期之中的代价,论文没有藏着掖着。
3. 关键的是,GIFT 依然保留了相对 FP32 的低精度加速中相当大的一部分,而且 GIFT 和欧氏基线加速比之间的差距会随 GPU 数量增加而收窄——意味着在测试过的最大规模上,通信效率最重要的时候,GIFT 的额外开销相对来说变得没那么举足轻重了。

仔细看这张图:两条曲线(欧氏基线和 GIFT)都随 GPU 数量增加而上升(相对 FP32 的加速越来越大),而两条曲线之间的垂直差距,就是 GIFT 为了换取下游质量优势所付出的系统代价。论文的论证是,这个垂直差距不会跟曲线本身上升得一样快——也就是说,GIFT 的绝对开销大致是每步固定的,而它所依托的基线通信节省却随规模持续增长,所以 GIFT 的开销在测试过的最大规模上,占整体的比例正在缩小。

验证损失:一个刻意”不刺激”的结果(Figure 6)

Figure 6 比较了预训练期间 FP32、欧氏基线、GIFT 的验证损失轨迹,两个模型规模都有。

这里的头条发现几乎让人有点扫兴:GIFT 和欧氏基线的验证损失曲线几乎无法区分,两者在整个训练过程中都跟 FP32 参考基准保持得比较接近。论文把这个当作一个重要的方法论要点,而不是一个令人失望的结果:它论证说,验证损失虽然是一个有用的信号,但对通信坐标系选择到底怎么影响模型实际下游有用性来说,是一个不完整的代理——不同通信方案之间的差异”可能在下游表现上比在预训练损失曲线本身上更明显”。这为论文的主要实证主张铺好了路,而这个主张活在下游任务表格里,不在损失曲线里。

下游任务表现:信号真正显现的地方(Table III)

Table III 报告了 14 个不同任务(BOOLQ、CB-ACC、COPA、MUL-RC、RCD-F1、RTE、WiC、WSC、LAMBADA、RACE、M-QA、PIQA、WinoGrande、LAMBADA-standard)的逐任务下游准确率,涵盖 FP32、BF16、per-block 欧氏 FP8、按层欧氏 FP8、完整 K-FAC、GIFT,300M 和 600M 两个规模都有。论文强调的汇总指标是”相对 FP32 的胜场数”——一个方法的准确率超过 FP32 参考基准的任务数量:

模型	方法	相对 FP32 的胜场(共 14)
600M	BF16	7
600M	per-block 欧氏 FP8	4
600M	按层欧氏(主基线)	5
600M	GIFT	7
300M	按层欧氏(主基线)	4
300M	完整 K-FAC	6
300M	GIFT	7

在两个规模上,GIFT 都赢了 14 个任务里的 7 个,明显强于按层欧氏基线(600M 是 5/14,300M 是 4/14)——而且值得注意的是,还强于 300M 规模测试的更重的完整 K-FAC变体(6/14),后者对每个符合条件的层都用完整双侧变换,而不是 GIFT 那种选择性、只用输入侧、秩 32 的设计。这最后一个对比,是论文”选择性、简化”设计理念的最强单点证据:处处都做更多几何感知计算,不仅更贵,实际上还不如有针对性地做更少的几何感知计算——很可能是因为完整 K-FAC 状态在相同的实践约束下(有限的画像窗口、周期性刷新)被估计和刷新,给那些原本就不需要几何修正的层引入了自己的估计噪声,却没有换回任何东西。

论文很小心地把这个结论框定正确:“这些结果不应该被解读为 GIFT 在每个任务上都匹配 FP32”——没有任何方法能做到这一点,逐任务的结果确实是真实混杂的(看 Table III 的原始数字,GIFT 在某些单个任务上是输给欧氏基线的,即便在汇总的”相对 FP32 胜场数”上赢了,比如 600M 的 CB-ACC:欧氏 0.4107 对 GIFT 0.3036)。这个主张具体是关于跨任务保持轮廓的——GIFT 在整个任务套件上更一致地保持接近或高于 FP32 参考基准,不是说它在每个单独任务上都占优。

内存开销(第六节 E 小节)

相对欧氏基线,GIFT 让 300M 模型的内存使用增加了 3.33%,600M 模型增加了 8.98%。这部分开销来自存储秩 32 的低秩输入侧因子和误差反馈残差缓冲区,但只针对小的被选中层子集 $S$——论文明确把这个框定为”是有限的,而不是像完整的双侧逐层 K-FAC 状态那样按比例增长”,这再次强调了为什么选择性部署设计(简化三)对实用性很重要,不只是对前面讨论的精度结果重要。

端到端系统收益:7.6% 时间缩减,明确小于基线的 10.79%

把各个部分拼到完整预训练规模上(64 张 GH200 超级芯片,Llama-600M):GIFT 相对 FP32 减少了 75.0% 的梯度通信载荷体积(用 FP8 代替 FP32 的机械结果,任何 FP8 方案都能拿到),这转化成了 7.6% 的端到端预训练时间缩减。论文很直接地承认,这小于同规模下纯欧氏 FP8 基线本身的 10.79% 缩减——GIFT 明确牺牲了一部分原本纯通信时间加速,来换取 Table III 记录的更好下游保真度轮廓。这个取舍到底值不值,完全取决于一个实践者有多看重那最后几个百分点的任务保持质量,相对于预训练的实际时钟时间——这是一个论文提出来、而不是替读者解决的真实成本收益问题。论文还给出了一个”为什么哪怕是一个不起眼的百分比在规模上也很重要”的论证:现代 LLM 是在数万张 GPU 上跑好几个月,所以”即便 7.6% 的时间缩减也可能折算成数百万美元”,而且释放出来的互联带宽本身就是一种宝贵的共享资源,因为算力增长的速度比互联带宽快。

这对实践者意味着什么

跳出论文具体的数字往回看,值得说清楚一个做大规模 LLM 预训练的实践者到底应该从这项工作里学到什么,毕竟”在通信时用 K-FAC 白化坐标”目前还不是任何主流训练框架里的一个勾选项。

如果你已经在用 FP8 梯度通信、但还没考虑过坐标系效应,第一个可以直接上手的收获,是那套画像流程本身(简化三),跟你要不要采纳白化变换完全无关:对你自己的模型、在你自己的欧氏 FP8 基线下跑一次代价很小的 100 步边界命中率画像,直接告诉你现在到底是哪些层在真实训练里承受着最严重的量化扭曲——这个诊断本身就有用,哪怕你还没决定要不要搭建白化的那套机制。

如果你在权衡额外的工程复杂度到底值不值,论文给出的诚实成本收益是:预期个位数百分比的内存开销(实验里 3%-9%),预期要牺牲纯 FP8 相对 FP32 原本能拿到的加速的大约四分之一到三分之一(600M/64-GPU 场景下是 7.6% 对 10.79%),换来的是明显更好的下游任务保持轮廓(14 个任务里 7 个 vs. 4-5 个)。这个取舍值不值,取决于你的训练流程更受限于”在固定算力预算下的时钟时间”,还是更受限于”在固定算力预算下能拿到的最好模型质量”——这是两个不同的优化目标,而 GIFT 明确是在为后者做优化,为此付出了对前者的一些代价。

如果你在构建或者扩展一个分布式训练框架(Megatron-LM 式、DeepSpeed 式、或者自研栈),GIFT 增加的实际工程面是:(1)一次离线画像来选出脆弱层,理想情况下模型架构或训练配方有意义地变化时就该重新跑一遍;(2)对每个被选中层,周期性(每 $K$ 步)维护一个低秩的输入侧统计量,这需要自己的计算/内存预算,理想情况下应该跟其他训练步骤的工作重叠,而不是阻塞它;(3)在通信路径里,专门为被选中的层插入一对额外的矩阵乘法(进来的白化变换、出去的映射回去变换),这需要跟已经在用的任何分片/并行策略(ZeRO、张量并行、流水线并行)干净地组合起来——论文没有明确讨论这一点,这本身就值得指出,是任何想在比论文自己的 32-GPU 纯数据并行配置更复杂的并行场景下,把这个想法用起来的人需要面对的一个集成问题。

作者自己声明的局限性

论文在自己的”局限性”小节里明确列出了三点:

1. 规模:端到端实验限于两个中等规模(300M/600M 参数)的 Llama 模型,固定的预训练配方;测试更大的模型、更长的训练周期、更广泛的硬件环境,留给了未来工作。
2. 评测广度:虽然下游结果已经显示出 GIFT 相对欧氏方案的一致优势,论文承认这会”通过更多的随机种子和更大的基准测试套件进一步加强”——承认目前 14 个任务、看起来是单一种子的对比,附带着真实的统计不确定性,而论文没有在任何地方量化这个不确定性。
3. 压缩栈的范围:这项工作只是通信层面的低精度压缩,不是一个完全量化的预训练栈(不涉及权重、激活、优化器状态)——这种几何感知原则跟这些其他形式的低精度预训练组合起来的时候,是否依然有用,被明确留作了开放问题。

批判性分析:弱点与改进空间

弱点与不足

• 下游结果的单种子不确定性。 Table III 给出的是 14 个下游任务的点估计,没有报告方差、置信区间,或者多种子平均。有几个报告的”胜场”差距很窄——比如 600M 规模上,GIFT 的 WiC 分数(0.5031)对欧氏基线(0.4969),只差 0.62 个百分点,而这个任务本身是个二元性质的词义消歧分类基准,更广义的文献里通常报告说这类任务的跑次间噪声相当大。没有种子或方差范围,确实很难判断 Table III 里到底有多少个单独的任务差异能在重复运行中站得住脚,尽管汇总的”相对 FP32 胜场”模式(7 对 4-5)差距足够大,大概率是能站得住的。
• K-FAC 因子刷新周期(每 50 步)是断言的,没有做消融。 论文把这个当作固定超参数报告,没有展示比如每 10 步(更新鲜、更贵)或每 200 步(更陈旧、更便宜)刷新会发生什么——留下了一个悬而未决的问题:50 是不是接近最优,还是仅仅是”第一个试出来还算凑合的值”。考虑到整个方法都建立在 K-FAC 因子是对当前输入激活几何的一个相当准确的局部近似这个前提上,而梯度/激活统计量在一次训练里可能发生有意义的漂移(尤其是训练早期,模型表示变化最快的时候),这个周期选择理应有一个论文没有提供的消融实验。
• “13 个脆弱层”是单次运行的画像快照,没有做鲁棒性检验的选择规则。 画像流程(Figure 3)只跑了单次 32-GPU 600M 运行的前 100 个训练步。论文没有报告换一个随机种子、换一个数据顺序、或者稍微晚一点的画像窗口,是否会选出同样的前 13(或者任意前 $k$)层集合——也就是说,这个脆弱度排序到底是稳定的,还是对画像窗口的选择足够敏感,以至于换一个 100 步的快照可能选出一个有意义地不同的层集合。
• 没有把 GIFT 额外花的时间具体分解到底去哪了。 Figure 5 显示 GIFT 在总单步时间上比欧氏基线慢,但论文从来没有把这个开销拆解成”花在白化/映射回去矩阵乘法上的时间”对”花在维护/刷新 K-FAC 因子上的时间”对”因为额外同步而损失的时间”。这在实践上很重要:如果因子维护占主导,一个更聪明的刷新计划,或者一个异步/重叠的因子更新实现,就可能补上跟朴素欧氏基线之间的大部分差距;如果变换本身的矩阵乘法占主导,那是一个难得多的、工程上没法绕开的代价,而论文给出的数字没法让读者判断到底是哪种情况。

作者低估或遗漏的局限性

• 论文从来没有对真实训练梯度直接测量变换前后的各向异性,尽管各向异性是论文整个因果故事的核心。Figure 1 被描述为示意性的(“真实梯度样本”的二维投影),但结果部分没有任何地方报告一个定量的各向异性指标(比如 $A$ 或者原始梯度协方差的特征值比/条件数度量,跟踪整个训练过程)。这是一个真实的空白:整个机制都依赖于梯度真的有意义地各向异性、而且在整个训练过程中一直保持”白化可操作”的各向异性,但论文对这个主张的证据完全是间接的(通过应用变换的下游效果),而不是直接的(通过测量各向异性本身、展示它在白化之后减小了)。
• 选择性层选择跟 Muon 优化器之间的具体相互作用没有被探索。 GIFT 只用 Muon 优化器评测过,而 Muon 本身有一些有趣的几何性质(Muon 通过 Newton-Schulz 迭代对更新做正交化,这本身就是权重空间里一种”偏好方向”的概念)。论文没有讨论 GIFT 基于 K-FAC 的白化,跟 Muon 自己的几何预条件之间,到底是相互促进、中性、还是相互冲突——两个不同的”几何感知”机制,在流水线的不同位置(通信时刻的白化 vs. 更新时刻的正交化)对同一个梯度起作用,这正是可能真正有影响的那种相互作用,而想要拿着 GIFT 换另一个优化器(Adam、Lion、Shampoo)的读者,没有任何证据可以据此推理。
• GIFT 能推广到未来 FP4 通信的说法有说,但没有测试。 结论部分说”GIFT 背后的坐标变换原则不是 FP8 专属的……同样的想法可以应用到未来的 FP4 通信”,但 FP4 的动态范围比 FP8 窄得多,白化变换的收益跟一个已经精度饥渴得多的格式相互作用,可能完全不一样——比如,完全有可能(不管哪个方向都没测过)FP4 的粗糙会把坐标系条件更好带来的边际收益给淹没掉,或者反过来,FP4 比 FP8 更需要几何感知。不管哪个方向,都是论文当作一个前瞻性主张提出来、却没有任何支撑证据的真实实证问题。

具体的改进建议

• 给下游评测(Table III)报告至少 3 个种子的逐任务方差,或者至少从现有的单种子评测集里做自举置信区间,这样读者才能分辨”GIFT 稳健地赢了 7/14”和”GIFT 赢了 7/14,但其中 2-3 个胜场在噪声范围内”这两种情况。
• 对 K-FAC 刷新周期做消融($K=10, 25, 50, 100, 200$),同时报告保真度(像 Table I 那样的 RelL2/Cos)和随 $K$ 变化的时钟开销,给实践者一个真实的成本-精度前沿,而不是一个固定的运行点。
• 直接测量并报告梯度/激活的各向异性(比如 $A$ 的条件数或者特征值谱),在训练过程的几个时间点上,既是为了验证论文的核心因果故事,也是为了检验各向异性在训练中大致是平稳的(这样固定刷新周期就说得通),还是会随训练大幅漂移(这样就应该用自适应的刷新周期)。
• 测试层集合的稳定性,用不同的随机种子或者更晚的画像窗口(比如第 500-600 步而不是第 0-100 步)重新跑一遍 Figure 3 的画像流程,报告两次选出来的前 13 层集合之间的重合度——这能直接回答选择性部署设计到底是鲁棒的,还是对画像运行的方差很脆弱。
• 把 Figure 5 的开销分解成变换计算时间、因子维护时间、同步时间,找出 GIFT 增加的成本里哪一部分是缩小跟欧氏基线原始加速比之间差距的最有希望的工程目标。
• 至少再测试一个优化器(Adam 或者 Shampoo 式的方法,两者都跟 Muon 一样在 LLM 预训练里被广泛使用),来确定输入侧占主导这个发现(简化一)和 13 个脆弱 fc2 层这个发现(简化三)到底是 Muon 特有的产物,还是更普遍成立的——这大概是确立这个方法更广泛适用性的、单点价值最高的补充实验,因为目前的结果跟论文一般化的框定是一致的,但没有排除比论文暗示的范围窄得多的适用性。

一个手算的玩具数值例子:给一个各向异性玩具梯度做白化

式(5)-(14)的推导比较抽象,亲眼看一下白化变换作用在具体数字上会有帮助,所以这一节构造了一个小的、完全手算、独立验证过的玩具例子(不在论文里),用简化一里输入侧变换的一个 $2\times 2$ 类比。

设定。 假设一个层的输入激活,为了图示方便投影到只有 2 维,二阶矩矩阵是

$$A = \begin{pmatrix} 4 & 3 \\ 3 & 4 \end{pmatrix}. \tag{A1}$$

这个 $A$ 正是论文 Figure 1 图示的那种各向异性对象:它的特征值是 $\lambda_1 = 7$(特征向量 $\tfrac{1}{\sqrt2}(1,1)^\top$)和 $\lambda_2 = 1$(特征向量 $\tfrac{1}{\sqrt2}(1,-1)^\top$)——“长”方向和”短”方向的比例是 7:1,也就是一个中等程度的雪茄状激活云,沿着 $(1,1)$ 对角线方向拉长。

第一步——Cholesky 因子。 一个有效的 Cholesky 分解 $A = L_A L_A^\top$ 是

$$L_A = \begin{pmatrix} 2 & 0 \\ 1.5 & \sqrt{4-2.25} \end{pmatrix} = \begin{pmatrix} 2 & 0 \\ 1.5 & 1.3229 \end{pmatrix}. \tag{A2}$$

(验算:$L_A L_A^\top = \begin{pmatrix} 4 & 3 \\ 3 & 2.25+1.75\end{pmatrix} = \begin{pmatrix}4 & 3\\3&4\end{pmatrix} = A$,跟式 A1 一致。)

第二步——一个玩具梯度行。 假设某一层权重梯度的某一行(也就是某个输出神经元在 2 个输入维度上的梯度向量)是 $g = (5, 1)^\top$——这个向量恰好主要沿着标准基方向,没有跟 $A$ 的主导特征向量对齐。这正是欧氏量化的对抗情形:$A$ 方差最大的方向是 $(1,1)$ 对角线,但 $g$ 主要沿着原始 $x$ 轴,所以一个按 $g$ 自身量级校准的 per-tensor 欧氏缩放系数,完全没告诉你 $g$ 每个分量按 $A$ 加权到底有多”重要”。

第三步——应用输入侧白化变换。 按照简化一里输入侧的简化(式 16),白化后的梯度行是 $\tilde g^\top = g^\top L_A^{-\top}$,也就是 $\tilde g = L_A^{-1} g$。从式(A2)算出 $L_A^{-1}$:

$$L_A^{-1} = \begin{pmatrix} 0.5 & 0 \\ -0.5669 & 0.7559 \end{pmatrix}, \qquad \tilde g = L_A^{-1} g = \begin{pmatrix} 0.5 \times 5 + 0 \times 1 \\ -0.5669\times 5 + 0.7559 \times 1\end{pmatrix} = \begin{pmatrix} 2.5 \\ -2.079\end{pmatrix}. \tag{A3}$$

第四步——比较模拟的量化误差。 假设我们的玩具 FP8 式量化器把每个坐标圆整到 $\Delta = s / 8$ 的最近倍数,其中 $s = \max(|\cdot|)$ 是被量化向量的最大绝对值(一个每边 8 级的、FP8 离散网格的玩具替身)。对原始梯度 $g=(5,1)$:$s_\text{raw} = 5$,$\Delta_\text{raw} = 0.625$,所以 $g$ 圆整到 $(5.0, 1.25)$(1 距离 0.625 的倍数 0.625 和 1.25 几乎正好在中间,取 $\hat g = (5.0, 1.25)$,误差 $= (0, 0.25)$,$\|\hat g - g\|_2 = 0.25$)。对白化后的 $\tilde g = (2.5, -2.079)$:$s_\text{white} = 2.5$,$\Delta_\text{white} = 0.3125$,所以 $\tilde g$ 圆整到 $(2.5, -2.1875)$(离 $-2.079$ 最近的 0.3125 倍数),误差 $=(0, -0.1085)$,$\|\hat{\tilde g} - \tilde g\|_2 = 0.1085$。把这个白化坐标系里的量化误差(像式 14 那样)映射回去,跟原始梯度空间里的误差做同口径比较:映射回去之后的误差向量范数,对这个具体的 $L_A$ 来说,比原始的 $0.25$ 误差要小,因为承载这个圆整误差的那个白化坐标,映射回去之后对应的正是 $A$ 的一个低方差方向,而不是一个未加区分的原始坐标误差。这个玩具计算刻意做得足够小、可以手算验证,但精确地捕捉到了机制本身:白化重新分配了量化误差落在哪里,把它从 K-FAC 度量认为高度重要的方向,挪到了它认为不那么重要的方向——这正是式(9)椭球到球的映射对一般情形所形式化的机制。

GIFT 跟以往低精度通信方案的对比

把 GIFT 的设计直接放到它引用的那些替代方案旁边很有用,因为”用 FP8”不是单一的方案,而是一整族设计点,搞清楚 GIFT 在其中的位置,能说清楚它的贡献到底是什么(以及不是什么):

方法	改变的是什么	坐标系	处理各向异性吗?	每层额外状态
朴素欧氏 FP8(按层或按块缩放)	缩放 + 量化	原始参数/梯度坐标系	否——每个张量/块一个标量缩放系数	无
FP8-LM	权重、激活、梯度、一阶矩都做缩放 + 量化	原始欧氏,加动态范围技巧	否	每张量动态范围统计量
COAT	混合粒度的激活/二阶矩量化	原始欧氏	否	每张量/组的缩放统计量
SDP4Bit	通过傅里叶变换 + 分层 all-to-all 实现 4-bit	傅里叶变换后的坐标系	部分(通过另一种变换,不是曲率相关的)	傅里叶域缓冲区
PowerSGD	低秩梯度近似	原始欧氏,但降秩	否——降秩不改变几何形状	低秩因子(用于近似,不是白化)
完整 K-FAC(本文自己的强消融)	双侧白化,每个符合条件的层	K-FAC 白化(双侧)	是,完全处理	每层完整的 $A$、$G$ Cholesky 因子
GIFT(本文最终设计)	单侧白化,只选中的层	K-FAC 白化(只输入侧)	是,在脆弱层上	秩 32 因子 + 误差缓冲区,只在选中的层

这张表让论文的定位一目了然:GIFT 不是在竞争用哪种数值格式(它明确保持 FP8 不变,不像 SDP4Bit 转向 4-bit,或者 PowerSGD 降秩)——它是这个对比里唯一一个核心机制是在量化之前重塑坐标系的方法,而且即便在这个想法内部,它也刻意做得更少(单侧、低秩、选择性),而不是理论上”更完整”的完整 K-FAC 替代方案,因为 Table III 的下游结果显示,做得少但打得准,胜过处处都上完整变换。这是论文最不动声色、却最重要的一个实证主张:更多的几何修正,一旦把估计噪声和开销算进去,并不必然等于更好的几何修正。

复现笔记

• 论文给出了具体的超参数(batch size 512、micro-batch 4、Muon 优化器、学习率 $5\times10^{-4}$ 余弦衰减到 $5\times10^{-5}$、300M/600M 模型分别用序列长度 4096/2048、秩 32 输入侧因子、每 50 步刷新一次、600M 画像里选出前 13 个脆弱层),这些合起来让核心配方是可以重建的。
• 具体的脆弱层集合被描述成一个流程(把欧氏基线画像 100 步,按 FP8 边界命中率排序,取分数明显下降之前的那些层),而不是一个固定的通用列表——意味着在不同的模型/配方上复现,需要重新跑这个画像步骤,而不是直接复用论文具体的层索引。
• 本次评审能看到的论文文本里没有提到代码或制品的发布;引用的 Megatron-LM 框架(广泛可获取)和 OpenWebText 数据集(公开可获取)都是标准的、可复现的组件,但 GIFT 特有的变换/误差反馈实现本身,在没有公开代码发布的情况下,需要从算法描述(算法 1)和第四节的公式重新实现。
• Vista 超级计算机 / GH200 硬件设置是一个具体的、不是普遍可获取的环境;复现精确的 7.6%/10.79% 时间数字,需要类似的互联和 GPU 代际,不过 GIFT 和欧氏基线之间的相对比较,应该比绝对百分比对硬件的敏感度低得多。

常见混淆点辨析

GIFT 设计里有几个地方很容易跟表面上类似的想法混淆在一起;这一节直接处理最可能引起困惑的几个点。

“这不就是 K-FAC 吗?” 不是——作为优化器的 K-FAC 用 $F^{-1}$(或它的近似)在参数更新之前预条件化梯度,永久地改变了优化轨迹本身。GIFT 用的是同一套分解($A \otimes G$),但只在通信步骤前后作为一个临时的换基(式 11-14):梯度被映射到白化坐标、量化、通信、然后直接映射回去——等优化器看到它时,它(除了 FP8 量化误差)在数值上跟普通欧氏方案下本该得到的梯度是一样的。优化轨迹没变,变的只是平均后的梯度经过 FP8 往返有多保真。

“白化梯度不会改变模型学到的东西吗?” 不会像自然梯度方法那样。因为式(14)的变换精确地逆转了式(11)的变换(除了中间引入的 FP8 量化噪声),对优化器输入的净效应是:同样的梯度,只是带着一个不同(希望更小)的 FP8 扭曲。如果 GIFT 用一个无损的通信格式(比如在白化坐标里用 FP32 通信),那映射回去就是一个数学上的恒等变换,GIFT 就可以证明等价于朴素基线。GIFT 的白化坐标之所以有影响,唱一的原因就是量化到一个像 FP8 这样有损、窄窄的格式发生在两个变换之间——白化只改变了FP8 往返会摧毁多少信息,不改变梯度代表什么。

“为什么不直接用对角(每维度)缩放代替完整的几何变换?” 这正是 Table II 的消融实验回答的问题:对 $A$ 做对角近似几乎不比不变换的欧氏基线强(RelL2 $5.262\times10^{-2}$ 对基线的 $5.282\times10^{-2}$),因为一个对角矩阵只能缩放单个坐标轴,永远不能旋转或混合它们之间的关系——而论文的证据(对比秩 8 就能带来大得多的收益)是,各向异性修正有用的部分具体需要修正跨维度相关,而一个纯对角(按轴)缩放在数学上无法表示这种关系。

“GIFT 需要每个 worker 自己算自己的 K-FAC 因子,还是大家共享一个因子?” 论文的算法(“对每个 l 属于 S,输入侧分解……每 K 个训练步刷新一次”)暗示着每层一个因子,周期性刷新;但它没有明确说清楚这个因子到底是从那个 worker 自己的激活 batch 本地算出来的,还是从全局聚合的激活统计量算出来、广播/同步给各 worker 的。这是一个真实的实现细节,同时影响因子的保真度(全局聚合的 $A$ 应该比任何单一 worker 的本地 batch 更能代表”真实”的总体统计量)和维护它的同步开销(全局同步的因子需要自己的周期性集合通信,独立于梯度 all-reduce)——论文的表述在这个地方留下了模糊地带,值得任何想精确复现这个系统的人注意。

结论

GIFT 的贡献在概念上很窄,但确实有用:它精确定位并修复了 FP8 梯度通信一个之前没怎么被仔细审视过的失效模式——各向异性梯度几何跟坐标轴对齐量化之间的不匹配——用的工具(K-FAC 的 Kronecker 分解曲率近似)是从另一个子领域(自然梯度优化)整体借来的,但被重新用在了一个真正不同的活儿上(临时的通信坐标,不是优化器预条件器)。论文最有说服力的证据,不是那个头条的 7.6% 时间缩减(论文老老实实地承认它比纯 FP8 基线自己的加速比要小),而是那个谨慎的三步简化过程——只输入侧、秩 32、选择性层部署——每一步都用一个针对性的消融来证明这次简化几乎不损失保真度,最终得到的设计在下游任务上甚至打败了更重的”处处都上完整 K-FAC”变体,而额外内存开销只有个位数百分比。真正开放的问题,与其说是这个几何想法本身是不是站得住(这里的推导和消融都很仔细),不如说是这些具体的实证发现到底能推广多远:输入侧占主导、以及具体的脆弱层轮廓,到底是这个 Muon/Llama/OpenWebText 组合特有的属性,还是 LLM 预训练梯度更普遍的一个属性;以及同样的坐标变换原则,应用到动态范围更窄的 FP4 上时,是会失效、退化、还是变得更有价值。这些正是接下来该做的实验,论文自己的局限性小节也指向了其中的大部分。