EDiT通过分层同步与伪梯度惩罚,提升大模型分布式训练效率与稳定性。
EDiT: A Local-SGD-Based Efficient Distributed Training Method for Large Language Models
- 分层前向同步参数,降低通信与内存开销
- 伪梯度惩罚抑制损失突增,保障训练稳定
- 支持异步训练,适合异构计算集群
分布式训练对大语言模型至关重要。现有方法常受限于通信瓶颈、慢节点及弹性不足,尤其在异构或大规模环境中。局部SGD虽可缓解部分问题,但因额外内存开销和效率稳定性不足,仅适用于小规模训练。为此,我们提出EDiT——一种结合定制化局部SGD与模型分片的高效分布式训练方法。EDiT在前向传播中进行分层参数同步,减少通信与内存开销,并实现重叠;同时引入伪梯度惩罚策略,抑制损失尖峰,确保训练稳定并提升性能。此外,我们提出A-EDiT,一种完全异步的变体,适配异构集群。基于EDiT/A-EDiT,我们开展系列实验验证大模型异步训练可行性,并进行深入分析。结果表明,EDiT/A-EDiT表现优越,为多样化计算生态中的分布式大模型训练提供了鲁棒解决方案。代码已开源:https://github.com/intelligent-machine-learning/atorch/tree/main/atorch/local_sgd。
原文摘要 · Abstract (English)
Distributed training methods are crucial for large language models (LLMs). However, existing distributed training methods often suffer from communication bottlenecks, stragglers, and limited elasticity, particularly in heterogeneous or large-scale environments. Local SGD methods have been proposed to address these issues, but their effectiveness remains limited to small-scale training due to additional memory overhead and lack of concerns on efficiency and stability. To tackle these issues, we propose EDiT, an innovative Efficient Distributed Training method that combines a tailored Local SGD approach with model sharding techniques to enhance large-scale training efficiency. EDiT performs layer-wise parameter synchronization during forward pass, reducing communication and memory overhead and enabling overlap. Besides, EDiT employs a pseudo gradient penalty strategy to suppress loss spikes, which ensures training stability and improves performance. Additionally, we introduce A-EDiT, a fully asynchronous variant of EDiT that accommodates heterogeneous clusters. Building on EDiT/A-EDiT, we conduct a series of experiments to validate large-scale asynchronous training for LLMs, accompanied by comprehensive analyses. Experimental results demonstrate the superior performance of EDiT/A-EDiT, establishing them as robust solutions for distributed LLM training in diverse computational ecosystems. The code is available at Atorch codebase: https://github.com/intelligent-machine-learning/atorch/tree/main/atorch/local_sgd.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。