💻 computer science

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

TACO is a robust FP8-based framework that employs adaptive quantization and a fused compression operator to efficiently compress intermediate tensors in tensor-parallel LLM training, achieving up to 1.87x throughput improvement while maintaining near-lossless accuracy.

Original authors: Man Liu, Xingchen Liu, Xingjian Tian, Bing Lu, Shengkay Lyu, Shengquan Yin, Wenjing Huang, Zheng Wei, Hairui Zhao, Guangming Tan, Dingwen Tao

Published 2026-04-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Man Liu, Xingchen Liu, Xingjian Tian, Bing Lu, Shengkay Lyu, Shengquan Yin, Wenjing Huang, Zheng Wei, Hairui Zhao, Guangming Tan, Dingwen Tao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a giant team of robots (a Large Language Model) to write a story. To do this, you split the work among hundreds of robots working in parallel. This is called Tensor Parallelism.

However, there's a major problem: these robots need to constantly whisper their progress to each other. They pass notes back and forth thousands of times a second. The problem is, these notes are huge, and the hallway they use to pass them (the network) is narrow and slow. The robots spend more time waiting for notes than actually writing the story.

TACO is a new system designed to shrink these notes so the robots can talk faster without losing the meaning of the story.

Here is how TACO works, explained through simple analogies:

1. The Problem: The "Zero" Crowd

The notes the robots pass (called intermediate tensors) have a weird shape. Most of the numbers in them are tiny, tiny numbers clustered right around zero. A few are huge, but they are rare.

The Old Way (INT8): Imagine trying to describe a crowd where 99% of people are standing still in a tiny circle, and 1% are running far away. If you use a standard ruler (INT8) that has equal spacing for every inch, you can't measure the tiny circle accurately. Everyone in the circle gets squished into the same spot, and the information is lost. The robots get confused, and the training fails.
The TACO Solution (FP8): TACO uses a special ruler (FP8) that has very fine, tiny markings near zero and wider markings far away. This is perfect for the "zero crowd." But even this ruler has limits.

2. The Magic Trick: The "Adaptive Shuffle" (ASH Transform)

Even with the special FP8 ruler, the notes are still too crowded near zero. TACO performs a clever magic trick called the Adaptive Scale–Hadamard Transform.

The Analogy: Imagine a room full of people huddled tightly in one corner (the zero values). If you just ask them to stand in a line, they are still too close together to count accurately.
What TACO does: It first measures how "energetic" each small group of people is. Then, it gently pushes the quiet groups (low values) apart to make them easier to see, while pulling the loud groups (high values) closer together so they don't run off the edge of the room.
The Result: The crowd is now spread out perfectly across the room, fitting exactly into the "sweet spot" of the FP8 ruler. This prevents the "Zero-Collapse" where information gets lost.

3. The Safety Net: "Dual-Scale" (DS)

Sometimes, even after shuffling, a few numbers might get too big for the FP8 ruler, or too small.

The Analogy: Imagine you are packing a suitcase. You have a main scale for the heavy clothes, but you also need a tiny scale for the jewelry.
What TACO does: It uses two scales at once. One scale adjusts the whole group to fit the suitcase (preventing overflow), and the other scale ensures the tiny jewelry (the small values) doesn't get crushed. This keeps the training stable and prevents the robots from getting "divergent" (going off the rails).

4. The Speed Boost: "The Fusion Kitchen"

Usually, to shrink these notes, the computer has to do three separate steps: measure the crowd, shuffle them, and then pack them. This is like a chef chopping, then stirring, then plating, but having to walk to the fridge between every step. It's slow.

TACO's Innovation: TACO builds a "fused kitchen." It combines chopping, stirring, and plating into one single, super-fast motion. The computer does all three steps at once without stopping to write data to memory. This makes the process incredibly fast and allows the robots to talk while they are still thinking (overlapping communication with computation).

The Results

The paper tested TACO on giant models (like GPT and Qwen) and found:

Speed: It made the training 1.87 times faster in some cases.
Accuracy: Despite shrinking the data so much, the robots learned just as well as if they had sent the full, uncompressed notes. The "loss" (error) was almost non-existent.
Scalability: It works great even when you have 8, 16, or more robots working together.

In short: TACO is a smart, fast way to shrink the massive data traffic between AI training robots. It uses a special "reshuffling" trick to make the data fit perfectly into a smaller format, ensuring the AI learns quickly without losing its mind.

1. Problem Statement

Large Language Model (LLM) training at scale relies heavily on 3D Parallelism (Data Parallelism, Pipeline Parallelism, and Tensor Parallelism). While Data Parallelism (DP) and Pipeline Parallelism (PP) have seen successful communication compression techniques, Tensor Parallelism (TP) remains a critical bottleneck.

Communication Overhead: TP requires frequent, tightly synchronized communication of intermediate tensors (activations and gradients) at every layer and iteration. In large-scale training, TP communication accounts for 50–60% of total communication time and lies on the critical execution path, making it difficult to overlap with computation.
The Compression Challenge:
- Error Accumulation: Unlike DP/PP, TP involves high-frequency synchronization. Minor quantization errors in intermediate tensors propagate through attention and residual connections, amplifying during backpropagation and causing training divergence.
- Data Distribution: TP intermediate tensors exhibit a dense, near-zero distribution with a long tail. Standard quantization methods (like INT8) fail here because they use uniform spacing, leading to massive information loss for the small-magnitude values that dominate the tensor.
- Inefficiency of Existing Methods: Prior inference-oriented compression (e.g., SmoothQuant) or DP-focused methods (e.g., SDP4bit) cannot be directly applied to TP training without catastrophic failure due to numerical instability.

2. Methodology: TACO Framework

The authors propose TACO (Tensor-parallel Adaptive COmmunication compression), a robust FP8-based framework designed specifically for TP intermediate tensors.

A. Why FP8?

The paper analyzes the statistical properties of TP tensors and concludes that FP8 is superior to INT8:

INT8 Limitations: Uniform quantization causes "collisions" where distinct small values map to the same integer, leading to high relative error in the dense zero-centered region.
FP8 Advantage: FP8 (specifically E4M3 format) uses an exponential scale, providing high precision for small values (near zero) and a wide dynamic range for outliers. This matches the statistical distribution of TP tensors perfectly.

B. Core Techniques

TACO introduces three key innovations to make FP8 compression viable for TP training:

Adaptive Scale–Hadamard (ASH) Transform:
- Problem: Standard Hadamard transforms are isometric (preserve energy) and cannot disperse the dense zero-clusters of TP tensors effectively; they remain trapped near zero, underutilizing FP8's precision.
- Solution: TACO employs a block-wise adaptive rescaling before the Hadamard transform. It calculates the second-order raw moment (energy) of each block and applies a scaling factor to normalize the energy to a target level.
- Effect: This "unfolds" the dense zero-clusters, spreading them across the FP8 representable range, followed by an orthogonal Hadamard rotation to create a Gaussian-like distribution suitable for quantization.
Dual-Scale (DS) Quantization:
- Problem: Even with ASH, extreme values can cause overflow, or small values can underflow, leading to instability.
- Solution: TACO uses two distinct scalars per block:
  - Pre-rotation scale ( $\alpha_k$ ): Controls the energy of the block before transformation.
  - Post-rotation scale ( $s_k$ ): Ensures the transformed values fit strictly within the FP8 dynamic range ( $Q_{max}$ ).
- Reconstruction: During decompression, these two scales are used to reverse the process, ensuring high-fidelity reconstruction and preventing error amplification.
System-Level Kernel Fusion & Overlap:
- Fused Operator: To minimize the computational overhead of compression (which is critical since TP is on the critical path), TACO fuses variance computation, ASH transform, and DS quantization into a single GPU kernel. This eliminates redundant global memory accesses and kernel launch overhead.
- Communication-Computation Overlap: TACO integrates with the COCCL (Compressible Collective Communication Library) and uses a "two-shot" AllReduce strategy (ReduceScatter + AllGather). It overlaps the compression/decompression kernels with communication, effectively hiding latency.

3. Key Contributions

Systematic Analysis: Identified the "Zero-Collapse" phenomenon in TP tensors, proving that standard INT8 and fixed transforms are inadequate, and establishing FP8 as the optimal format.
TACO Framework: The first 3D-parallel communication compression system to achieve stable convergence in large-scale TP training. It introduces the ASH transform and Dual-Scale quantization to handle the unique distribution of TP tensors.
High-Performance Implementation: Developed a highly fused compression operator that reduces memory traffic and enables efficient overlap with communication, mitigating the overhead of frequent compression.
Full 3D Integration: Successfully integrated TACO with state-of-the-art DP (SDP4bit) and PP (TahQuant) compressors to create a fully compression-enabled 3D parallel training system.

4. Experimental Results

The authors evaluated TACO on GPT models (2.7B, 6.7B, 13B) and Qwen2.5-7B using NVIDIA H100 GPUs.

Accuracy (Convergence):
- TACO achieves near-lossless accuracy. On GPT-350M, the validation loss degradation is only +0.25% compared to the BF16 baseline, whereas existing methods (TahQuant) degrade by ~2.9%.
- On larger models (GPT-6.7B, Qwen-2.5-7B), the degradation remains negligible (<0.7%), proving stability across scales.
- Ablation studies confirm that removing ASH or DS leads to catastrophic divergence or significant accuracy loss.
Performance (Throughput):
- TP Compression: TACO achieves up to 1.87× end-to-end throughput improvement for TP communication compared to the uncompressed baseline.
- Full 3D Parallelism: When applied to a full 3D parallel setup (TP=4, PP=2, DP=2), TACO delivers up to 1.53× speedup over the baseline.
- Scalability: The performance gains increase with the degree of Tensor Parallelism (TP), as TP communication becomes the dominant bottleneck.

5. Significance

Solving the TP Bottleneck: TACO addresses the most difficult aspect of distributed LLM training—compressing the high-frequency, numerically sensitive intermediate tensors in Tensor Parallelism without sacrificing model quality.
Enabling Larger Models: By reducing communication overhead by nearly 50% (via FP8) while maintaining stability, TACO allows for training larger models on existing hardware clusters or training existing models significantly faster.
Hardware Agnostic: While optimized for FP8 (supported by modern GPUs like H100), the framework is designed to degrade gracefully to INT8 on hardware lacking native FP8 support, ensuring broad applicability.
Future Impact: This work paves the way for exascale training by demonstrating that aggressive communication compression is feasible even on the critical path of training, provided the compression is mathematically tailored to the data distribution.