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

arXiv: https://arxiv.org/abs/2506.17615 · HTML: https://mlsanity.com/paper/2506.17615/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 32 references that are in the map.
- First posted on arXiv 21 Jun 2025.

## Builds on (most cited first, 25 of 25)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [GPQA: A Graduate-Level Google-Proof Q&A Benchmark](https://mlsanity.com/paper/2311.12022/) (2023)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference](https://mlsanity.com/paper/1712.05877/) (2017)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [What Disease does this Patient Have? A Large-scale Open Domain Question Answering Dataset from Medical Exams](https://mlsanity.com/paper/2009.13081/) (2020)
- [A Survey of Quantization Methods for Efficient Neural Network Inference](https://mlsanity.com/paper/2103.13630/) (2021)
- [AGIEval: A Human-Centric Benchmark for Evaluating Foundation Models](https://mlsanity.com/paper/2304.06364/) (2023)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [Efficiently Scaling Transformer Inference](https://mlsanity.com/paper/2211.05102/) (2022)
- [TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings](https://mlsanity.com/paper/2304.01433/) (2023)
- [Integer Quantization for Deep Learning Inference: Principles and Empirical Evaluation](https://mlsanity.com/paper/2004.09602/) (2020)
- [ZeRO++: Extremely Efficient Collective Communication for Giant Model Training](https://mlsanity.com/paper/2306.10209/) (2023)
- [Evaluating Modern GPU Interconnect: PCIe, NVLink, NV-SLI, NVSwitch and GPUDirect](https://mlsanity.com/paper/1903.04611/) (2019)
- [SDP4Bit: Toward 4-bit Communication Quantization in Sharded Data Parallelism for LLM Training](https://mlsanity.com/paper/2410.15526/) (2024)
- [Quantized Distributed Training of Large Models with Convergence Guarantees](https://mlsanity.com/paper/2302.02390/) (2023)
- [gZCCL: Compression-Accelerated Collective Communication Framework for GPU Clusters](https://mlsanity.com/paper/2308.05199/) (2023)

## Built on by (most cited first, 2 of 2)

- [A Switch-Centric In-Network Architecture for Accelerating LLM Inference in Shared-Memory Network](https://mlsanity.com/paper/2603.28239/) (2026)
- [Echo: Error-Compensated 4-Bit Tensor-Parallel Communication for LLM Training and Inference](https://mlsanity.com/paper/2610.14277/) (2026)

## Abstract (arXiv)

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.