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

2026-04-27

AI Summary

This work addresses the high communication overhead, significant compression errors, and heavy computational burden associated with intermediate tensors in large-scale tensor-parallel training of large language models. To tackle these challenges, the authors propose TACO, a novel framework that integrates data-driven tensor reordering with an adaptive scaled Hadamard transform to enable efficient, high-fidelity FP8 quantization. TACO employs a dual-scale mechanism to ensure numerical stability during training and introduces a highly fused compression operator to reduce memory traffic and kernel launch overhead. Seamlessly integrated into 3D parallel training pipelines, TACO achieves up to 1.87× end-to-end throughput improvement on GPT and Qwen models while maintaining near-lossless model accuracy.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• communication compression
• tensor-parallel training
• intermediate tensors
• large language models
• communication overhead

Innovation

• communication compression
• FP8 quantization
• tensor-parallel training
• adaptive scale-hadamard transform
• 3D-parallelism