TACO通过高效压缩中间张量，显著提升大模型张量并行训练速度。
TACO: Efficient Communication Compression of Intermediate Tensors for Scalable Tensor-Parallel LLM Training
- 基于FP8的自适应压缩框架，提升量化精度与稳定性。
- 压缩后训练吞吐最高提升1.87倍，几乎无精度损失。
- 适合大规模语言模型训练，尤其适用于多维并行场景。
大规模张量并行训练中的通信开销问题，源于中间张量的密集且接近零的分布特性，在频繁通信下加剧误差，并带来显著压缩计算开销。为此，我们提出TACO（Tensor-parallel Adaptive COmmunication compression），一种基于FP8的鲁棒压缩框架，用于压缩张量并行中的中间张量。首先，采用数据驱动的重排策略结合自适应尺度-哈达玛变换，实现高保真度的FP8量化；其双尺度量化机制确保训练全过程数值稳定。其次，设计高度融合的压缩算子，减少内存流量和核启动开销，支持与通信高效重叠。最后，将TACO集成至现有先进数据并行与流水线并行方法中，构建支持压缩的三维并行训练框架。在GPT与Qwen模型上的实验表明，端到端吞吐最高提升1.87倍，同时保持近乎无损的精度，验证了TACO在大规模训练中的有效性与高效性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。