arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2408.04307cs.DCcs.LG
收藏

MoC-System:面向稀疏专家混合模型训练的高效容错机制

MoC-System: Efficient Fault Tolerance for Sparse Mixture-of-Experts Model Training

The Hong Kong University of Science and Technology (Guangzhou)(香港科技大学(广州))

Weilin Cai, Le Qin, Jiayi Huang

AI总结:

本文提出 MoC-System,通过部分专家检查点机制和两级异步检查点管理,在保持模型精度的前提下,将稀疏 MoE 模型训练中的检查点开销最高降低 98.9%。

摘要:

随着大语言模型持续扩展,分布式训练系统已突破一万个节点,容错的重要性日益凸显。检查点已成为主流的容错策略,大量研究致力于优化其效率。然而,稀疏专家混合(Mixture-of-Experts,MoE)模型的出现带来了新的挑战,因为尽管其计算需求与稠密模型相当,但模型规模却大幅增加。在本文中,我们提出了混合检查点系统(Mixture-of-Checkpoint System,MoC-System),用于协调分布式训练系统中产生的大量检查点分片。MoC-System 具有一种新颖的部分专家检查点(Partial Experts Checkpointing,PEC)机制,这是一种算法-系统协同设计,策略性地保存选定的专家子集,从而有效地将 MoE 检查点大小降低到与稠密模型相当的水平。结合混合并行策略,MoC-System 采用完全分片检查点策略,将工作负载均匀分配到各个分布式 rank 上。此外,MoC-System 引入了一种两级检查点管理方法,异步处理内存快照和持久化过程。我们在 Megatron-DeepSpeed 框架上构建了 MoC-System,在使用 ZeRO-2 数据并行和专家并行进行 MoE 模型训练时,与原始方法相比,每次检查点过程的开销最多降低了 98.9%。此外,大量的实证分析证实,我们的方法在提高效率的同时保持了相当的模型精度,甚至在下游任务上实现了平均 1.08% 的精度提升。

英文摘要:

As large language models continue to scale up, distributed training systems have expanded beyond 10k nodes, intensifying the importance of fault tolerance. Checkpoint has emerged as the predominant fault tolerance strategy, with extensive studies dedicated to optimizing its efficiency. However, the advent of the sparse Mixture-of-Experts (MoE) model presents new challenges due to the substantial increase in model size, despite comparable computational demands to dense models. In this work, we propose the Mixture-of-Checkpoint System (MoC-System) to orchestrate the vast array of checkpoint shards produced in distributed training systems. MoC-System features a novel Partial Experts Checkpointing (PEC) mechanism, an algorithm-system co-design that strategically saves a selected subset of experts, effectively reducing the MoE checkpoint size to levels comparable with dense models. Incorporating hybrid parallel strategies, MoC-System involves fully sharded checkpointing strategies to evenly distribute the workload across distributed ranks. Furthermore, MoC-System introduces a two-level checkpointing management method that asynchronously handles in-memory snapshots and persistence processes. We build MoC-System upon the Megatron-DeepSpeed framework, achieving up to a 98.9% reduction in overhead for each checkpointing process compared to the original method, during MoE model training with ZeRO-2 data parallelism and expert parallelism. Additionally, extensive empirical analyses substantiate that our methods enhance efficiency while maintaining comparable model accuracy, even achieving an average accuracy increase of 1.08% on downstream tasks.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑