通过几何感知坐标系提升低精度梯度通信精度,加速大模型训练
GIFT: Geometry-Informed Low-precision Gradient Communication for LLM Pretraining

- 将梯度变换到近各向同性空间再量化,减少方向依赖失真
- 在64张GH200芯片上使Llama-600M预训练时间缩短7.6%
- 无需改动优化器和通信协议,适合大规模分布式训练场景
大语言模型预训练中梯度通信是主要的扩展瓶颈。使用FP8、NVFP4等低精度格式传输梯度可显著降低通信量。现有方法在欧氏空间中进行线性或非线性映射量化,因梯度高度各向异性导致方向依赖失真,影响模型性能。本文提出GIFT,一种几何感知的低精度梯度通信方法,在几何感知坐标系中进行量化。通过将梯度转换至近各向同性空间,GIFT使低精度表示更贴近高精度版本。GIFT仅改变低精度通信的坐标系,不修改优化器、训练流程、通信集合或低精度格式。我们还设计了基于低秩近似和选择性应用的简化算法,平衡计算开销与通信压缩效果。在Llama-300M和Llama-600M模型上验证收敛性:在64张NVIDIA GH200 Superchips上,GIFT使Llama-600M预训练总时长减少7.6%,且在相同优化器和通信路径下,下游任务保留能力优于直接欧氏空间FP8通信。
原文摘要 · Abstract (English)
Gradient communication is a primary scaling bottleneck in large language model (LLM) pretraining. Communicating gradients in low-precision formats, such as FP8 and NVFP4, can significantly reduce the communication volume. Existing methods quantize gradients via linear or nonlinear mappings in Euclidean space, often degrading model performance because highly anisotropic gradients incur direction-dependent distortion. We present GIFT, a geometry-informed gradient scaling method that performs low-precision communication in geometry-aware coordinates. By transforming gradients into a near-isotropic space before quantization, GIFT makes low-precision representations substantially more faithful to their high-precision counterparts. GIFT only changes the coordinate system used for low-precision gradient communication and does not change the optimizer, training recipe, communication collective, or low-precision format. We also develop a simplified geometry-aware transformation algorithm with low-rank approximation and selective application to balance the computation overhead and communication reduction. We examine the empirical convergence of GIFT using Llama-300M and Llama-600M models. Our results show that GIFT reduces the end-to-end pretraining time of Llama-600M by 7.6% on 64 NVIDIA GH200 Superchips, while improving the downstream task preservation profile over direct Euclidean FP8 communication under the same optimizer and communication path.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。