Malleus: Straggler-Resilient Hybrid Parallel Training of Large-scale Models via Malleable Data and Model Parallelization

TL;DR

Malleus is a novel framework that enhances large-scale model training efficiency by dynamically detecting and mitigating GPU stragglers through adaptive parallelization and seamless model state migration.

Contribution

It introduces a per-GPU straggler quantification and a re-planning algorithm for adaptive hybrid parallel training, improving resilience and efficiency.

Findings

Achieves 2.63-5.28x efficiency gains over existing methods.

Effectively adapts to dynamic straggler situations during training.

Supports large language models up to 110B parameters.

Abstract

As the scale of models and training data continues to grow, there is an expanding reliance on more GPUs to train large-scale models, which inevitably increases the likelihood of encountering dynamic stragglers that some devices lag behind in performance occasionally. However, hybrid parallel training, one of the de facto paradigms to train large models, is typically sensitive to the stragglers. This paper presents Malleus, a straggler-resilient hybrid parallel training framework for large-scale models. Malleus quantifies the stragglers at the nuanced, per-GPU granularity during training, and develops a novel planning algorithm to deduce the optimal parallelization of GPU devices, pipeline stages, model layers, and training data, maximizing training efficiency when stragglers exist. In addition, once a shift in the straggler situation is detected, Malleus adaptively adjusts the parallelization via a re-planning process, and seamlessly and efficiently migrates the model states on the fly, without sacrificing the stability of the training tasks. Empirical results on large language models with up to 110B parameters show that Malleus consistently outperforms existing parallel training frameworks under various straggler situations, delivering on average 2.63-5.28 times of efficiency improvement.