EQuARX：用于分布式机器学习加速的 XLA 高效量化 AllReduce
机器学习
2025-06-24 v1
摘要
虽然大语言模型（LLMs）已变得极具影响力，但其庞大的规模带来了重大的部署挑战。高效地服务这些模型通常需要将它们分布在众多加速器设备上，这引入了由于设备间通信（集合通信）带来的显著性能开销。虽然模型量化已被广泛采用，以在质量影响最小的情况下减少 LLM 权重和激活的内存与计算需求，但将量化直接应用于 AllReduce 等集合通信操作本质上是困难的，因为其中涉及的设备间求和可能导致数值不稳定或显著的误差累积。在这项工作中，我们提出了在用于 TPU 的 XLA 编译器内实现的原生动态分块高效量化 AllReduce（EQuARX）。通过使用对 TPU 友好的量化以及通信与计算的深度流水线化，EQuARX 在 int8 精度下在各种网络拓扑上比基线 BF16 AllReduce 实现了 1.8 倍的加速。此外，EQuARX 将 Gemma 3 27B 的预填充阶段加速了 1.25 倍，将 Gemma 3 12B 加速了 1.1 倍，而对质量的影响微乎其微。
引用
@article{arxiv.2506.17615,
title = {EQuARX: Efficient Quantized AllReduce in XLA for Distributed Machine Learning Acceleration},
author = {Ibrahim Ahmed and Clemens Schaefer and Gil Tabak and Denis Vnukov and Zenong Zhang and Felix chern and Anatoliy Yevtushenko and Andy Davis},
journal= {arXiv preprint arXiv:2506.17615},
year = {2025}
}