LLMem: Estimating GPU Memory Usage for Fine-Tuning Pre-Trained LLMs

1 Introduction

Since the introduction of the Transformer model [Vaswani et al., 2017], researchers have proposed numerous language models based on it. As the model's performance has improved, its size has grown exponentially, necessitating a substantial dataset for training. However, training emerging large language models (LLMs) is infeasible without a dedicated infrastructure with high-performance hardware due to memory constraints. Instead, it is preferred to utilize a small dataset to fine-tune a pre-trained model for a specific application.

To efficiently handle small datasets and reduce training time, the conventional method of data parallelism places the entire model on each GPU, splits the dataset, and trains simultaneously. Nevertheless, the model size remains huge, potentially causing GPU out-of-memory (OOM) issues. Therefore, it is necessary to reduce the amount of memory a GPU uses by splitting the model and distributing it to each GPU.

∗This work is based on Taeho Kim's internship at AWS.

†Seyeon Kim is the corresponding author.

‡This work is unconnected to current role at AWS.

ZeRO [Rajbhandari et al., 2020] Stage 3 is an advanced data parallelism method that partitions the model parameters, gradients, and optimizer states to each GPU for memory advantage while maintaining the distribution of the dataset across GPUs. Although ZeRO Stage 3 saves memory by using only partitioned model data on each GPU during noncomputation phases, there are limitations in preventing GPU OOM issues because partitioned parameters/gradients must be all-gathered during computation.

Tensor parallelism divides each parameter tensor in the model into rows or columns and distributes them to each GPU, using only partitioned parameters on each GPU during computation. For example, Megatron-LM [Shoeybi et al., 2019], a representative tensor parallelism method, splits a tensor along its rows or columns considering the position and connection of operators. By doing so, it can reduce GPU memory usage more than data parallelism when the model size is large.

As we described above, various distributed fine-tuning methods have been proposed, but the GPU memory usage and fine-tuning time required for each are different. For instance, conventional data parallelism provides the shortest fine-tuning time but requires the highest GPU memory usage. On the other hand, tensor parallelism has no benefit in saving fine-tuning time but can significantly reduce GPU memory usage. Users may want to select an appropriate method that avoids GPU OOM and has a short fine-tuning time. However, it is difficult to determine in advance whether there is enough GPU memory to fine-tune a given pre-trained LLM.

DNNMem [Gao et al., 2020] is the most recent work detailing procedures to estimate GPU memory usage on a single GPU. DNNMem provides key equations for GPU memory estimation when training various DNN models by analyzing the connections between operators and live tensors in the forward and backward passes. However, it has limitations for fine-tuning LLMs. GPU memory estimation for fine-tuning transformer-based LLM is challenging for two reasons.

First, when fine-tuning an LLM in multi-GPU, distributed fine-tuning methods should be used to overcome GPU memory constraints due to large model sizes. Depending on the method used, the distribution of parameters, gradients, and optimizer states to each GPU is different, as is the amount of GPU memory used during the calculation process. Therefore, GPU memory usage estimates from a single GPU cannot be used in a multi-GPU environment.

Second, GPU memory consumption must be predicted by distinguishing between transformer and language modeling head (lm head) parts. The transformer part is the central part of fine-tuning, where chunk memory management for memory sharing of model parameters and gradients is applied, and parameters are updated. On the other hand, the lm head part requires separate analysis because it does not apply distributed methods directly and consumes a lot of memory due to its large dictionary size.

To address these challenges, we propose LLMem that estimates the GPU memory consumption when applying distributed fine-tuning methods to multiple GPUs. LLMem considers several factors to estimate GPU memory usage for each method, including recombining parameters prior to computation when applying advanced data parallelism and the output driven by all-gather in the backward pass when using tensor parallelism. Additionally, LLMem analyzes the difference in memory allocation method between the transformer and the lm head part and reflects it in GPU memory estimation. To the best of our knowledge, this is the first work to estimate the peak GPU memory consumption for LLM fine-tuning.

In summary, our contributions are:

• We propose a GPU memory usage estimation method for LLM fine-tuning on single and multiple GPUs.
• We provide an algorithm to determine the most efficient distributed fine-tuning method based on GPU memory usage estimation.
• Experimental results show that LLMem estimates peak GPU memory usage to fine-tune LLM on a single GPU with error rates of up to 1.6%, which is significantly smaller than the state-of-the-art DNNMem's average error rate of 42.6%. When applying distributed fine-tuning methods to LLMs with over a billion parameters on multiple GPUs, LLMem successfully estimates GPU memory usage with an average error rate of 3.0%.

Our source code repository can be found at https://github.com/taehokim20/LLMem.

2 RelatedWorks

2.1 GPUMemory Estimation

There have been several attempts to avoid GPU OOM issues by predicting the GPU memory usage that will be used to train a given model in advance. DNNMem [Gao et al., 2020] sequentially traverses the computation graph of a DL model and computes the GPU memory consumption by taking into account previously allocated but still in-use tensors, newly allocated tensors for the currently visited operator, and resident buffers of the CUDA context and allocator reservation. Our LLMem is inspired by DNNMem, whose mechanism is described in more detail in Section 3. TSplit [Nie et al., 2022] also calculates the total size of live tensors for the visiting operator. However, TSplit lacks an explanation of the detailed memory estimation process and its accuracy. SchedTune [Albahar et al., 2022] predicts GPU memory usage not only based on DL model characteristics but also on different GPU types running the job. However, using measured GPU memory as data for prediction does not align with the purpose of estimating memory usage before fine-tuning a model.

2.2 Distributed Fine-Tuning with GPUs

Data parallelism can enhance fine-tuning speed in proportion to the number of GPUs. However, LLM often runs into memory constraints, so the ZeRO optimizer [Rajbhandari et al., 2020], described in Section 1, is widely used as an alternative. The ZeRO optimizer selectively gathers only the model parameters or gradients required during the computation process and utilizes reduce-scatter after the computation to maintain their partitioning on each GPU.

Tensor parallelism can further reduce peak GPU memory usage by sharding tensors under certain conditions, eliminating the need to gather all model parameters and gradients even during computation. Tensor parallelism results in each GPU producing only partial results, necessitating that all GPUs receive the same input data. The widely adopted tensor parallelism method, Megatron-LM [Shoeybi et al., 2019], splits each model parameter tensor by row or column. Other proposed methods [Xu et al., 2021] [Wang et al., 2021] [Bian et al., 2021] achieve additional memory savings by sharding both input and model parameters.

If GPU memory constraints cannot be met with any distributed fine-tuning method on GPUs alone, we can use heterogeneous fine-tuning utilizing CPU memory. ZeROoffload [Ren et al., 2021] manages gradients, optimizer states, and optimizer computation on the CPU while retaining parameters and forward and backward computation on the GPU.

3 Motivation

To select distributed fine-tuning methods, it is crucial to estimate GPU memory usage accurately. Existing approaches for estimating GPU memory usage do not consider scenarios where advanced data parallelism, such as the ZeRO Stage 3 optimizer [Rajbhandari et al., 2020], or tensor parallelism is applied across multiple GPUs. Relying solely on estimated GPU memory usage on a single GPU when estimating on multiple GPUs can lead to significant errors. In this section, we implement the existing DNNMem [Gao et al., 2020], validate the implementation results, and discuss factors causing substantial GPU memory estimation errors during the finetuning of pre-trained transformer-based language models.

3.1 DNNMem Implementation

DNNMem source codes are not publicly available and are mainly based on TensorFlow, so we implement DNNMem based on the description in the paper [Gao et al., 2020]. First, we extract the corresponding computation graph from a given pre-trained DL model to identify the output size in each operator based on parameters, batch size (bs), and sequence length (sl). We also compute pre-allocated GPU memory, including CUDA context and weight tensors of the model, before operator execution. In particular, since PyTorch does not release the loaded model parameters until the end of the fine-tuning, the initial GPU memory is retained throughout the fine-tuning process. The next step is to compute peak GPU memory usage at each operator while traversing the graph. We compute additional GPU memory with the input/output tensors and previously unreleased tensors in each operator during the forward propagation. Additionally, we reflect that PyTorch aligns with multiples of 512 bytes for internal tensor fragmentation, and DNNMem treats the buffer size as a constant (64 MB by default) as memory block management.

To validate our DNNMem implementation, we compare GPU memory estimation results for the BERT [Devlin et al., 2018] model on the GLUE benchmark [Wang et al., 2018] with the experimental results from the paper. The environment we used in the experiment was PyTorch 2.0.1 with CUDA 11.7 on NVIDIA RTX2060, which differs from PyTorch 1.2.0 with CUDA 9.0 on NVIDIA Tesla P40 used in the DNNMem paper. In bs = 32, sl = 32, our DNNMem shows 34.38%, and the DNNMem shows 31.42% error rates. In bs = 64, sl = 32, our DNNMem shows 20.48%, and the DNNMem shows 19.12% error rates. Considering similar error rates, we use our DNNMem implementation for singleGPU comparisons.

3.2 Limitations of DNNMem for LLM Fine-Tuning Memory Estimation

DNNMem [Gao et al., 2020] does not handle mixed precision, which is commonly used in fine-tuning pre-trained language models. In addition, it does not consider howmemory chunks are managed to ensure that forward pass parameters and backward pass gradients share the same GPU memory space [Fang et al., 2022]. Furthermore, DNNMem overlooks extra GPU memory usage during the initial fine-tuning iteration due to optimizer states.

Comparison results for estimating peak GPU memory usage of our proposed LLMem and DNNMem on a single GPU are shown in Figure 1. The experimental environment is summarized in Section 7.1. LLMem predicts peak GPU memory usage with minimal error rates compared to ground truth, outperforming DNNMem. DNNMem exhibits larger errors as the total parameter size increases. Furthermore, DNNMem fails to predict GPU memory consumption in the context of distributed fine-tuning methods across multiple GPUs. As a

Table 1: Notation

result, existing approaches for estimating GPU memory usage face challenges when using the current transformer-based LLM for distributed fine-tuning.

4 Single-GPUMemory Usage Estimation

This section outlines considerations for estimating GPU memory usage of transformer-based language models on a single GPU. The symbols used in the explanation are organized in Table 1.

4.1 Workflow for Fine-Tuning Pre-TrainedModels

Initialization phase. The initialization phase preceding finetuning involves allocating memory for the CUDA context, responsible for managing information to control GPU devices, and memory for applying chunk-based memory management [Fang et al., 2022]. The initially used GPU memory is denoted as mbase. The chunk manager determines the optimal chunk size to minimize GPU memory waste based on the parameters of the provided pre-trained language model. GPU memory spaces for param fp16 (float-16) and param fp32 (float-32) are allocated in units of the chunk size ( 1 in Figure 4).

Fine-tuning phase. During the fine-tuning phase, param fp16 goes through forward and backward passes, and param fp16

is converted to gradient fp16, as illustrated in Figure 2. Consequently, param fp16 and gradient fp16 share the same GPU memory space. After the backward pass, the ADAM optimizer updates parameters using optimizer states, including param fp32, momentum fp32, and variance fp32 tensors. Momentum fp32 and variance fp32 tensors, which are not allocated memory during the initialization process before finetuning, consume GPU memory based on the actual tensor size, not the chunk size. GPU memory occupied by these tensor types is allocated in the first iteration for fine-tuning ( 5 in Figure 4). Subsequently, similar to chunk-based parameters, the GPU memory is retained until the fine-tuning process is complete.

4.2 Memory Consumption with Structure of Transformer-based Decoder Model

The peak GPU memory usage on a single GPU (mspeak ) is

Each variable in this formula, except mbase described in Section 4.1, is calculated as follows.

First, mp is the GPU memory used by param/gradient fp16 and param fp32. Considering that param fp16 and fp32 use 2

bytes (B16) and 4 bytes (B32) per value, respectively, mp is

, where embedp is the input embedding param size, cs is the chunk size, otherp is the remaining param size, and cup is the CUDA memory page size, typically 2 × 10242 bytes. Transformer-based decoder models [Vaswani et al., 2017] are largely divided into a transformer model for fine-tuning and lm head for output, as shown in Figure 3. The part that uses the chunk memory is the transformer model in which the parameters are updated. The embedp is huge because the input embedding has a large dictionary. Therefore, embedp is managed separately.

Second, mos is the GPU memory used by momentum fp32 and variance fp32 of optimizer states. mos is

, where t is the operator of the given transformer model, E is Embedding, L is Linear, and tp is the parameter size of t. The system allocates GPU memory based on the actual size of each momentum fp32 and variance fp32, so GPU memory must be calculated for each tensor of each operator. Since the amount of GPU memory consumed by Bias or LayerNorm is very small, they can use space with other memory fragmentation. Therefore, we only calculate the GPU memory usage due to Embedding or Linear operator parameters.

Third, mout is the peak GPU memory usage due to output tensors. If the number of Embedding, layers, and model's output features are en, ln, and on, respectively, then mout is

PyTorch provides gradient checkpointing 1 as an option to save memory during fine-tuning. Therefore, we support estimating GPU memory usage due to each operator's input/output tensors considering gradient checkpointing. Since the output tensors of the current operator are the input tensors of the next operator, we focus on the output. It is challenging to accurately predict GPU memory consumption due to the outputs of operators within a model. We observed that the layer and embedding outputs of the transformer model are kept in GPU memory for efficient gradient checkpointing, which minimizes the increase in fine-tuning time. The estimation error rate is reduced using the mout equation, which accounts for our observation.

Lastly, mlm is the GPU memory used in the lm head part including the loss calculation part. If the size of the embedding dictionary is dictn, mlm is

1 Gradient checkpointing reduces GPU memory usage by clearing specific outputs and recomputing them during a backward pass.

, where lmp is the lm head param size, and B is either B16 or B32 depending on the model type. The lm head converts the transformer outputs into logits. Then, the value obtained by shifting the sequence length of the logits by one space is stored in a separate temporary variable and used for the loss calculation. mout +mlm is the output phase ( 3 in Figure 4).

5 Multi-GPUMemory Usage Estimation

This section outlines the factors for estimating peak GPU memory usage during distributed fine-tuning on multiple GPUs and summarizes the estimation process.

Conventional data parallelism (CDP). Since CDP places the entire model on each GPU, its peak GPU memory usage estimation equals the peak single-GPU memory usage mspeak , as shown in Figure 4.

Advanced data parallelism (ADP). The peak GPU memory usage with ADP on multiple GPUs (mdp peak ) is

, where mp,16 and mp,32 are the GPU memory consumed by the entire param fp16/fp32, and gpun is the number of GPUs in use. ZeRO-3 optimizer, a method of advanced data parallelism, evenly distributes parameters, gradients, and optimizer states by gpun, reducing GPU memory usage. Among these, gradient fp16 shares GPU memory with param fp16 as explained in Section 4, so we only need to divide the GPU memory usage of parameters and optimizer states by gpun. However, during the calculation process, each GPU must have all the values of param fp16 (Figure 5a and 2 in Figure 4), so the whole mp,16 is allocated to the GPU memory.

Tensor parallelism (TP). The peak GPU memory usage with 1D TP on multiple GPUs (mtp peak ) is

, where mtp back ( 4 in Figure 4 and Figure 6) is the additional GPU memory usage due to the temporary buffer through the backward all-gather. If the number of GPUs used for tensor parallelism is tpn, mtp back is

Tensor parallelism divides the parameter values of each operator by gpun and does not combine them again, as shown in Figure 5b. It splits each model parameter tensor by row or column to apply tensor parallelism to multiple pre-trained language models. We call this one-dimension tensor parallelism (1D TP). Let us assume that we apply 1D TP to a linear operation on four GPUs. The linear operator's equation is y = xAT + b, where y is output, x is input, A is params/gradients, and b is bias. The linear matrix multiplication process when each parameter tensor is split into columns is shown in Figure 6. We shard parameters by column because the output size after multiplication is the same as the size of the bias without sharding, so it is not affected by the use of bias. In the backward pass, the fine-tuning goes through an all-gather process. mtp back is the total temporary buffer size for tensors imported from the other GPUs, calculated by multiplying the output size of each layer by the number of layers.

Combination of DP+TP. The peak GPU memory usage with the combination of DP+TP on multiple GPUs (mdp+tp peak ) is

, as shown in Figure 4. It is possible to achieve hybrid parallelism by fine-tuning through a combination of data and tensor parallelism.

Algorithm 1 Distributed Fine-Tuning Method Decision

Input: Pre-trained model M, gpun, and sl

Output: Selected fine-tuning method and the optimal bs 1: eval = [0, 0, 0, 0] and bslist = [0, 0, 0, 0] 2: Measure total GPU memory capacity (mtotal) 3: for i in range(4) do 4: Set up the configure of the ith method 5: bs = 1 and compute mbase, mp, and mos 6: [1]: Compute mout and mlm 7: if i = 0 then 8: Repeat bs = bs + 1 until mspeak > mtotal after [1] 9: eval[i] = (bs - 1) × gpun × 1.5 10: else if i = 1 then 11: Repeat bs = bs + 1 until mdp peak >mtotal after [1] 12: eval[i] = (bs - 1) × gpun 13: else if i = 2 then 14: Repeat bs = bs + 1 until mtp peak >mtotal after [1] 15: eval[i] = bs - 1 16: else if i = 3 then 17: Repeat bs = bs + 1 until mdp+tp peak >mtotal after [1] 18: eval[i] = (bs - 1) × dpn 19: end if 20: bslist[i] = bs - 1 21: end for 22: Save the index of the maximum score to idx 23: return idx, bslist[idx] if eval[2] ̸= 0 else 4, 0

6 Distributed Fine-TuningMethod Decision

Algorithm 1 describes the process for selecting the optimal method to fine-tune a pre-trained model based on the results of estimating the peak GPU memory usage. In Sections 4 and 5, We estimated mspeak , mdp peak , mtp peak , and mdp+tp peak . Here, mspeak represents CDP, and the remaining estimations are connected to ADP, TP, and DP+TP, respectively. Of these methods, the optimal one is the method that requires the shortest time for fine-tuning while avoiding GPU OOM.

LLMem takes a pre-trained model M, the total number of GPUs to fine-tune gpun, and the maximum sequence length sl. eval is a list that stores the performance evaluation score of each method. eval[0], eval[1], eval[2], and eval[3] correspond to CDP, ADP, TP, and DP+TP, respectively. LLMem increments the batch size bs for each method and gets the value of bs when it reaches the total GPU memory capacity. Then, bs - 1 is the largest batch size to avoid GPU OOM. CDP uses (bs - 1) × gpun amount of data for fine-tuning in one iteration. In addition, since the ZeRO-3 optimizer increases the total communication volume of a baseline DP to 1.5× [Rajbhandari et al., 2020], the performance score of CDP is (bs - 1) × gpun × 1.5. In one iteration, ADP uses (bs-1)×gpun, TP uses bs-1, and DP+TP uses (bs-1)×dpn of data for fine-tuning. dpn is the number of GPUs used for DP. These values become the performance scores of each method. Finally, LLMem selects the method with the highest performance score (If the scores are tied, select CDP, ADP, TP, and DP+TP in that order). If the performance scores of all methods are 0, heterogeneous training using CPU memory is selected as an alternative to avoid GPU OOM.

7 Experiments

In this section, we compare the peak GPU memory usage estimate of LLMem with the ground truth data when applying various distributed fine-tuning methods. In addition, our DNNMem implementation is included in comparing GPU memory usage estimation on a single GPU.

7.1 Experimental Setup

For a multi-GPU environment, we use a Tesla V100 (total GPU memory capacity: 16384 MB) with 4 GPUs in CloudLab [CloudLab, 2024]. We also use the Colossal-AI [Li et al., 2023] 2 , a widely used framework for applying distributed fine-tuning methods, and PyTorch 2.0.1 with CUDA 11.7. The models we used in the experiment are OPT [Zhang et al., 2022], BLOOM [Workshop et al., 2022], CodeGen [Nijkamp et al., 2022], BioGPT [Luo et al., 2022], GPTBigCode [Allal et al., 2023], GPT Neo [Black et al., 2021], and LLaMA [Touvron et al., 2023]. The dataset used is alpaca data [Taori et al., 2023], which is 52K instruction-following data. For the ground truth data, we measure peak GPU memory usage using only the maximum sequence length of 512.

7.2 Estimation of Single-GPUMemory Usage

First, we compare the peak GPU memory usage estimate from LLMem for a single GPU with the DNNMem estimate and the actual peak GPU memory usage. Since we used gradient checkpointing for LLM fine-tuning, the same approach was applied to DNNMem. Figure 7 compares the peak GPU memory usage estimation results of LLMem and DNNMem for various pre-trained LLMs that cause GPU OOM during fine-tuning on a single GPU. LLMem predicts GPU OOM for all models, while DNNMem predicts peak GPU memory usage that falls short of mtotal. Table 2 shows the predicted and actual GPU memory usage peaks when applying the maximum batch size to obtain the ground truth data for each model during fine-tuning on a single GPU. DNNMem underestimates the peak GPU memory usage for all models because it does not account for factors considered when finetuning Transformer-based LLM, as explained in Section 3.2. LLMem's GPU memory estimation helps approximate the peak GPU memory usage close to the ground truth.

2 LLMem's approach can extend to other platforms such as DeepSpeed and PyTorch-FSDP with the Hugging Face library.

7.3 Estimation of Multi-GPUMemory Usage

CDP. The experimental results are the same as the memory usage estimation results on a single GPU in Section 7.2.

ADP. Figure 8 shows the predicted and actual GPU memory usage peaks when applying the maximum batch size to obtain ground truth data for each model during fine-tuning with ADP on four GPUs. The error rate between the predicted value of LLMem and the actual GPU memory usage tends to increase on multi-GPU setups. One reason is the gap in memory usage between the GPUs. ADP places tensors separately on each GPU instead of being sharded, so not all GPUs can use precisely the same amount of memory. Second, the error tends to be slightly larger when the model size is large. A larger number of layers and outputs in large models can lead to larger error rates due to memory allocator characteristics. TP and DP+TP. Figure 9 shows the predicted and actual GPU memory usage peaks when applying the maximum batch size to obtain ground truth data for each model during fine-tuning with 4TP or 2DP+2TP on four GPUs. 4TP uses 4 GPUs in TP, and 2DP+2TP uses 2 GPUs in DP and 2 GPUs in TP for hybrid parallelism. We focus on estimating the peak GPU memory usage of the large-size model for TP because LLMem can select DP for quick fine-tuning of models that are small and do not have OOM problems. TP applies the all-gather operation in the backward pass, as shown in Figure 6. The all-gather operation allocates temporary buffers in GPU memory and collects values in those buffers, consuming additional GPU memory. If the model size is large and the possible batch size is small, the system can use the allocated but currently empty memory space for a temporary buffer.

Model

Table 3: Distributed fine-tuning method selection by LLMem on four GPUs and the actual amount of time it takes to fine-tune each method to the largest possible batch size (s)

Therefore, the GPU memory consumed due to the temporary buffer does not increase excessively, leading to smaller errors as shown in Figure 9. 2DP+2TP shows slightly larger errors than 4TP in most cases. This is because GPU memory usage due to the temporary buffer may be additionally affected in 2 and 4 of Figure 4 while applying both DP and TP.

7.4 Fine-TuningMethod Selection with LLMem

Table 3 assesses whether LLMem finds the optimal finetuning method to achieve the fastest fine-tuning while avoiding GPU OOM for various models. When measuring the time taken for each method, we applied the maximum batch size that can prevent GPU OOM. LLMem typically selects TP when DP causes GPU OOM. It is challenging for LLMem to choose DP+TP because only 4 GPUs were used in the experiment. DP+TP allows for more diverse combinations depending on the number of GPUs used and is more likely to be selected. LLMem also suggests CPU offloading when GPU memory is insufficient.

8 Conclusion

This paper introduces LLMem, a method for estimating GPU memory consumption during fine-tuning of large language models (LLMs) on multi-GPU setups. We analyze factors affecting GPU memory usage, considering different memory allocation methods for the transformer and output sections. Experimental results demonstrate that LLMem achieves accurate peak GPU memory usage estimation on both single and multiple GPUs with minimal error rates.

Acknowledgements

This work was supported by the Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korean government (MSIT) (No.RS2023-00277060, Development of open edge AI SoC hardware and software platform).