Automatic parallelism strategy generation with minimal memory redundancy

Abstract

Large-scale deep learning models are trained distributedly due to memory and computing resource limitations. Few existing strategy generation approaches take optimal memory minimization as the objective. To fill in this gap, we propose a novel algorithm that generates optimal parallelism strategies with the constraint of minimal memory redundancy. We propose a novel redundant memory cost model to calculate the memory overhead of each operator in a given parallel strategy. To generate the optimal parallelism strategy, we formulate the parallelism strategy search problem into an integer linear programming problem and use an efficient solver to find minimal-memory intra-operator parallelism strategies. Furthermore, the proposed algorithm has been extended and implemented in a multi-dimensional parallel training framework and is characterized by high throughput and minimal memory redundancy. Experimental results demonstrate that our approach achieves memory savings of up to 67% compared to the latest Megatron-LM strategies; in contrast, the gap between the throughput of our approach and its counterparts is not large.

1. Introduction

Large-scale deep learning models have been a tremendously popular area of study in recent years owing to their excellent performance improvements in various fields, such as image recognition, speech recognition, dialogue systems, natural language processing, and recommendation systems (Krizhevsky et al., 2012; Naumov et al., 2019; Brown et al., 2020; Dan et al., 2023), which is a result of increasing model sizes and dataset sizes (Zhuang et al., 2017). For example, PaLM with 540 billion parameters was trained with a corpus of 780 billion tokens representing a wide range of natural language use cases (Chowdhery et al., 2022). As a result, training these models is computation-intensive and time-consuming, requiring a lot of computing resources. However, resources on a single node are limited and insufficient to train these models (Lan et al., 2018; Guan et al., 2020). To address this challenge, researchers have proposed various parallelism strategies to accelerate the process of training large-scale deep learning models. However, designing parallel strategies is not a straightforward task, and it often requires considerable expertise and trial-and-error exploration to achieve high performance (Mo, 2018; He et al., 2023).

To relieve us from the parallelism design procedure, researchers have proposed auto-parallelism methods (Jia et al., 2018; Cai et al., 2022; Zheng et al., 2022; Liu YL et al., 2023) that can find decent strategies given a specific model and environment. These methods model the communication costs of different parallelism strategies and then use a search algorithm to find the optimal strategy. However, existing studies typically aim to minimize the communication cost for a given model. Under this search objective, the resulting parallelism strategies are often not memory-optimal. To fill in this gap, we propose an algorithm to search for memory-optimized operator-level parallelism strategies by analyzing the memory overhead of neural networks in distributed training. This enables us to further increase the hidden layer or batch size on hardware resources with limited memory. Initially, we instantiate a memory analyzer, enabling the computation of memory overhead for any specified strategy. Subsequently, we formalize the problem of strategy search. To accomplish this, we leverage an external solver specifically designed for integer linear programming (ILP) and employ it as a basis for arriving at the decisions, pertaining to optimal multi-dimensional intra-operator parallelism strategies that minimize memory usage.