All is Not Lost: LLM Recovery without Checkpoints

Abstract

Training LLMs on decentralized nodes or on-spot instances, lowers the training cost and enables model democratization. The inevitable challenge here is the transient churns of nodes due to failures and the operator's scheduling policies, leading to losing parts of the model (some layers). The conventional approaches to recover from failures is to either use checkpointing, where periodically a copy of the entire model is sent to an additional storage, or redundant computation. These approaches yield significant communication and/or computation overhead even in non-failure cases and scale poorly in settings with large models. In this paper we propose CheckFree, an efficient recovery method where a failing stage is substituted by weighted averaging of the closest neighboring stages. In contrast to the state of the art, CheckFree requires no additional computation or storage. However, because of the nature of averaging neighbouring stages, it can only recover failures of intermediate stages. We further extend our method to CheckFree+ with out-of-order pipeline execution to tolerate crashes of the first and last stages. Thanks to out-of-order pipelining, behaviour of the first and last stages are mimicked by their neighboring ones, which allows CheckFree+ to recover them by copying the neighboring stages. To recover the (de-)embedding layers, CheckFree+ copies those layers in the neighboring stages, which requires relatively small storage overhead. We extensively evaluate our method on LLaMa models of model sizes from 124M to 1.5B with varying failure frequencies. In the case of low and medium failure rates (5-10%), CheckFree and CheckFree+ outperform both checkpointing and redundant computation in terms of convergence wall-clock time, achieving up to 12% improvement over redundant computation. Both of our proposals can be ran via our code available at: https://github.com/gensyn-ai/CheckFree

Figures (4)

• Figure 1: Three strategies to recover the stage failure using a new node. A stage failure can happen when either each stage is only run by a single node or all the nodes of the same stage in DP fails simultaneously. a) In checkpointing strategy, an external non-faulty storage is used to periodically store checkpoints. Upon a failure, a new node receives stale weights from the storage. b) In redundant computation (RC), each node redundantly stores and performs the computations of the subsequent stage as well bamboo, and in the case of failure, new node can obtain the stage from the previous one. c) In our methods, CheckFree and CheckFree+, we use the weights of the neighbouring stages to initialize the failed stage. CheckFree+ also recovers the first and last stages by copying the corresponding neighbouring stage which is trained to simulate the same behaviour.
• Figure 2: Varying reinitialization strategies for failed stages of a LLaMa 500M model.
• Figure 3: Convergence of the models with 10% failure rate for small and medium models & 16% for large ones.
• Figure 4: Ablation studies of CheckFree+ on medium models.