Communication-Efficient Multi-Device Inference Acceleration for Transformer Models

TL;DR

ASTRA is a communication-efficient multi-device inference framework for Transformer models that significantly reduces latency and bandwidth requirements while maintaining accuracy, enabling faster AI applications in bandwidth-limited environments.

Contribution

ASTRA introduces a novel integration of sequence parallelism and mixed-precision attention with compression techniques to enable efficient multi-device Transformer inference.

Findings

Achieves up to 2.64X speedup over single-device inference.

Achieves up to 15.25X speedup over existing multi-device methods.

Operates effectively under bandwidths as low as 10 Mbps.

Abstract

Transformer models power many AI applications but suffer from high inference latency, limiting their use in real-time settings. Multi-device inference can reduce latency by parallelizing computation. Yet, existing methods require high inter-device bandwidth, making them impractical for bandwidth-constrained environments. We propose ASTRA, a communication-efficient framework that accelerates Transformer inference through a novel integration of sequence parallelism and a Mixed-Precision Attention mechanism designed to minimize inter-device communication. ASTRA compresses non-local token embeddings via vector quantization and preserves task accuracy through two optimizations, Noise-Augmented Quantization and Distributed Class Tokens. Experiments on ViT and GPT2 across vision and NLP tasks show that ASTRA achieves up to 2.64X speedups over single-device inference and up to 15.25X speedups over state-of-the-art multi-device inferences, while operating under bandwidths as low as 10 Mbps. ASTRA is open-sourced at https://github.com/xl1990/Astra.