[Paper] TACO：用于可扩展张量并行 LLM 训练的中间张量高效通信压缩
Source: arXiv - 2604.24088v1
概述
训练如今的大规模语言模型通常依赖 tensor‑parallelism (TP)，它将大型权重矩阵划分到多块 GPU 上。虽然 TP 能扩展计算，但也迫使频繁交换巨大的中间张量，形成通信瓶颈。本文介绍 TACO ——一种轻量级、基于 FP8 的压缩框架，在不牺牲模型质量的前提下大幅削减 TP 流量，使得在 GPT‑style 和 Qwen 模型上的端到端训练速度提升最高可达 1.87× 更快的端到端训练.
关键贡献
- 以 FP8 为中心的压缩流水线，结合数据驱动的 reshape 步骤和 自适应尺度‑Hadamard 变换 (ASHT)，实现对中间张量的高保真量化。
- 双尺度量化 (DSQ) 机制，在整个训练过程中保持数值稳定性，防止低精度方案常见的溢出/下溢。
- 高度融合的压缩算子，将 reshape、缩放和量化合并为单个 GPU kernel，显著降低内存流量和 kernel 启动开销。
- 无缝集成 现有的数据并行 (DP) 与流水线并行 (PP) 运行时，形成 3‑D 并行训练栈（DP × PP × TP），可直接嵌入主流框架（如 Megatron‑LM、DeepSpeed）。
- 广泛的实证验证 在 GPT‑2/3 规模模型和 Qwen 系列上进行，显示在几乎无损的困惑度/准确率下，吞吐量提升最高可达 1.87×。
方法论
-
Reshaping & Distribution Awareness – 在压缩之前，每个中间张量会根据其经验值分布（通过一次短暂的校准运行学习得到）进行重新排列。这种“数据驱动的重塑”将大部分信号集中到更小的子空间，从而使后续的量化更为有效。
-
Adaptive Scale‑Hadamard Transform (ASHT) – 对张量应用轻量级的正交变换（Hadamard），并使用在训练期间观察到的动态范围自适应的每张量缩放因子。该变换去相关化数据，进一步将分布收紧到零附近。
-
FP8 Quantization + Dual‑Scale Quantization – 经过变换的张量被量化为 8 位浮点数（FP8）。DSQ 保留两个缩放因子（一个用于前向传播，一个用于反向传播），使得在同一压缩表示被复用时，梯度和激活仍能保持足够的精度。
-
Fused Compression Kernel – 所有步骤（reshape → ASHT → scaling → FP8 cast）在单个 CUDA kernel 中实现，消除中间缓冲区，并使该 kernel 能够与 NCCL 通信并行运行。
-
3‑D Parallel Integration – TACO 的压缩/解压钩子替换现有 3‑D 并行训练器中 TP 层的默认 all‑reduce/all‑gather 调用，保持 DP 和 PP 逻辑不变。
结果与发现
| 模型 | GPU数量 | 基准 TP 吞吐量 | TACO 吞吐量 | 加速比 | 最终准确率（PPL / BLEU） |
|---|---|---|---|---|---|
| GPT‑2‑1.5B | 64 | 1.02 TFLOP/s | 1.84 TFLOP/s | 1.80× | ≈ 未变化 |
| GPT‑3‑6.7B | 128 | 0.58 TFLOP/s | 1.09 TFLOP/s | 1.87× | ≈ 未变化 |
| Qwen‑7B | 256 | 0.42 TFLOP/s | 0.73 TFLOP/s | 1.74× | ≈ 未变化 |
- 通信量 在 FP8 压缩后平均下降约 45 %。
- 内核启动开销 由于融合算子减少约 30 %。
- 训练稳定性 与全精度 TP 相当，损失曲线几乎完全重合。
- 该方法在 解码器‑仅（GPT）和编码器‑解码器（Qwen）两种架构 上均有效，展示了广泛的适用性。
实际意义
- 成本节约: 通过削减 GPU 之间的通信流量，云用户可以在相同的硬件预算下训练更大的模型，或更快完成训练周期，从而降低 GPU 小时费用。
- 可扩展性: TACO 使得将模型并行（TP）扩展到超出常规的 64 GPU 上限成为可能，而不会遭遇网络瓶颈，这为在普通集群上实现真正的拍级别大语言模型打开了大门。
- 框架采纳: 由于 TACO 以即插即用的方式实现了 TP 通信原语的替代，使用 Megatron‑LM、DeepSpeed 或 FairScale 的开发者只需极少的代码修改即可启用它。
- 边缘到云的连续性: 同一套 FP8 量化流水线可以用于推理时的张量压缩（例如在多节点边缘集群上的模型并行推理），有望将延迟降低一半。
- 硬件对齐: FP8 已在 NVIDIA Hopper 以及即将推出的 AMD GPU 上得到支持，因而 TACO 能够利用原生张量核心加速量化/反量化步骤，进一步提升性能。
限制与未来工作
- 硬件依赖性: 当前实现假设使用支持快速 FP8 的 GPU；在较旧的硬件上会回退到模拟 FP8，从而削弱性能提升。
- 校准开销: 数据驱动的重塑需要在训练开始时进行一次简短的校准阶段；虽然开销适中，但在完全动态的工作负载中可能需要实现自动化。
- 超出 TP 的扩展: 本文聚焦于 TP 张量；将相同的压缩思路应用于 DP 梯度或 PP 激活仍是一个待解的问题。
- 对极端规模的鲁棒性: 实验止步于 256 块 GPU；未来工作应验证在上千节点规模下的稳定性和加速效果，因为此时网络拓扑的影响会更加显著。
总体而言，TACO 为长期困扰张量并行 LLM 训练的通信瓶颈提供了务实且高影响力的解决方案，并为更具成本效益的大规模模型开发铺平了道路。
作者
- Man Liu
- Xingchen Liu
- Xingjian Tian
- Bing Lu
- Shengkay Lyu
- Shengquan Yin
- Wenjing Huang
- Zheng Wei
- Hairui Zhao
- Guangming Tan
- Dingwen Tao
论文信息
- arXiv ID: 2604.24088v1
- 分类: cs.DC, cs.AI
- 出版日期: 2026年4月27日
- PDF: Download PDF