发表机构
Rutgers University（罗格斯大学）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
研究针对LLM预训练中梯度通信瓶颈，提出几何感知梯度缩放方法GIFT，通过变换梯度到近各向同性空间进行低精度通信，开发简化算法平衡开销与减少量，经实验验证可减少预训练时间并改善下游任务保留情况。
AI 中文摘要
梯度通信是大语言模型（LLM）预训练中的主要扩展瓶颈。以低精度格式（如FP8和NVFP4）通信梯度可显著减少通信量。现有方法在欧几里得空间中通过线性或非线性映射量化梯度，常因高度各向异性梯度导致方向相关失真而降低模型性能。我们提出GIFT，一种在几何感知坐标中进行低精度通信的几何感知梯度缩放方法。通过在量化前将梯度变换到近各向同性空间，GIFT使低精度表示更忠实于高精度对应物。GIFT仅改变低精度梯度通信所用坐标系，不改变优化器、训练方法、通信集合或低精度格式。我们还开发了一种简化的几何感知变换算法，通过低秩近似和选择性应用来平衡计算开销和通信减少。我们使用Llama-300M和Llama-600M模型检验了GIFT的经验收敛性。结果表明，在64个NVIDIA GH200超级芯片上，GIFT将Llama-600M的端到端预训练时间减少了7.6%，同时在相同优化器和通信路径下，相较于直接的欧几里得FP8通信，改善了下游任务保留情况。
英文摘要
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 GeoFP8, 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, GeoFP8 makes low-precision representations substantially more faithful to their high-precision counterparts. GeoFP8 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 GeoFP8 using Llama-300M and Llama-600M models. Our results show that GeoFP8 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.
Comments12 pages, 6 figures, 3 tables
