Extensive system scales (i.e. thousands of GPU/TPUs) and prolonged training periods (i.e. months of pretraining) significantly escalate the probability of failures when training large language models (LLMs). Thus, efficient and reliable fault-tolerance methods are in urgent need. Checkpointing is the primary fault-tolerance method to periodically save parameter snapshots from GPU memory to disks via CPU memory. In this paper, we identify the frequency of existing checkpoint-based fault-tolerance being significantly limited by the storage I/O overheads, which results in hefty re-training costs on restarting from the nearest checkpoint. In response to this gap, we introduce an in-memory fault-tolerance framework for large-scale LLM pretraining. The framework boosts the efficiency and reliability of fault tolerance from three aspects: (1) Reduced Data Transfer and I/O: By asynchronously caching parameters, i.e., sharded model parameters, optimizer states, and RNG states, to CPU volatile memory, Our framework significantly reduces communication costs and bypasses checkpoint I/O. (2) Enhanced System Reliability: Our framework enhances parameter protection with a two-layer hierarchy: snapshot management processes (SMPs) safeguard against software failures, together with Erasure Coding (EC) protecting against node failures. This double-layered protection greatly improves the survival probability of the parameters compared to existing checkpointing methods. (3) Improved Snapshotting Frequency: Our framework achieves more frequent snapshotting compared with asynchronous checkpointing optimizations under the same saving time budget, which improves the fault tolerance efficiency. Empirical results demonstrate that Our framework minimizes the overhead of fault tolerance of LLM pretraining by effectively leveraging redundant CPU resources.
翻译:大规模系统规模(即数千GPU/TPU)和长期训练周期(即数月的预训练)显著增加了大语言模型训练过程中发生故障的概率。因此,亟需高效可靠的容错方法。检查点保存作为主要的容错手段,通过CPU内存定期将参数快照从GPU内存持久化至磁盘。本文指出,现有基于检查点的容错机制其执行频率受限于存储I/O开销,导致从最近检查点重启时产生高昂的重新训练代价。针对这一不足,我们提出面向大规模LLM预训练的内存容错框架。该框架从三个方面提升容错效率与可靠性:(1)减少数据传输与I/O:通过将分片模型参数、优化器状态及随机数生成器状态异步缓存至CPU易失性内存,显著降低通信开销并绕过检查点I/O;(2)增强系统可靠性:采用双层参数保护架构——快照管理进程抵御软件故障,结合纠删码应对节点故障。相较于现有检查点方法,这种双层保护极大提升参数的生存概率;(3)提升快照频率:在相同保存时间预算下,本框架相较于异步检查点优化方法实现更高频率的快照,从而提高容错效率。实验结果表明,通过有效利用冗余CPU资源,本框架将LLM预训练的容错开销降至最低。