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

TL;DR

This work addresses dynamic stragglers in hybrid parallel training of very large models by introducing Malleus, which quantifies per-GPU straggling through rates $x_i$ and uses a bi-level optimization to compute non-uniform partitions across GPUs, pipeline stages, model layers, and training data. It includes a malleable training workflow with asynchronous re-planning and on-the-fly model migration to maintain high efficiency without restarts. The planning leverages ILP/MINLP formulations to produce optimal or near-optimal configurations within tens of seconds, and experiments on LLMs up to 110B parameters show substantial speedups (average $2.63$–$5.28 imes$, up to $6.73 imes$) over strong baselines. The results demonstrate robust, scalable training on heterogeneous GPU clusters, highlighting the practical impact of fine-grained straggler resilience for future large-scale model development.

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.

Figures (11)

• Figure 1: An example of 3-Dimensional (3D) parallel encompassing data parallel (DP), tensor parallel (TP), and pipeline parallel (PP). The model consists of 32 layers and the global batch size is 64.
• Figure 2: Key characteristics of the parallelization plans in our work, illustrated with the example workload in Figure \ref{['fig:3d_parallel']}.
• Figure 3: Architecture overview of Malleus.
• Figure 4: Overview of the planning routine, illustrated with an example where 3 out of 16 GPUs are stragglers (depicted in the lower left). The model consists of 32 layers and the global batch size is 64. The planning routine partitions the GPUs into 6 groups (upper left), organizes them into 2 pipelines (upper right), and assigns the model layers across different pipeline stages as well as assigns the training data across the pipelines (lower right).
• Figure 5: An example of GPU grouping with splitting.

Theorems & Definitions (10)

• Theorem 1
• Theorem 2
• Theorem 3
• Theorem 1
• Theorem 2
• Theorem 3
• Proposition 1
• Proposition 2
• Proposition 3
• Proposition 4