MoFa: A Unified Performance Modeling Framework for LLM Pretraining

TL;DR

MoFa addresses the challenge of predicting and optimizing large-scale LLM pretraining performance by integrating a unified cost model with fault-tolerance overhead. It introduces CostModel_base, CostModel_Mo, and a fault-tolerance model that uses historical reliability data, enabling accurate end-to-end performance predictions and bottleneck analysis. The framework also includes a MoFa tuning system for joint optimization of parallel strategies, optimization features, and fault tolerance, demonstrated across multiple model scales with high prediction accuracy. The work provides actionable guidance for designing and deploying ultra-large pretraining systems on thousands to tens of thousands of GPUs.

Abstract

The exponential growth in LLM scales, with parameters soaring from billions to trillions, has necessitated distributed pretraining across large clusters comprising thousands to tens of thousands of devices. While hybrid parallelization strategies enable such pretraining, the vast combinatorial strategy space introduces significant optimization challenges. Traditional manual tuning methods incur prohibitive trial-and-error costs, and existing performance modeling approaches exhibit critical limitations: they fail to comprehensively account for prevalent optimization features and ignore the substantial overhead imposed by essential fault tolerance mechanisms like checkpoint recovery in long-duration pretraining. To address these gaps, we propose MoFa, a novel pretraining performance modeling framework that unifies multi-dimensional optimization features and fault tolerance. MoFa incorporates an enhanced cost model to accurately capture the effects of key optimizations and integrates a fault tolerance model based on historical cluster reliability data. Besides, a MoFa-based tuning system is developed to explore optimal pretraining performance and potential bottlenecks in various scenarios. Extensive modeling evaluations demonstrate that MoFa can achieve high prediction accuracy across various scenarios. In addition, through comprehensive tuning experiments, our framework systematically reveals the key factors influencing pretraining performance under different configurations, which provides solid a priori guidance for LLM pretraining system design and deployment.

Figures (10)

• Figure 1: Three-level Strategy for Fault Tolerance.
• Figure 2: Timing diagram of one pretraining step.
• Figure 3: MoFa Tuning System.
• Figure 4: Trend of ETTR with $I_{ckpt}$.
• Figure 5: Trend of ETTR with other factors.