EQuARX: Efficient Quantized AllReduce in XLA for Distributed Machine Learning Acceleration

📄 arXiv: 2506.17615v1 📥 PDF

作者: Ibrahim Ahmed, Clemens Schaefer, Gil Tabak, Denis Vnukov, Zenong Zhang, Felix chern, Anatoliy Yevtushenko, Andy Davis

分类: cs.LG

发布日期: 2025-06-21

💡 一句话要点

提出EQuARX，在XLA中实现高效量化AllReduce，加速分布式机器学习。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 量化 AllReduce 分布式机器学习 XLA编译器 TPU 大型语言模型 通信优化

📋 核心要点

  1. 大型语言模型部署面临设备间通信开销大的挑战，直接量化AllReduce易导致数值不稳定和误差累积。
  2. EQuARX通过TPU友好的量化和通信计算深度流水线，在XLA编译器中实现高效量化AllReduce。
  3. 实验表明，EQuARX在int8精度下实现了1.8倍的加速，并加速了Gemma模型的预填充阶段。

📝 摘要（中文）

大型语言模型（LLMs）规模庞大，给部署带来了巨大挑战。高效服务这些模型通常需要在多个加速器设备上进行分布式处理，这会因设备间通信（collectives）而产生显著的性能开销。模型量化已被广泛采用，以减少LLM权重和激活的内存和计算需求，同时对质量的影响最小。然而，由于涉及设备间求和，直接将量化应用于AllReduce等collectives本质上是困难的，这可能导致数值不稳定或显著的误差累积。本文提出了一种原生的、动态的、块状的、高效的量化AllReduce方法EQuARX，它在TPU的XLA编译器中实现。通过使用TPU友好的量化以及通信和计算的深度流水线，EQuARX在int8精度下，在各种网络拓扑中实现了相对于基线BF16 AllReduce 1.8倍的加速。此外，EQuARX分别将Gemma 3 27B的预填充阶段加速了1.25倍，将Gemma 3 12B的预填充阶段加速了1.1倍，对质量的影响很小甚至可以忽略不计。

🔬 方法详解

问题定义：论文旨在解决分布式机器学习中，由于大型语言模型规模庞大，在TPU集群上进行AllReduce通信时效率低下的问题。现有的BF16 AllReduce通信开销大，而直接对AllReduce进行量化容易导致数值不稳定和误差累积，影响模型精度。

核心思路：论文的核心思路是设计一种TPU友好的量化方案，并结合通信和计算的深度流水线，在XLA编译器中实现高效的量化AllReduce（EQuARX）。通过动态块状量化，降低通信带宽需求，同时减少量化带来的精度损失。深度流水线设计则可以隐藏通信开销，进一步提升整体性能。

技术框架：EQuARX的技术框架主要包含以下几个阶段：1) 数据分块：将输入数据划分为多个块。2) 动态量化：对每个块进行动态量化，将BF16或FP32数据转换为INT8数据。3) AllReduce通信：使用量化后的数据进行AllReduce通信。4) 反量化：将AllReduce后的数据反量化回原始数据类型。5) 计算：使用反量化后的数据进行后续计算。整个流程在XLA编译器中实现，可以与TPU硬件进行深度集成和优化。

关键创新：EQuARX的关键创新在于：1) 提出了一种TPU友好的动态块状量化方案，该方案能够有效地降低通信带宽需求，同时保持较高的模型精度。2) 设计了一种通信和计算的深度流水线，该流水线能够有效地隐藏通信开销，进一步提升整体性能。3) 将量化AllReduce集成到XLA编译器中，实现了与TPU硬件的深度集成和优化。

关键设计：EQuARX的关键设计包括：1) 动态量化范围的确定：采用动态的方式确定每个块的量化范围，以最大限度地减少量化误差。2) 量化和反量化操作的实现：使用TPU指令集优化量化和反量化操作，以提高计算效率。3) 流水线深度的调整：根据网络拓扑和硬件配置，调整流水线深度，以实现最佳的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，EQuARX在int8精度下，相对于基线BF16 AllReduce，在各种网络拓扑中实现了1.8倍的加速。此外，EQuARX分别将Gemma 3 27B的预填充阶段加速了1.25倍，将Gemma 3 12B的预填充阶段加速了1.1倍，对模型质量的影响很小甚至可以忽略不计。这些结果表明EQuARX是一种高效且实用的量化AllReduce方法。

🎯 应用场景

EQuARX可广泛应用于需要大规模分布式训练和推理的场景，尤其是在资源受限的环境下，例如边缘设备或低功耗服务器。通过降低通信带宽需求和加速计算，EQuARX能够显著提升大型语言模型等应用的性能和效率，降低部署成本，并推动AI技术在更广泛领域的应用。

📄 摘要（原文）

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.