针对大模型训练梯度通信瓶颈，提出自适应压缩方法提升效率。
TAGC: Optimizing Gradient Communication in Distributed Transformer Training
- 基于Transformer结构设计分层动态压缩策略
- 相比标准FSDP提速最高达15%，模型质量几乎不变
- 已开源集成至PyTorch FSDP，适合大规模模型训练者
大型语言模型（LLMs）的复杂性日益增加，亟需高效训练策略以缓解分布式训练带来的高计算成本。其中，多个GPU间梯度同步是主要瓶颈，尤其在零冗余并行模式下更为突出。本文提出面向Transformer的梯度压缩算法（TAGC），在无损同态压缩基础上，针对分片模型进行适配，并引入层选择性压缩与动态稀疏化等优化。实验表明，相较于标准全分片数据并行（FSDP）方法，TAGC可使训练加速最高达15%，对模型质量影响极小。我们已将TAGC集成至PyTorch FSDP框架中，代码公开于https://github.com/ipolyakov/TAGC。
原文摘要 · Abstract (English)
The increasing complexity of large language models (LLMs) necessitates efficient training strategies to mitigate the high computational costs associated with distributed training. A significant bottleneck in this process is gradient synchronization across multiple GPUs, particularly in the zero-redundancy parallelism mode. In this paper, we introduce Transformer-Aware Gradient Compression (TAGC), an optimized gradient compression algorithm designed specifically for transformer-based models. TAGC extends the lossless homomorphic compression method by adapting it for sharded models and incorporating transformer-specific optimizations, such as layer-selective compression and dynamic sparsification. Our experimental results demonstrate that TAGC accelerates training by up to 15% compared to the standard Fully Sharded Data Parallel (FSDP) approach, with minimal impact on model quality. We integrate TAGC into the PyTorch FSDP framework, the implementation is publicly available at https://github.com/ipolyakov/TAGC.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。