Salvato in:
| Autori principali: | , , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2025
|
| Soggetti: | |
| Accesso online: | https://arxiv.org/abs/2506.17615 |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866909655987388416 |
|---|---|
| author | Ahmed, Ibrahim Schaefer, Clemens Tabak, Gil Vnukov, Denis Zhang, Zenong chern, Felix Yevtushenko, Anatoliy Davis, Andy |
| author_facet | Ahmed, Ibrahim Schaefer, Clemens Tabak, Gil Vnukov, Denis Zhang, Zenong chern, Felix Yevtushenko, Anatoliy Davis, Andy |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2506_17615 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | EQuARX: Efficient Quantized AllReduce in XLA for Distributed Machine Learning Acceleration Ahmed, Ibrahim Schaefer, Clemens Tabak, Gil Vnukov, Denis Zhang, Zenong chern, Felix Yevtushenko, Anatoliy Davis, Andy Machine Learning 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. |
| title | EQuARX: Efficient Quantized AllReduce in XLA for Distributed Machine Learning Acceleration |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2506.17615 |