3-4比特量化让大模型流水线训练快4.3倍,还能稳定收敛。
TAH-QUANT: Effective Activation Quantization in Pipeline Parallelism over Slow Network
- 按小通道块分块量化,动态分配比特数
- 3-4比特下训练速度比浮点快4.3倍,比AQ-SGD快1.33倍
- 适合分布式大模型训练,尤其网络慢的场景
大规模语言模型的去中心化训练可汇聚地理分布的计算资源,但在流水线并行设置下常受网络通信瓶颈制约。流水线并行将模型层分置于不同设备以处理超大规模模型,但需频繁传输中间激活值,在带宽受限时尤为困难。为此,我们提出TAH-Quant(逐块自适应哈达玛量化):一种面向流水线并行的新型激活量化框架。TAH-Quant结合细粒度逐块量化、基于熵的逐块自适应比特分配,以及基于哈达玛变换与枢纽交换的异常值抑制机制。相比词元级分配,逐块分配在每个词元内的小通道窗口中进行精度分配,相同比特预算下降低量化误差。我们证明,采用TAH-Quant的流水线并行训练保持O(1/sqrt(T))的收敛速率,与原始随机梯度下降一致。大量实验表明,TAH-Quant实现3-4比特的激活动量压缩,在不损失训练收敛性的同时,较未压缩的FP32提升最高4.3倍吞吐量,较AQ-SGD提升最高1.33倍实际运行速度,并避免了其激活缓存开销,可在多种训练场景下良好泛化。
原文摘要 · Abstract (English)
Decentralized training of large language models offers the opportunity to pool computational resources across geographically distributed participants, but is often bottlenecked by network communication, particularly under pipeline parallel settings. While pipeline parallelism partitions model layers across devices to handle large-scale models, it necessitates frequent communication of intermediate activations, creating challenges when network bandwidth is limited. To address these issues, we propose TAH-Quant (Tile-wise Adaptive Hadamard Quantization), a novel activation quantization framework for pipeline parallelism. TAH-Quant integrates fine-grained tile-wise quantization, entropy-guided tile-wise adaptive bit allocation for optimal bit usage, and a Hadamard-based transformation with pivot swapping to effectively suppress outliers. Compared with token-level allocation, the tile-wise allocator assigns precision at the granularity of small channel windows within each token, reducing quantization error under the same bit budget. We prove that pipeline parallel training equipped with TAH-Quant maintains a convergence rate of O(1/sqrt(T)), matching that of vanilla stochastic gradient descent. Extensive experiments demonstrate that TAH-Quant achieves an aggressive activation quantization ratio of 3-4 bits, providing up to 4.3x throughput speedup over uncompressed FP32 and up to 1.33x wall-clock speedup over AQ-SGD, while preserving training convergence, avoiding AQ-SGD's activation-cache overhead, and generalizing well across various training scenarios.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。