A method to accelerate the training of large language models

By mathematically modeling and phase division of the big model, combining tensor row optimization and data parallelism, and using synchronous communication and CPU offloading technology, the problems of long training time, memory limitation and nonlinearity of computing scalability in large model training are solved, and efficient resource utilization and reduced video memory usage are achieved.

CN118779095BActive Publication Date: 2025-04-22CHINA ACADEMY OF ELECTRONICS AND INFORMATION TECHNOLOGY OF CHINA ELECTRONICS TECHNOLOGY GROUP CORPORATION
View PDF 2 Cites 0 Cited by

Patent Information

Application Number
CN202410747978.8
Authority / Receiving Office
CN · China
Patent Type
Patents(China)
Current Assignee / Owner
Filing Date
2024-06-11
Publication Date
2025-04-22
Estimated Expiration
2044-06-11

AI Technical Summary

Technical Problem

Large model training faces the problems of huge number of computational operations, resulting in long training time, limited memory capacity, unlinear computational scalability, and the lack of detailed numerical indicators for modeling and analysis by existing optimization methods.

Method used

By mathematically modeling the different network layers of the large model, analyzing the parameter quantity and calculation quantity, combining calculation hardware parameters to divide the large model stage, tensor row optimization and data parallelism are adopted, and All-reduce synchronous communication time and CPU offload technology are used to optimize GPU video memory usage and CPU-GPU communication.

Benefits of technology

It realizes accurate evaluation of the resource occupancy of large models, finds training bottlenecks, improves GPU and CPU resource utilization, reduces GPU memory usage, and solves the problems of unclear division principles and unbalanced load in traditional methods.

✦ Generated by Eureka AI based on patent content.

Smart Images

  • Figure CN118779095B_ABST
    Figure CN118779095B_ABST
Patent Text Reader

Abstract

The present invention provides a large language model training acceleration method, comprising the following steps: mathematical modeling of different network layers of the large model, analyzing the parameter amount of the corresponding layer, and calculating the calculation amount of different layers: the model parameter amount refers to the number of parameters that can be learned and adjusted in a neural network or machine learning model, and these parameters include weights and biases, which are continuously updated during the training process to optimize the performance of the model, and the model calculation amount is expressed as the number of floating point operations in FLOPs, which measures the size of the calculation amount. The present invention provides a large language model training acceleration method, which can accurately evaluate the resource occupancy in large model training, further analyze to find the bottleneck in large model training, and help researchers propose large model acceleration training methods, which is of great significance for the analysis and optimization training in the field of large models.
Need to check novelty before this filing date? Find Prior Art

Description

Technical Field

[0001] The present invention relates to the field of large model training, and in particular to a large language model training acceleration method. Background Art

[0002] A large model refers to a neural network model that contains ultra-large-scale parameters. It has powerful expressive and learning capabilities, can handle more complex and larger data sets or tasks, and has stronger generalization and accuracy. Due to the scaling law and "emergence ability" of large models, the parameters of large models increase dramatically at an order of magnitude, and the training of large models becomes increasingly difficult.

[0003] Due to the scaling law of large models, larger models bring significant performance improvements, but models with huge training volumes bring many practical challenges. The current optimization methods for large model training are mainly: through modeling analysis of large models, such as communication scheduling and memory usage of large models, combined with multiple parallel technologies for optimization. The bottlenecks that restrict our efficient training of large models are mainly computing resources and network communications. Large model training is challenging, and the main problems it faces are:

[0004] like Figure 1 As shown, the number of computational operations required leads to impractically long training times; and, due to limited memory capacity, it is no longer possible to fit the entire model on a single accelerator;

[0005] Due to communication costs and memory redundancy, computational scalability is not linear. For example, if we use twice the number of GPUs, the training time cannot be reduced to half.

[0006] The existing large model training optimization methods do not have a complete set of numerical indicators to perform detailed modeling analysis on the parameter quantity, computational complexity, memory usage of each part of the large model training, as well as the entire process.

[0007] Therefore, it is necessary to provide a large language model training acceleration method to solve the above technical problems. Summary of the invention

[0008] The present invention provides a large language model training acceleration method, which solves the problems of inaccurate analysis of existing large model training and difficulty in large model training.

[0009] In order to solve the above technical problems, the present invention provides a large language model training acceleration method, comprising the following steps:

[0010] S1. Mathematically model different network layers of the large model, analyze the parameters of the corresponding layers, and calculate the computational complexity of different layers: Model parameters refer to the number of parameters that can be learned and adjusted in a neural network or machine learning model. These parameters include weights and biases, which are continuously updated during the training process to optimize the performance of the model. The model computational complexity is expressed in FLOPs, which represents the number of floating-point operations and measures the size of the computational complexity.

[0011] S2. After mathematically modeling the network layer of the large model, it is necessary to count the parameters and calculation amount of the large model as a whole. For the parameters of the large model as a whole, the parameters of the AttentionMask layer, the Embedding layer, and the Decoder layer are ignored here. The Transformer model will be overlapped and arranged in the large model. As the main body of the large model, the parameters of the main body of the large model are listed here to represent the overall parameters of the large model, which is convenient for later analysis.

[0012] The specific formula is as follows:

[0013] Eliminate the parameters of the Embedding layer and the Decoder layer:

[0014] φ=n layer *(3*d model *d attn +d model *d attn +2*d model *d ff )=2*n layer *d model *(2*d attn +d ff ) Each Token

[0015] Calculation Amount:

[0016]

[0017] S3. Consider the parameters of computing hardware in reality, and combine the mathematical modeling analysis of the large model to divide the large model into different stages so that it can run stably on GPUs. Assuming that there are n A100 GPUs to run a stage of the large model, and each stage of the large model contains t Transformer layers, it is necessary to ensure that the memory usage and computing power of this stage of the large model do not exceed the theoretical upper limit of n GPUs (generally 80% of the peak value is taken as the upper limit). Let φ′ represent the parameter quantity of this sub-model, and determine the number of Transformer layers n in this sub-model. l ' ayer .

[0018] For memory usage, n A100 GPUs contain n*40G memory, and the memory of the sub-model includes parameters 2φ′, gradients 2φ′, and optimizer states 12φ′, a total of 16φ′:

[0019] 16φ′=16*2*d model *(2*d attn +d ff )*n l ' ayer ≤n×40×80%

[0020] For the peak computing, the peak computing capacity of each A100 GPU is 19.5 TFLOPS, and the computing capacity of the sub-model is C′=2*φ′+2*n l ' ayer *n ctx *d model <<19.5×80%TFLOPS;

[0021] S4. In the training of large models, tensor row optimization is used for the calculation of the Transformer model. Tensor parallelism can be decomposed into row parallelism and column parallelism. Then, tensor parallelism is applied to the Attention structure and MLP structure respectively.

[0022] S5. Utilize the synchronous communication time of All-reduce and use data CPU offloading to save the parameters of the previous stage and the intermediate calculation results, thereby reducing the required GPU memory. At the same time, manage the GPU video memory and transfer the model parameters of the next stage to the GPU video memory to improve the GPU computing utilization. During the forward calculation process, enable CPU offloading when the GPU performs synchronous communication, and save the intermediate activation values ​​of the previous round of calculation results on the CPU. When updating the backward parameters, use the synchronous communication time during reverse communication to transfer the intermediate activation values ​​saved on the CPU in the forward calculation to the GPU for calculation, and save the updated parameters on the CPU.

[0023] Preferably, the analysis of the stored parameters and calculation amount of different network layers in S1 includes the following steps: S11. The parameter amount in the network layer mainly refers to the weight matrix in the network, and its number is determined by the matrix dimension, and the matrix dimension is related to its input and output dimensions. Here, d is used uniformly. model Represents the dimension of its output, while the other dimensions are related to the output or batch size batch;

[0024] S12. The computational complexity of large models mainly lies in the calculation of parameter matrices and inputs in the network layer, which can be divided into two categories: addition and subtraction and multiplication and division.

[0025] The Transformer structure can be divided into Attention QKV layer, Attention Mask layer, AttentionProject layer, FeedForward layer, Residual connection layer, etc. Generally speaking, there will be an Embedding layer and a Decoder layer at the beginning or end of the main body of the large model to complete the mapping of input to vector and vector to Tokens.

[0026] Preferably, the step of sequentially listing the parameter quantities and calculation quantities of different network layers in S12 comprises the following steps: S121, Embedding layer

[0027] The word embedding matrix parameter size is: n vocab *d model ,

[0028] The context input vector size is: n ctx *d model ,

[0029] The word embedding matrix calculation size is: n vocab *n ctx *d model ;

[0030] S122, AttentionQKV layer

[0031] The parameter size of the AttentionQKV layer is: 3*n layer *d model *d attn ,

[0032] The computational size of the AttentionQKV layer is: 6*n layer *d model *d attn ;

[0033] S123, AttentionMask layer

[0034] The parameter size of the AttentionMask layer is: d attn ,

[0035] The calculation size of the AttentionMask layer is: 2*n layer *d model *d attn ;

[0036] S124, AttentionProject layer

[0037] The size of the AttentionProject layer parameters is: n layer *d model*d attn ,

[0038] The calculation size of the AttentionProject layer is: 2*n layer *d model *d attn ;

[0039] S125, Feedforward layer

[0040] The parameter size of the Feedforward layer is: 2*n layer *d model *d ff ,

[0041] The calculation amount of the Feedforward layer is: 4*n layer *d model *d ff ;

[0042] S126, Decoder layer

[0043] The size of the Decoder layer parameters is: n vocab *d model ,

[0044] The computational size of the decoder layer is: 2*n vocab *d model .

[0045] Preferably, the memory usage in the large model training in S2 mainly includes optimizer state, gradient and activation, and modeling analysis is performed: firstly, the large model optimizer state and gradient are analyzed. These two memory usages are closely related to the parameter amount of the large model, and are also affected by the memory optimization strategy. Firstly, assuming that the model parameter amount is φ, the relationship between the large model gradient and the optimizer state and the large model parameters is confirmed;

[0046] When the model is stored in FP32 data type, each element occupies 4 bytes, so the memory usage of the parameters is 4φ. The gradient is equal to the model parameters, and the memory usage is 4φ. The AdamW optimizer needs to save the parameters, first-order gradients, and second-order momentum during back propagation. Therefore, the memory usage of the optimizer state is 4φ (parameters) + 4φ (momentum) + 4φ (variance).

[0047] The activations here refer to all tensors calculated during the forward pass and needed in the backward pass;

[0048] The AttentionMask layer needs to complete the following operations:

[0049]

[0050] Therefore, it is necessary to store the intermediate values:

[0051] Q,K,V:[batch,n heads ,n ctx ,d attn ]

[0052] QK T :[batch,n heads ,n ctx ,n ctx ]

[0053]

[0054] The calculation formula of the Feedforward layer is: gelu (x attn W1)W2+x attn The intermediate values ​​that the Feedforward layer needs to store are:

[0055] x attn :[batch,n ctx ,d model ]

[0056] x attn W1:[batch,n ctx ,d ff ]

[0057] f gelu (x attn W1):[batch,n ctx ,d ff ]

[0058] f gelu (x attn W1)W2:[batch,n ctx ,d model ]

[0059] Dropout(f gelu (x attn W1)W2):[batch,n ctx ,d model ];

[0060] To summarize the above analysis of different layers, confirm the amount of parameters required to store for intermediate activation:

[0061] The parameter quantity of the intermediate activation of the AttentionQKV layer is:

[0062] batch*n ctx *dmodel ;

[0063] The parameter amount of the intermediate activation of the AttentionMask layer is:

[0064] 4*batch*n heads *n ctx *d attn +3*batch*n heads *n ctx *n ctx ;

[0065] The number of parameters activated in the middle of the AttentionProject layer is:

[0066] 2*batch*n ctx *d model ;

[0067] The number of parameters activated in the middle of the Feedforward layer is:

[0068] 2*batch*n ctx *(d model +d ff )+batch*n ctx *d model ;

[0069] The number of parameters activated in the middle of the LayerNorm layer is:

[0070] 2*batch*n ctx *d model ;

[0071] The total number of intermediate activation parameters is:

[0072] C activation =n layer *batch*n ctx *(14*d model +4*d ff +8*n heads *d attn +5*n heads *n ctx ).

[0073] Preferably, S4 includes the following steps: S41, the splitting method of the MLP module is: first, copy X into two copies and place them on GPU0 and GPU1 respectively, cut A into two copies by column and place them on two GPUs, and perform parallel operations on GPU0 and GPU1. GeLU is an operation for a single element, so there is no need to merge Y1 and Y2, so that Y1 and Y2 are obtained on two GPUs. Next, cut B into two halves by row, perform matrix multiplication on two GPUs respectively, and perform an All-Reduce operation to obtain Z;

[0074] S42, the method of dividing the Attention module is: divide it into two equal parts according to the set number of self-attention heads, and put them on one GPU respectively. The division method here is: divide K, Q, and V into two parts by column, and put one part on each GPU. In the two GPUs, each GPU gets half of Y, which is equivalent to cutting Y in half vertically. According to the previous MLP method, B is cut into two halves by row and put on two GPUs. Then Z1 and Z2 can be obtained on the two GPUs respectively, so finally an All-Reduce operation is performed on Z1 and Z2;

[0075] S43. After the model is decomposed into tensors, the decomposed model and corresponding model parameters are stored in corresponding devices instead of storing all large model parameters.

[0076] Preferably, in S5, when the Nth batch of Attention-MLP is calculated, the calculation results of the previous batch (N-1) are saved, the CPU unloading start time is synchronized with the GPU synchronous communication, and the GPU continues to calculate after completing the synchronous communication. The CPU will unload the intermediate activation value communication in the GPU, and it will end before the next GPU synchronization. From the perspective of time, the CPU unloads the calculation results of the Nth batch of Attention from the GPU synchronously with the calculation of the N+1th batch of Attention, and the CPU unloads the calculation results of the Nth batch of MLP from the GPU synchronously with the calculation of the N+1th batch of MLP;

[0077] During forward calculation, in addition to unloading data from the GPU to save GPU memory, the CPU also transfers the data and model parameters required for GPU calculation to the GPU to ensure that the GPU can continue to run. Two parameter cache intervals are managed in the GPU memory, one for the model parameters being executed, and the CPU caches the model parameters of the next stage in the other parameter cache interval, which is completed during the communication time between the CPU and the GPU. During reverse transmission, CPU-GPU communication is also used, but the communication direction is opposite. The GPU obtains the intermediate activation values ​​saved in the forward calculation from the CPU, performs gradient calculations and optimizer updates, and also performs data synchronization at the synchronization point when updating parameters in the reverse direction. The CPU also starts data transmission at the same time, transferring the intermediate activation values ​​of the previous batch to the GPU and unloading the updated model parameters to the CPU.

[0078] Compared with the related art, the large language model training acceleration method provided by the present invention has the following advantages:

[0079] Beneficial effects:

[0080] The present invention provides a large language model training acceleration method, which can accurately evaluate the resource usage in large model training, further analyze to find the bottleneck in large model training, and help researchers propose large model acceleration training methods, which is of great significance for the analysis and optimization training of large models.

[0081] It solves the problems of unclear division principles and unbalanced load in traditional methods.

[0082] Data parallelism and communication parallelism are adopted to solve the storage problem of a large number of intermediate activation values ​​and the frequent GPU synchronous communication problem in tensor parallelism. Through data parallelism, the parallelism of the multi-head self-attention mechanism is improved. The combination of CPU-GPU communication and synchronous communication greatly improves the resource utilization of GPU and CPU, and further reduces the GPU's video memory occupancy. BRIEF DESCRIPTION OF THE DRAWINGS

[0083] Figure 1 This is a schematic diagram of the original table;

[0084] Figure 2 A schematic diagram of a preferred embodiment of a large language model training acceleration method provided by the present invention;

[0085] Figure 3 Schematic diagram for the division of large models;

[0086] Figure 4 This is the first schematic diagram of the Transformer tensor parallelization;

[0087] Figure 5The second schematic diagram of Transformer tensor parallelization;

[0088] Figure 6 A schematic diagram of the computation and communication timing in large model training;

[0089] Figure 7 A schematic diagram of the timing of large model training. DETAILED DESCRIPTION

[0090] The present invention will be further described below in conjunction with the accompanying drawings and implementation modes.

[0091] Please refer to Figure 2 , Figure 3 , Figure 4 , Figure 5 , Figure 6 and Figure 7 ,in, Figure 1 A schematic diagram of a preferred embodiment of a large language model training acceleration method provided by the present invention; Figure 3 Schematic diagram for the division of large models; Figure 4 This is the first schematic diagram of the Transformer tensor parallelization; Figure 5 The second schematic diagram of Transformer tensor parallelization; Figure 6 A schematic diagram of the computation and communication timing in large model training; Figure 7 A method for accelerating the training of a large language model includes the following steps:

[0092] S1. Mathematically model different network layers of the large model, analyze the parameters of the corresponding layers, and calculate the computational complexity of different layers: Model parameters refer to the number of parameters that can be learned and adjusted in a neural network or machine learning model. These parameters include weights and biases, which are continuously updated during the training process to optimize the performance of the model. The model computational complexity is expressed in FLOPs, which represents the number of floating-point operations and measures the size of the computational complexity.

[0093] S2. After mathematically modeling the network layer of the large model, it is necessary to count the parameters and calculation amount of the large model as a whole. For the parameters of the large model as a whole, the parameters of the AttentionMask layer, the Embedding layer, and the Decoder layer are ignored here. The Transformer model will be overlapped and arranged in the large model. As the main body of the large model, the parameters of the main body of the large model are listed here to represent the overall parameters of the large model, which is convenient for later analysis.

[0094] The specific formula is as follows:

[0095] Eliminate the parameters of the Embedding layer and the Decoder layer:

[0096] φ=n layer *(3*d model *d attn +d model *d attn +2*d model *d ff )=2*n layer *d model *(2*d attn +d ff ) Each Token

[0097] Calculation Amount:

[0098]

[0099] S3. Consider the parameters of computing hardware in reality, and combine the mathematical modeling analysis of the large model to divide the large model into different stages so that it can run stably on GPUs. Assuming that there are n A100 GPUs to run a stage of the large model, and each stage of the large model contains t Transformer layers, it is necessary to ensure that the memory usage and computing power of this stage of the large model do not exceed the theoretical upper limit of n GPUs (generally 80% of the peak value is taken as the upper limit). Let φ′ represent the parameter quantity of this sub-model, and determine the number of Transformer layers n in this sub-model. l ' ayer .

[0100] For memory usage, n A100 GPUs contain n*40G memory, and the memory of the sub-model includes parameters 2φ′, gradients 2φ′, and optimizer states 12φ′, a total of 16φ′:

[0101] 16φ′=16*2*d model *(2*d attn +d ff )*n l ' ayer ≤n×40×80%

[0102] For the peak computing, the peak computing capacity of each A100 GPU is 19.5 TFLOPS, and the computing capacity of the sub-model is C′=2*φ′+2*n l ' ayer *n ctx *d model <<19.5×80%TFLOPS;

[0103] S4. In the training of large models, tensor row optimization is used for the calculation of the Transformer model. Tensor parallelism can be decomposed into row parallelism and column parallelism. Then, tensor parallelism is applied to the Attention structure and MLP structure respectively.

[0104] S5. Utilize the synchronous communication time of All-reduce and use data CPU offloading to save the parameters of the previous stage and the intermediate calculation results, thereby reducing the required GPU memory. At the same time, manage the GPU video memory and transfer the model parameters of the next stage to the GPU video memory to improve the GPU computing utilization. During the forward calculation process, enable CPU offloading when the GPU performs synchronous communication, and save the intermediate activation values ​​of the previous round of calculation results on the CPU. When updating the backward parameters, use the synchronous communication time during reverse communication to transfer the intermediate activation values ​​saved on the CPU in the forward calculation to the GPU for calculation, and save the updated parameters on the CPU.

[0105] The analysis of the parameters and calculation amount stored in different network layers in S1 includes the following steps: S11. The parameter amount in the network layer mainly refers to the weight matrix in the network, and its number is determined by the matrix dimension, which is related to its input and output dimensions. Here, d is used uniformly. model Represents the dimension of its output, while the other dimensions are related to the output or batch size batch;

[0106] S12. The computational complexity of large models mainly lies in the calculation of parameter matrices and inputs in the network layer, which can be divided into two categories: addition and subtraction and multiplication and division.

[0107] The Transformer structure can be divided into Attention QKV layer, Attention Mask layer, AttentionProject layer, FeedForward layer, Residual connection layer, etc. Generally speaking, there will be an Embedding layer and a Decoder layer at the beginning or end of the main body of the large model to complete the mapping of input to vector and vector to Tokens.

[0108] The step of sequentially listing the parameter quantities and calculation quantities of different network layers in S12 includes the following steps: S121, Embedding layer

[0109] The word embedding matrix parameter size is: n vocab *d model ,

[0110] The context input vector size is: n ctx *d model ,

[0111] The word embedding matrix calculation size is: n vocab *n ctx *d model ;

[0112] S122, AttentionQKV layer

[0113] The parameter size of the AttentionQKV layer is: 3*n layer *d model *d attn ,

[0114] The computational size of the AttentionQKV layer is: 6*n layer *d model *d attn ;

[0115] S123, AttentionMask layer

[0116] The parameter size of the AttentionMask layer is: d attn ,

[0117] The calculation size of the AttentionMask layer is: 2*n layer *d model *d attn ;

[0118] S124, AttentionProject layer

[0119] The size of the AttentionProject layer parameters is: n layer *d model *d attn ,

[0120] The calculation size of the AttentionProject layer is: 2*n layer *d model *d attn ;

[0121] S125, Feedforward layer

[0122] The parameter size of the Feedforward layer is: 2*n layer *d model *d ff ,

[0123] The calculation amount of the Feedforward layer is: 4*n layer *d model *d ff ;

[0124] S126, Decoder layer

[0125] The size of the Decoder layer parameters is: n vocab *d model ,

[0126] The computational size of the decoder layer is: 2*n vocab *d model .

[0127] In S2, the memory usage in the large model training, including mainly the optimizer state, gradient and activation, is modeled and analyzed: first, the large model optimizer state and gradient are analyzed. These two memory usages are closely related to the parameter amount of the large model. In addition, they are also affected by the memory optimization strategy. First, assuming that the model parameter amount is φ, the relationship between the large model gradient and the optimizer state and the large model parameters is confirmed;

[0128] When the model is stored in FP32 data type, each element occupies 4 bytes, so the memory usage of the parameters is 4φ. The gradient is equal to the model parameters, and the memory usage is 4φ. The AdamW optimizer needs to save the parameters, first-order gradients, and second-order momentum during back propagation. Therefore, the memory usage of the optimizer state is 4φ (parameters) + 4φ (momentum) + 4φ (variance).

[0129] The activations here refer to all tensors calculated during the forward pass and needed in the backward pass;

[0130] The AttentionMask layer needs to complete the following operations:

[0131]

[0132] Therefore, it is necessary to store the intermediate values:

[0133] Q,K,V:[batch,n heads ,n ctx ,d attn ]

[0134] QK T :[batch,n heads ,n ctx ,n ctx ]

[0135]

[0136] The calculation formula of the Feedforward layer is: gelu (x attn W1)W2+x attn The intermediate values ​​that the Feedforward layer needs to store are:

[0137] x attn :[batch,n ctx ,d model ]

[0138] x attn W1:[batch,n ctx ,d ff ]

[0139] f gelu (x attn W1):[batch,n ctx ,d ff ]

[0140] f gelu (x attn W1)W2:[batch,n ctx ,d model ]

[0141] Dropout(f gelu (x attn W1)W2):[batch,n ctx ,d model ];

[0142] To summarize the above analysis of different layers, confirm the amount of parameters required to store for intermediate activation:

[0143] The parameter amount of the intermediate activation of the AttentionQKV layer is:

[0144] batch*n ctx *d model ;

[0145] The parameter amount of the intermediate activation of the AttentionMask layer is:

[0146] 4*batch*n heads *n ctx *d attn +3*batch*n heads *n ctx *n ctx ;

[0147] The number of parameters activated in the middle of the AttentionProject layer is:

[0148] 2*batch*n ctx *d model ;

[0149] The number of parameters activated in the middle of the Feedforward layer is:

[0150] 2*batch*n ctx *(d model +d ff )+batch*n ctx *d model ;

[0151] The number of parameters activated in the middle of the LayerNorm layer is:

[0152] 2*batch*n ctx *d model ;

[0153] The total number of intermediate activation parameters is:

[0154] C activation =n layer *batch*n ctx *(14*d model +4*d ff +8*n heads *d attn +5*n heads *n ctx ).

[0155] The S4 includes the following steps: S41, the splitting method of the MLP module is: first, copy X into two copies and place them on GPU0 and GPU1 respectively, cut A into two copies by column and place them on two GPUs, and perform parallel operations on GPU0 and GPU1. GeLU is an operation for a single element, so there is no need to merge Y1 and Y2, so that Y1 and Y2 are obtained on two GPUs. Next, cut B into two halves by row, perform matrix multiplication on two GPUs respectively, and perform an All-Reduce operation to obtain Z;

[0156] S42, the method of dividing the Attention module is: divide it into two equal parts according to the set number of self-attention heads, and put them on one GPU respectively. The division method here is: divide K, Q, and V into two parts by column, and put one part on each GPU. In the two GPUs, each GPU gets half of Y, which is equivalent to cutting Y in half vertically. According to the previous MLP method, B is cut into two halves by row and put on two GPUs. Then Z1 and Z2 can be obtained on the two GPUs respectively, so finally an All-Reduce operation is performed on Z1 and Z2;

[0157] S43. After the model is decomposed into tensors, the decomposed model and corresponding model parameters are stored in corresponding devices instead of storing all large model parameters.

[0158] In the S5, when the Nth batch of Attention-MLP is calculated, the calculation results of the previous batch (N-1) are saved, the CPU unloading startup time is synchronized with the GPU synchronous communication, and the GPU continues to calculate after completing the synchronous communication. The CPU will unload the intermediate activation value communication in the GPU, and it will end before the next GPU synchronization. From the perspective of time, the CPU unloads the calculation results of the Nth batch of Attention from the GPU synchronously with the calculation of the N+1th batch of Attention, and the CPU unloads the calculation results of the Nth batch of MLP from the GPU synchronously with the calculation of the N+1th batch of MLP;

[0159] During forward calculation, in addition to unloading data from the GPU to save GPU memory, the CPU also transfers the data and model parameters required for GPU calculation to the GPU to ensure that the GPU can continue to run. Two parameter cache intervals are managed in the GPU memory, one for the model parameters being executed, and the CPU caches the model parameters of the next stage in the other parameter cache interval, which is completed during the communication time between the CPU and the GPU. During reverse transmission, CPU-GPU communication is also used, but the communication direction is opposite. The GPU obtains the intermediate activation values ​​saved in the forward calculation from the CPU, performs gradient calculations and optimizer updates, and also performs data synchronization at the synchronization point when updating parameters in the reverse direction. The CPU also starts data transmission at the same time, transferring the intermediate activation values ​​of the previous batch to the GPU and unloading the updated model parameters to the CPU.

[0160] parameter index Quantification Model Layers <![CDATA[n layer ]]> 24 Residual layer output dimension <![CDATA[d model ]]> 2048 Feedforward layer output dimension <![CDATA[d ff ]]> 8192 Attention layer output dimension <![CDATA[d attn ]]> 2048 Number of attention heads <![CDATA[n heads ]]> 16 Vocabulary size <![CDATA[n vocab ]]> 130528 Context length <![CDATA[n ctx ]]> 2048 batch batch 2

[0161] Example:

[0162] This experiment trains the large model GPT-2 (Large) on two A6000 GPUs. The specific implementation process is as follows:

[0163] Taking the large model GPT-2 (Large) as an example, the large model parameter configuration is as follows:

[0164] According to S1, the large model mathematical modeling method proposed in the present invention is used to model each time of the large model;

[0165] According to S2, a quantitative analysis is performed on the large model ChatGLM. First, the number of parameters of the large model is analyzed. The number of model parameters after excluding the Embedding layer and the Decoder layer is:

[0166] φ=2*n layer *d model *(2*d attn +d ff )=1,409,400,832

[0167] Add the model parameters of the Embedding layer and the Decoder layer to φ, and the total number of model parameters N is:

[0168] N=φ+d model *(n vocab +n ctx )=1,680,916,480

[0169] Next, we analyze the computational complexity of each Token in the large model:

[0170] C forward =2*φ+2*n layer *n ctx *d model =3,020,128,256

[0171] Through the above analysis, it can be determined that the number of parameters of the large model ChatGLM is 6B, and the amount of calculation for each token is 3.02GFLOPS.

[0172] Based on the above quantification results, the memory usage of the large model ChatGLM in training is analyzed, which is specifically divided into model parameter quantity, model gradient, optimizer state, and intermediate activation value. Here, all elements are stored in FP32 format during training, and each element occupies 4 bytes. When backpropagation is used to update parameters, FP32 data format is used. The memory usage is as follows:

[0173] Model parameters: 4φ≈6G

[0174] Model gradient: 4φ≈6G

[0175] Optimizer state: parameters + momentum + variance = 24φ≈36G

[0176] Intermediate activation value: C activation ≈14G.

[0177] After the above quantitative analysis, we can conclude that the 48G video memory of A6000 GPU can store all the parameters of the large model GPT-2, but if you want to complete pre-training, there may still be a problem of video memory exceeding the limit; in order to support training, it is best to divide the model in parallel or optimize it.

[0178] According to S3, the present invention sets the model parallelism to MP=2 and divides the model into two stages, so that the existing hardware can meet the requirements of large model training. The memory usage of forward training in each stage is: (input data + parameters + intermediate activation values ​​+ part of optimizer state)*1 / 2≈23G<48G; the memory usage of reverse parameter update is: (optimizer state + gradient + intermediate activation)*1 / 2≈26G<48G.

[0179] By dividing the model into two stages, a single GPU can run one stage of a large model, which solves the problem of whether the large model can be trained, but does not solve the problem of accelerating the training of large models. Next, each stage is optimized and accelerated.

[0180] According to S4, tensor parallelism is applied in each stage. The specific parallel granularity also needs to be combined with the computing hardware. Tensor parallelism copies different parts of the model to different ranks, but different ranks will load the same data. Therefore, the granularity t of tensor parallelism is first determined to be equal to the number of GPUs n running each stage, set to t = n = 2; for each Transformer layer, the number of Attention heads is divided into two equal parts and placed on two GPUs respectively. Similarly, the parameters of the MLP layer are divided into two, and the corresponding model parameters are loaded to the corresponding devices; then the input X is copied to the two GPUs for operation. In this process, each device is equivalent to loading half of the model and model parameters, saving nearly half of the video memory.

[0181] According to S5, data parallelism copies the same model to different ranks, but different ranks will load different data and no longer perform data parallel division. Since there are only two A6000 GPUs and the tensor parallel granularity is t=n=2, the data parallel granularity is d=1, and data parallelism is no longer performed. However, due to the GPU video memory saved by tensor parallelism, the training batch of data can be increased to 4.

[0182] Experimental results, from the training time comparison, this large model training was carried out on two A6000GPU graphics cards, implementing the three methods of this paper (Ours), tensor parallelism (TP), and model parallelism (MP). Under the same conditions, the training was also 6 hours. The epoch number and loss value of the training model were compared:

[0183] Table 3 Comparison of different training methods

[0184] Method Epoch Loss TP (tp = 2, batch = 2) 94 4.54E+1 MP (mp=2, batch=2) 70 8.40E+1 Ours(mp=2, tp=2, batch=2) 131 8.91E+0

[0185] From the above experiments, it can be seen that model parallel training has the worst performance and the slowest training speed, but the model convergence is the same. Compared with model parallelism, pure tensor parallelism has better performance and faster training speed. Model convergence is related to the number of training iterations. The method in this paper has 131 iterations in the same time, but because this method sets batch=4, the number of iterations should theoretically be twice that of the TP method. In fact, each batch takes more time. This is due to the consumption caused by more communication time during training, but it performs best in model convergence.

[0186] Compared with the related art, the large language model training acceleration method provided by the present invention has the following advantages:

[0187] Beneficial effects:

[0188] The present invention provides a large language model training acceleration method, which can accurately evaluate the resource usage in large model training, further analyze to find the bottleneck in large model training, and help researchers propose large model acceleration training methods, which is of great significance for the analysis and optimization training of large models.

[0189] It solves the problems of unclear division principles and unbalanced load in traditional methods.

[0190] Data parallelism and communication parallelism are adopted to solve the storage problem of a large number of intermediate activation values ​​and the frequent GPU synchronous communication problem in tensor parallelism. Through data parallelism, the parallelism of the multi-head self-attention mechanism is improved. The combination of CPU-GPU communication and synchronous communication greatly improves the resource utilization of GPU and CPU, and further reduces the GPU's video memory occupancy.

[0191] The above descriptions are merely embodiments of the present invention and are not intended to limit the patent scope of the present invention. Any equivalent structure or equivalent process transformation made using the contents of the present invention specification and drawings, or directly or indirectly applied in other related technical fields, are also included in the patent protection scope of the present invention.

Claims

1. A large language model training acceleration method, characterized in that: The following steps are involved: S1. Mathematically model different network layers of the large model, analyze the parameters of the corresponding layers, and calculate the computational complexity of different layers: Model parameters refer to the number of parameters learned and adjusted in a neural network or machine learning model. These parameters include weights and biases, which are continuously updated during the training process to optimize the performance of the model. The model computational complexity is expressed in FLOPs, which represents the number of floating-point operations and measures the size of the computational complexity. S2. After mathematically modeling the network layer of the large model, it is necessary to count the parameters and calculation amount of the large model as a whole. For the parameters of the large model as a whole, the parameters of the Attention Mask layer, Embedding layer, and Decoder layer are ignored here. The Transformer model will be overlapped in the large model. As the main body of the large model, the parameters of the main body of the large model are listed here to represent the overall parameters of the large model, which is convenient for later analysis. The specific formula is as follows: Eliminate the parameters of the Embedding layer and the Decoder layer: φ=n layer *(3*d model *d attn +d model *d attn +2*d model *d ff )=2*n layer *d model *(2*d attn +d ff )Calculation amount for each Token: S3. Consider the parameters of computing hardware in reality and combine the mathematical modeling analysis of the large model to divide the large model into different stages so that it can run stably on GPUs. Assuming that there are n A100 GPUs to run a stage of the large model, and each stage of the large model contains t Transformer layers, it is necessary to ensure that the memory usage of this stage of the large model and the amount of calculation do not exceed the theoretical upper limit of n GPUs. φ′ represents the number of parameters of this sub-model and determines the number of Transformer layers n in this sub-model. l ' ayer , For memory usage, n A100 GPUs contain n*40G memory, and the memory of the sub-model includes parameters 2φ′, gradients 2φ′, and optimizer states 12φ′, a total of 16φ′: 16φ′=16*2*d model *(2*d attn +d ff )*n l ′ ayer ≤n×40×80% For the peak computing, the peak computing capacity of each A100 GPU is 19.5 TFLOPS, and the computing capacity of the sub-model is C′=2*φ′+2*n l ' ayer *n ctx *d model <<19.5×80%TFLOPS; S4. In the training of large models, tensor row optimization is used for the calculation of the Transformer model. Tensor parallelism can be decomposed into row parallelism and column parallelism. Then, tensor parallelism is applied to the Attention structure and MLP structure respectively. S5. Utilize the synchronous communication time of All-reduce and use data CPU offloading to save the parameters of the previous stage and the intermediate calculation results, thereby reducing the required GPU memory. At the same time, manage the GPU video memory and transfer the model parameters of the next stage to the GPU video memory to improve the GPU computing utilization. During the forward calculation process, enable CPU offloading when the GPU performs synchronous communication, and save the intermediate activation values ​​of the previous round of calculation results on the CPU. When updating the backward parameters, use the synchronous communication time during reverse communication to transfer the intermediate activation values ​​saved on the CPU in the forward calculation to the GPU for calculation, and save the updated parameters on the CPU.

2. A large language model training acceleration method according to claim 1, characterized in that: The analysis of the stored parameters and computational load of different network layers in S1 includes the following steps: S11. The number of parameters in the network layer mainly refers to the weight matrix in the network. Its number is determined by the matrix dimension, which is related to its input and output dimensions. Here we use d model Represents the dimension of its output, while the other dimensions are related to the output or batch size batch; S12. The computational complexity of large models mainly lies in the calculation of parameter matrices and inputs in the network layer, which can be divided into two categories: addition and subtraction and multiplication and division. The Transformer structure is divided into the Attention QKV layer, the Attention Mask layer, the Attention Project layer, the FeedForward layer, and the residual connection layer. The Embedding layer and the Decoder layer will be at the beginning or end of the main body of the large model to complete the mapping of input to vector and vector to Tokens.

3. A large language model training acceleration method according to claim 2, characterized in that: The step of sequentially listing the parameter quantities and calculation quantities of different network layers in S12 includes the following steps: S121, Embedding layer The word embedding matrix parameter size is: n vocab *d model , The context input vector size is: n ctx *d model , The word embedding matrix calculation size is: n vocab *n ctx *d model ; S122, Attention QKV layer The parameter size of the Attention QKV layer is: 3*n layer *d model *d attn , The calculation size of the Attention QKV layer is: 6*n layer *d model *d attn ; S123, Attention Mask layer The parameter size of the Attention Mask layer is: d attn , The calculation size of the Attention Mask layer is: 2*n layer *d model *d attn ; S124, Attention Project layer The size of the Attention Project layer parameters is: n layer *d model *d attn , The calculation size of the Attention Project layer is: 2*n layer *d model *d attn ; S125, Feedforward layer The parameter size of the Feedforward layer is: 2*n layer *d model *d ff , The calculation amount of the Feedforward layer is: 4*n layer *d model *d ff ; S126, Decoder layer The size of the Decoder layer parameters is: n vocab *d model , The computational size of the decoder layer is: 2*n vocab *d model .

4. A large language model training acceleration method according to claim 1, characterized in that: In S2, the memory usage in the large model training, including mainly the optimizer state, gradient and activation, is modeled and analyzed: first, the large model optimizer state and gradient are analyzed. These two memory usages are closely related to the parameter amount of the large model. In addition, they are also affected by the memory optimization strategy. First, assuming that the model parameter amount is φ, the relationship between the large model gradient and the optimizer state and the large model parameters is confirmed; When the model is stored in FP32 data type, each element occupies 4 bytes, so the memory usage of the parameters is 4φ. The gradient is equal to the model parameters, and the memory usage is 4φ. The AdamW optimizer needs to save the parameters, first-order gradients, and second-order momentum during back propagation. Therefore, the memory usage of the optimizer state is 4φ (parameters) + 4φ (momentum) + 4φ (variance). The activations here refer to all tensors calculated during the forward pass and needed in the backward pass; The Attention Mask layer needs to complete the following operations: Therefore, it is necessary to store the intermediate values: Q,K,V:[batch,n heads ,n ctx ,d attn ] QK T :[batch,n heads ,n ctx ,n ctx ] The calculation formula of the Feedforward layer is: gelu (x attn W1)W2+x attn The intermediate values ​​that the Feedforward layer needs to store are: x attn :[batch,n ctx ,d model ] x attn W1:[batch,n ctx ,d ff ] f gelu (x attn W1):[batch,n ctx ,d ff ] f gelu (x attn W1)W2:[batch,n ctx ,d model ] Dropout(f gelu (x attn W1)W2):[batch,n ctx ,d model ]; To summarize the above analysis of different layers, confirm the amount of parameters required to store for intermediate activation: The parameter quantity of the intermediate activation of the Attention QKV layer is: batch*n ctx *d model ; The number of parameters activated in the middle of the Attention Mask layer is: 4*batch*n heads *n ctx *d attn +3*batch*n heads *n ctx *n ctx ; The number of parameters activated in the middle of the Attention Project layer is: 2*batch*n ctx *d model ; The number of parameters activated in the middle of the Feedforward layer is: 2*batch*n ctx *(d model +d ff )+batch*n ctx *d model ; The number of parameters activated in the middle of the LayerNorm layer is: 2*batch*n ctx *d model ; The total number of intermediate activation parameters is: C activation =n layer *batch*n ctx *(14*d model +4*d ff +8*n heads *d attn +5*n heads *n ctx )。 5. A large language model training acceleration method according to claim 1, characterized in that: The S4 includes the following steps: S41, the splitting method of the MLP module is: first, X is copied into two copies and placed on GPU0 and GPU1 respectively, A is cut into two copies by column and placed on two GPUs, and parallel operations are performed on GPU0 and GPU1. GeLU is an operation for a single element, so there is no need to merge Y1 and Y2, so that Y1 and Y2 are obtained on the two GPUs, then B is cut into two halves by row, matrix multiplication is performed on the two GPUs respectively, and Z can be obtained by performing an All-Reduce operation; S42, the method of dividing the Attention module is: divide it into two equal parts according to the set number of self-attention heads, and put them on one GPU respectively. The division method here is: divide K, Q, and V into two parts by column, and put one part on each GPU. In the two GPUs, each GPU gets half of Y, which is equivalent to cutting Y in half vertically. According to the previous MLP method, B is cut into two halves by row and put on two GPUs. Then Z1 and Z2 can be obtained on the two GPUs respectively, so finally an All-Reduce operation is performed on Z1 and Z2; S43. After the model is decomposed into tensors, the decomposed model and corresponding model parameters are stored in corresponding devices instead of storing all large model parameters.

6. A large language model training acceleration method according to claim 1, characterized in that: In the S5, when the Nth batch of Attention-MLP is calculated, the calculation results of the previous batch (N-1) are saved, the CPU unloading startup time is synchronized with the GPU synchronous communication, and the GPU continues to calculate after completing the synchronous communication. The CPU will unload the intermediate activation value communication in the GPU, and it will end before the next GPU synchronization. From the perspective of time, the CPU unloads the calculation results of the Nth batch of Attention from the GPU synchronously with the calculation of the N+1th batch of Attention, and the CPU unloads the calculation results of the Nth batch of MLP from the GPU synchronously with the calculation of the N+1th batch of MLP; During forward calculation, in addition to unloading data from the GPU to save GPU memory, the CPU also transfers the data and model parameters required for GPU calculation to the GPU to ensure that the GPU can continue to run. Two parameter cache intervals are managed in the GPU memory, one for the model parameters being executed, and the CPU caches the model parameters of the next stage in the other parameter cache interval, which is completed during the communication time between the CPU and the GPU. During reverse transmission, CPU-GPU communication is also used, but the communication direction is opposite. The GPU obtains the intermediate activation values ​​saved in the forward calculation from the CPU, performs gradient calculations and optimizer updates, and also performs data synchronization at the synchronization point when updating parameters in the reverse direction. The CPU also starts data transmission at the same time, transferring the intermediate activation values ​​of the previous batch to the GPU and unloading the updated model parameters to the CPU.

Citation Information

Patent Citations

  • Parallel training acceleration method and system for large model parameter partitioning

    CN117744838A

  • Supercritical unit wide-load rapid modeling method combined with Swinin-Transformer

    CN117787077A