An Evaluation of Memory Optimization Methods for Training Neural Networks

To improve accuracy and generalization, deep neural networks (DNNs) have rapidly grown in size (Brown et al., 2020; Chowdhery et al., 2022; Touvron et al., 2023). The further growth of these models is constrained by hardware limitations. In particular, the limited memory on GPUs restricts how large a model we can train. In response, researchers have developed various memory optimization methods (MOMs) to reduce the memory footprint of training DNNs. Popular MOMs include gradient checkpointing, swapping, and activation reduced training (ACT). We describe these in detail in subsection 2.2.

In this study, we conduct a survey of MOMs used in DNN training to understand when MOMs prove advantageous. Specifically, we aim to address several key questions: Does the application of these techniques truly improve the training process? Which technique achieves the best time-memory trade-off? Are there MOMs that demonstrate superior performance on particular architectural configurations?

In short, consistent with previous literature, we find that MOMs help in two ways. First, they reduce the peak memory compared to standard training for a fixed batch size. Second, they enable larger maximum batch sizes compared to standard training for a given hardware configuration. However, our central finding is that it is difficult to answer our motivating questions with existing ways of evaluating MOMs. Most of these issues stem from the absence of standard networks, metrics, and experimental practices. Crucially, previous evaluation metrics only consider a single aspect of training and ignore the application scenarios of MOMs. Some of these metrics even conflict with each other. For example, as shown in section 4, synchronous swapping performs better than gradient checkpointing on the Swin-Large model when evaluated with maximum batch size, but gradient checkpointing is more efficient when considering the execution overhead at a given batch size compared with the original training.

In this paper, we provide a comprehensive overview of various MOMs and outline their respective application scenarios. Recognizing the diversity of these scenarios, we propose using distinct evaluation metrics tailored to each scenario. By employing these evaluation metrics, we conduct a thorough assessment of numerous mainstream network architectures. Surprisingly, our findings reveal that MOMs only demonstrate benefits in particular settings. Additionally, we provide a detailed discussion of the conditions under which applying MOMs is advantageous, thereby aiding future researchers in making informed decisions regarding the selection and adoption of MOMs.

In summary, the paper makes the following contributions: (1) We provide a summary of MOM application scenarios and propose tailored evaluation metrics for each scenario. (2) We conduct a comprehensive evaluation of popular MOMs in widely-used frameworks, analyzing their performance in both single GPU and distributed settings.

Previous evaluations of various MOMs can be categorized as follows: (1) Chen et al. (2021); Peng et al. (2020); Liu et al. (2022a) compare the largest batch size the MOM enables given the same memory budget as shown in Figure 1(a), which is limited by the memory capacity of the hardware. (2) Similarly, Chen et al. (2016, 2021); Liu et al. (2022a) compare the largest model the MOM allows given the same batch size and memory budget, using the largest depth and width of the model as a measure of efficacy. (3) Kirisame et al. (2020); Chen et al. (2021); Huang et al. (2020); Jain et al. (2020); Liu et al. (2022a) compare the performance overhead of a MOM with the original training under different memory thresholds, given the same batch size on a specified model, as shown in Figure 1 (b).

However, there is an existing gap between these evaluation metrics and practical use. For those wanting faster training, a larger batch size may not result in higher throughput due to the MOM’s overhead. For those wanting to train larger models, looking at just the largest trainable model is not sufficient as it must be trained using a realistic throughput, which the MOM hinders. Thus we argue that MOMs should be evaluated in the context of their application. We list three main scenarios in which practitioners may find it useful to employ MOMs.

Training larger models: Fixed hardware, evaluated with the maximum model size that satisfies a minimum throughput threshold: The first scenario involves a fixed hardware setting with the goal of training a larger model within the memory limit. Under this scenario, evaluating based on the model size is essential. However, it is equally important to consider the training speed. To address this, we propose a new evaluation metric: the maximum model size that satisfies a minimum throughput threshold. This metric takes into account both the model size expansion achieved through MOM and the practical feasibility of training in terms of throughput.

Faster training: Fixed hardware, evaluated with the maximum throughput: The second scenario involves a fixed hardware setting and aims to enable a larger batch size, reducing the number of iterations required to train a model while increasing device utilization and speeding up the training process. In this case, the evaluation metric to be used is the maximum throughput. As demonstrated in section 4, throughput increases with larger batch sizes. Therefore, the maximum throughput is attained when the model is trained using the maximum batch size. This metric provides a measure of the highest achievable training speed, indicating the effectiveness of MOM in improving the overall throughput of the training process.

Cheaper training: Fixed training setting, evaluated with the cost-latency trade-off curve: Finally, the third scenario involves a fixed training setting, but the hardware configuration is varied to find the best cost/time trade-off with and without MOM. In this setting, the monetary cost of training a model is directly proportional to the number of GPUs employed and the time taken for training. By leveraging MOM, it is possible to reduce the required number of GPUs, resulting in a potentially lower total training cost, despite an extended training time. Therefore, when the model, batch size, and training setting are held constant, a trade-off frontier is established, as illustrated in Figure 1 (c). The practitioner’s preferences (time vs cost) determine what point on the frontier is chosen for training.

In this section, we present two key findings regarding the optimal utilization of MOMs for model training. Subsequently, we conduct experiments to validate and substantiate these findings.

MOMs are more beneficial for larger models. In our study, we concentrate on comparing models within the same family (e.g., different layers of Bert models, GPT models with varying numbers of encoder/decoder blocks, etc.), as models from different families are not directly comparable. As models grow larger in size, the likelihood of the original training being underutilized increases. Therefore, the application of MOM becomes more advantageous as it enables an increase in batch size and improves hardware utilization.

To validate this implication, we conducted experiments by varying the depth and width of ResNet and GPT-3 models on V100 GPUs. The results, as depicted in Figure 6, demonstrate that when the model is shallow enough (with fewer than 280 layers for ResNet or fewer than 18 layers for GPT-3) or narrow (with a hidden size smaller than 248 for ResNet or 1500 for GPT-3), the application of MOMs actually slows down the training process. However, as the model becomes deeper or wider, MOMs start to accelerate the training. In the most extreme case, the original training fails to execute even with a batch size of one, emphasizing the necessity of applying MOMs to enable single-GPU training. These findings provide evidence supporting the claim that the benefits of MOMs are more pronounced with increasing model depth and width.

MOMs are more beneficial for multi-GPU data parallel training. The intuition behind this is that communication makes MOMs less expensive. Communication and synchronization across multiple devices result in a substantially higher fixed execution cost.

We compare the maximum throughput ratio between gradient checkpointing and the original training on Bert-Large and Swin-Large on distributed data-parallel training with 1, 2, 4, 8 V100 GPUs on a single machine. As shown in Figure 6, the ratio increases by adding more GPUs. Specifically, on Bert-Large, gradient checkpointing slows down training on a single GPU setting. However, as we increase the number of GPUs, gradient checkpointing achieves a higher maximum throughput than the original training, and can speed up the training process by 1.7× when there are 8 GPUs.

In this study, we present a comprehensive analysis of the benefits of applying various memory optimization methods (MOMs) in different scenarios. By considering distinct application scenarios, we propose to utilize specific evaluation metrics tailored to each scenario. To assess the effectiveness of existing popular MOMs in mainstream frameworks, we conduct experiments using the newly introduced evaluation metrics. The results demonstrate that the application of MOMs does not always yield favorable outcomes. Thus we urge practitioners to exhibit careful judgement and reflect on their use case before applying MOMs.