TACO: Efficient Communication Compression of Intermediate Tensors for Scalable Tensor-Parallel LLM Training

TL;DR

TACO is a novel FP8-based compression framework that reduces communication overhead in large-scale tensor-parallel training, significantly improving throughput while maintaining accuracy.

Contribution

We introduce TACO, a robust FP8 compression method with adaptive scaling and fused operators, enabling efficient, high-fidelity tensor communication in large-scale LLM training.

Findings

Achieves up to 1.87X throughput improvement

Maintains near-lossless accuracy in GPT and Qwen models

Effectively reduces memory traffic and kernel overhead

Abstract

Handling communication overhead in large-scale tensor-parallel training remains a critical challenge due to the dense, near-zero distributions of intermediate tensors, which exacerbate errors under frequent communication and introduce significant computational overhead during compression. To this end, we propose TACO (Tensor-parallel Adaptive COmmunication compression), a robust FP8-based framework for compressing TP intermediate tensors. First, we employ a data-driven reshaping strategy combined with an Adaptive Scale-Hadamard Transform to enable high-fidelity FP8 quantization, while its Dual-Scale Quantization mechanism ensures numerical stability throughout training. Second, we design a highly fused compression operator to reduce memory traffic and kernel launch overhead, allowing efficient overlap with communication. Finally, we integrate TACO with existing state-of-the-art methods for Data and Pipeline Parallelism to develop a compression-enabled 3D-parallel training framework. Detailed experiments on GPT models and Qwen model demonstrate up to 1.87X end-to-end throughput improvement while maintaining near-lossless accuracy, validating the effectiveness and efficiency of TACO in large-scale training.