Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MeCeFO: Enhancing LLM Training Robustness via Fault-Tolerant Optimization

Rizhen Hu, Yutong He, Ran Yan, Mou Sun, Binghang Yuan, Kun Yuan

TL;DR

This work addresses the rising challenge of hardware failures in large-scale distributed LLM training by proposing MeCeFO, a fault-tolerant optimization framework that achieves robustness with minimal overhead. It introduces a neighbor-do-both strategy and three efficiency techniques—MHA backward skip-connections, selective FFN activation recomputation, and a low-rank gradient approximation—to sharing load among neighboring nodes while preserving convergence. The authors prove a convergence rate of $ rac{1}{T+1} ext{E}[ orm{ abla f(m{w}^{(t)})}_2^2] = Oig(1/ ext{sqrt}(nT)ig)$, matching standard distributed SGD, and validate the approach empirically on LLaMA-350M/1B/7B pre-training under frequent failures, showing only a $4.18 ext{%}$ throughput drop and $5.0 imes$ to $6.7 imes$ greater resilience than prior methods. The results demonstrate that carefully orchestrated algorithmic efficiency can substantially improve fault tolerance in practical, large-scale transformer training, with negligible impact on final model quality and downstream tasks.

Abstract

As distributed optimization scales to meet the demands of Large Language Model (LLM) training, hardware failures become increasingly non-negligible. Existing fault-tolerant training methods often introduce significant computational or memory overhead, demanding additional resources. To address this challenge, we propose Memory- and Computation-efficient Fault-tolerant Optimization (MeCeFO), a novel algorithm that ensures robust training with minimal overhead. When a computing node fails, MeCeFO seamlessly transfers its training task to a neighboring node while employing memory- and computation-efficient algorithmic optimizations to minimize the extra workload imposed on the neighboring node handling both tasks. MeCeFO leverages three key algorithmic designs: (i) Skip-connection, which drops the multi-head attention (MHA) module during backpropagation for memory- and computation-efficient approximation; (ii) Recomputation, which reduces activation memory in feedforward networks (FFNs); and (iii) Low-rank gradient approximation, enabling efficient estimation of FFN weight matrix gradients. Theoretically, MeCeFO matches the convergence rate of conventional distributed training, with a rate of $\mathcal{O}(1/\sqrt{nT})$, where n is the data parallelism size and T is the number of iterations. Empirically, MeCeFO maintains robust performance under high failure rates, incurring only a 4.18% drop in throughput, demonstrating 5.0$\times$ to 6.7$\times$ greater resilience than previous SOTA approaches. Codes are available at https://github.com/pkumelon/MeCeFO.

MeCeFO: Enhancing LLM Training Robustness via Fault-Tolerant Optimization

TL;DR

This work addresses the rising challenge of hardware failures in large-scale distributed LLM training by proposing MeCeFO, a fault-tolerant optimization framework that achieves robustness with minimal overhead. It introduces a neighbor-do-both strategy and three efficiency techniques—MHA backward skip-connections, selective FFN activation recomputation, and a low-rank gradient approximation—to sharing load among neighboring nodes while preserving convergence. The authors prove a convergence rate of , matching standard distributed SGD, and validate the approach empirically on LLaMA-350M/1B/7B pre-training under frequent failures, showing only a throughput drop and to greater resilience than prior methods. The results demonstrate that carefully orchestrated algorithmic efficiency can substantially improve fault tolerance in practical, large-scale transformer training, with negligible impact on final model quality and downstream tasks.

Abstract

As distributed optimization scales to meet the demands of Large Language Model (LLM) training, hardware failures become increasingly non-negligible. Existing fault-tolerant training methods often introduce significant computational or memory overhead, demanding additional resources. To address this challenge, we propose Memory- and Computation-efficient Fault-tolerant Optimization (MeCeFO), a novel algorithm that ensures robust training with minimal overhead. When a computing node fails, MeCeFO seamlessly transfers its training task to a neighboring node while employing memory- and computation-efficient algorithmic optimizations to minimize the extra workload imposed on the neighboring node handling both tasks. MeCeFO leverages three key algorithmic designs: (i) Skip-connection, which drops the multi-head attention (MHA) module during backpropagation for memory- and computation-efficient approximation; (ii) Recomputation, which reduces activation memory in feedforward networks (FFNs); and (iii) Low-rank gradient approximation, enabling efficient estimation of FFN weight matrix gradients. Theoretically, MeCeFO matches the convergence rate of conventional distributed training, with a rate of , where n is the data parallelism size and T is the number of iterations. Empirically, MeCeFO maintains robust performance under high failure rates, incurring only a 4.18% drop in throughput, demonstrating 5.0 to 6.7 greater resilience than previous SOTA approaches. Codes are available at https://github.com/pkumelon/MeCeFO.
Paper Structure (22 sections, 5 theorems, 23 equations, 7 figures, 11 tables, 3 algorithms)

This paper contains 22 sections, 5 theorems, 23 equations, 7 figures, 11 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries and related Works
  3. Method
  4. MeCeFO overview
  5. Key technique I: skip-connection
  6. Key technique II: selective activation recomputation
  7. Key technique III: low-rank gradient approximation
  8. Convergence analysis
  9. Experimental results
  10. Experimental setup
  11. Training throughput under failures
  12. Training performance under failures
  13. Conclusions and limitations
  14. Proof of Theorem \ref{['thm:convergence']}
  15. Discussions
  16. ...and 7 more sections

Key Result

Theorem 1

Under Assumptions asp:smooth-asp:grad-err, if momentum parameter $\beta_1\in(1-\delta/(24-12\delta),1)$ and learning rate $\eta\le\min\{1/(2L),$$\sqrt{(\delta(1-\beta_1)^2)/(8L^2)}\}$, MeCeFO (with momentum SGD) converges as where $\Delta:=f(\bm{w}^{(0)})-\inf_{\bm{w}}f(\bm{w})$, and $\Delta_1:=\|\bm{m}^{(0)}-\nabla f(\bm{w}^{(0)})\|_2^2$. (Proof is in Appendix app:pfthm)

Figures (7)

  • Figure 1: Overview of the MeCeFO framework. During both forward (left) and backward (right) propagation, the workload of a failed node is offloaded to a neighboring node within the same data parallel (DP) group.
  • Figure 2: The skip-connection technique in MHA layers. We only skip the MHA's connection in the backward propagation.
  • Figure 3: Ablation of module skipping in LLaMA-130M pre-training.
  • Figure 4: Single-batch relative error of pre-training LLaMA-1B on the C4 dataset.
  • Figure 5: Full-batch relative error of pre-training LLaMA-1B on the C4 dataset.
  • ...and 2 more figures

Theorems & Definitions (8)

  • Theorem 1
  • Corollary 1
  • Lemma 1: Descent lemma
  • proof : Proof of Lemma \ref{['lm:descent']}
  • Lemma 2: Momentum-gradient gap
  • proof : Proof of Lemma \ref{['lm:mom-grad-gap']}
  • Theorem 2: Convergence of MeCeFO
  • proof : Proof of Theorem \ref{['thmres:convergence']}