[论文解读] Quantized Distributed Training of Large Models with Convergence Guarantees

Ilia Markov, Adrian Vladu+2 more|arXiv (Cornell University)|Feb 5, 2023

Advanced Neural Network Applications被引用 5

一句话总结

本文提出QSDP，一种通信高效的全分片数据并行（FSDP）训练变体，可在理论收敛保证下对梯度和模型权重实现完整量化。通过在全部到全部通信过程中对权重和梯度应用无偏量化，QSDP消除了FSDP的通信瓶颈，同时保持模型精度，在参数量高达13亿的GPT模型上实现了高达2.2倍的端到端加速。

ABSTRACT

Communication-reduction techniques are a popular way to improve scalability in data-parallel training of deep neural networks (DNNs). The recent emergence of large language models such as GPT has created the need for new approaches to exploit data-parallelism. Among these, fully-sharded data parallel (FSDP) training is highly popular, yet it still encounters scalability bottlenecks. One reason is that applying compression techniques to FSDP is challenging: as the vast majority of the communication involves the model's weights, direct compression alters convergence and leads to accuracy loss. We present QSDP, a variant of FSDP which supports both gradient and weight quantization with theoretical guarantees, is simple to implement and has essentially no overheads. To derive QSDP we prove that a natural modification of SGD achieves convergence even when we only maintain quantized weights, and thus the domain over which we train consists of quantized points and is, therefore, highly non-convex. We validate this approach by training GPT-family models with up to 1.3 billion parameters on a multi-node cluster. Experiments show that QSDP preserves model accuracy, while completely removing the communication bottlenecks of FSDP, providing end-to-end speedups of up to 2.2x.

研究动机与目标

• 为解决大规模语言模型FSDP训练中的通信瓶颈问题，该瓶颈源于频繁的全部到全部权重交换。
• 在不损害收敛性或模型精度的前提下，实现在FSDP中对模型权重和梯度的完整量化。
• 为在非凸量化点集合上进行训练时，SGD在完整量化模型状态下的收敛性提供理论保证。
• 设计一种实用且低开销的实现方案，可无缝集成到PyTorch等现有深度学习框架中。

提出的方法

• 提出QSDP，一种修改后的FSDP变体，在每次全部到全部通信步骤前对梯度和权重应用量化。
• 对权重使用无偏量化算子，定义为最近量化级别的随机舍入，确保期望下误差为零。
• 对梯度量化应用标准的无偏压缩技术（如Alistarh等人，2017年提出的方案）。
• 将训练过程建模为带有量化作为投影机制的稀疏恢复问题，从而支持理论收敛性分析。
• 通过证明在标准光滑性和有界梯度假设下，QSDP可收敛至量化格点上损失函数的极小值点，推导出收敛性保证。
• 在PyTorch中高效实现QSDP，内存和计算开销极低，无需额外节点内存或复杂的误差校正机制。

实验结果

研究问题

• RQ1 是否可以在不降低模型收敛性或精度的前提下，对FSDP中的模型权重和梯度实现完整量化？
• RQ2 尽管使用有偏梯度估计器，是否仍可理论保证在非凸量化权重格点上训练的收敛性？
• RQ3 是否可通过量化完全消除FSDP中的通信瓶颈，同时保持端到端训练速度提升？
• RQ4 QSDP在大规模语言模型（如参数量高达13亿的GPT）上的实际表现如何？
• RQ5 在FSDP分布式训练中，量化精度与收敛稳定性之间的权衡关系如何？

主要发现

• 在多节点集群上训练GPT系列模型（参数量最高达13亿）时，QSDP相比标准FSDP实现了高达2.2倍的端到端加速。
• 该方法在所有评估模型中均保持了模型精度，包括13亿参数的GPT模型，性能无明显下降。
• 标准FSDP中的通信瓶颈被QSDP完全消除，因为量化显著降低了带宽需求。
• 理论分析证明，在标准光滑性和有界梯度假设下，QSDP可收敛至量化格点上损失函数的极小值点。
• 权重量化通过无偏随机舍入算子实现，确保量化权重的期望值等于原始权重。
• 实现几乎不引入额外计算或内存开销，使其适用于生产规模训练。