EQuARX让TPU上的分布式训练更快，量化通信不丢精度。
EQuARX: Efficient Quantized AllReduce in XLA for Distributed Machine Learning Acceleration
- 在XLA编译器中实现块级动态量化AllReduce，适配TPU架构。
- 8位整数量化下比基线BF16快1.8倍，不同网络拓扑均有效。
- 部署Gemma模型时预填充阶段提速1.1~1.25倍，质量几乎不变。
大型语言模型（LLMs）规模庞大，部署时需跨多个加速器设备分摊计算，但设备间通信（如AllReduce）带来显著性能开销。尽管模型量化已被广泛用于降低权重与激活的内存和计算需求且影响较小，但直接对AllReduce等集合操作进行量化因涉及设备间求和，易引发数值不稳定或误差累积。本文提出XLA编译器中针对TPU的原生动态块级高效量化AllReduce方法（EQuARX）。通过使用适合TPU的量化方式及通信与计算深度流水线化，EQuARX在int8精度下相比基线BF16 AllReduce在多种网络拓扑上实现1.8倍加速。此外，该方法使Gemma 3 27B的预填充阶段提速1.25倍，Gemma 3 12B提速1.1倍，质量影响极小。
原文摘要 · Abstract (English)
While Large Language Models (LLMs) have become highly influential, their enormous scale presents significant deployment challenges. Efficiently serving these models typically requires distributing them across numerous accelerator devices, which introduces substantial performance overhead from inter-device communication (collectives). While model quantization has been widely adopted to reduce the memory and compute requirements of LLM weights and activations with minimal quality impact, applying quantization directly to collectives like AllReduce is inherently difficult due to the inter-device summation involved, which can lead to numerical instability or significant error accumulation. In this work, we present a native dynamic block-wise efficient quantized AllReduce within the XLA compiler for TPUs (EQuARX). By using TPU-friendly quantization and deep pipelining of communication and compute, EQuARX with int8 precision achieves a 1.8X speedup over baseline BF16 AllReduce across various network topologies. Furthermore, EQuARX accelerates the prefill stage of Gemma 3 27B by 1.25X and Gemma 3 12B by 1.1X, respectively, with small to negligible impact on quality.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。