A low-cost model training method based on block coordinate descent method

Abstract

The application belongs to the technical field of artificial intelligence and machine learning, and discloses a low-cost model training method based on block coordinate descent method, and the specific technical scheme is as follows: model parameters are divided into multiple sub-blocks, only one sub-block is activated to participate in forward propagation, backward propagation and parameter updating in each round of training, and the remaining sub-blocks remain in a frozen state, which significantly reduces the memory occupation and training calculation overhead, the frozen parameters do not participate in gradient calculation and state storage, and the resource utilization efficiency is further improved; the method can complete effective pre-training of a large model on a consumer-level GPU cluster, and the precision maintaining ability and resource saving effect are verified on multiple public data sets and mainstream models, and the method has good universality and engineering deployment value.

Description

Technical Field

[0001] This invention belongs to the field of artificial intelligence and machine learning technology, and in particular relates to an optimization method for training deep learning models in low-resource environments. Background Technology

[0002] As the scale of pre-trained models continues to expand, the GPU memory requirements and economic costs for training Large Language Models (LLMs) are growing exponentially, becoming a significant hurdle for small and medium-sized research teams and enterprises. During full parameter pre-training, training a model with W parameters to process 1024 tokens using the Adam optimizer requires storing gradients (W), optimizer states (3W), and activation values ​​(approximately 7W), totaling about 57W of GPU memory. Even with recomputation techniques, the memory requirement remains as high as 5W. For example, the total pre-training cost for a 20B model can reach millions of dollars, severely limiting the widespread adoption and application of large language model technology.

[0003] Traditional 3D distributed parallel training overcomes the limitations of single-card memory by distributing model parameters and optimizer states across multiple GPUs, but its cost is becoming increasingly prominent. On the one hand, as the model size increases, the frequency of communication between nodes and the data throughput increase dramatically, and communication overhead consumes a lot of training time; on the other hand, the computing power cost of continuously running dozens or even hundreds of high-end A100 / A800 GPUs is extremely high, driving the total economic investment to continue to rise.

[0004] In contrast, high-performance clusters with high cost-performance ratios, such as the RTX 4090 or DCU7000, have gradually emerged. However, their single-card bandwidth and memory are relatively low, making it difficult to support large-scale full parameter pre-training, thus creating a serious gap between cost and performance (see Table 1).

[0005] Table 1. Comparison of performance parameters of different GPU hardware platforms

[0006]

[0007] However, there is currently no mature algorithm that can make full use of these low-cost clusters under resource-constrained conditions. Common methods such as DeepSpeed's OffLoad can train models on limited GPUs, but due to the frequent migration of parameters, optimizer states, activation values ​​and gradients between GPU and host memory, it causes serious performance degradation and increased time costs, and may not necessarily reduce the overall cost. Summary of the Invention

[0008] To address the technical challenge of training large language models under resource-constrained environments, this invention provides a low-cost model training method based on block coordinate descent. This method can effectively improve training efficiency on resource-constrained devices, reducing memory usage and economic costs while ensuring the training accuracy and convergence of the model.

[0009] To achieve the above objectives, the technical solution adopted by this invention is: a low-cost model training method based on block coordinate descent, comprising the following steps:

[0010] Step 1: Obtain the neural network model to be pre-trained, the training dataset, and the convergence criterion. According to the network layer order or the boundary of the computation module of the neural network model, divide all the model parameters of the neural network model into K parameter sub-blocks, where K is an integer greater than or equal to 2.

[0011] Step 2: According to the network layer order of the neural network model, divide the K parameter sub-blocks into multiple pipelined parallel stages, and deploy different pipelined parallel stages on multiple graphics processors. Transmit the activation data generated by forward propagation and the gradient data generated by backpropagation between adjacent pipelined parallel stages.

[0012] Step 3: Following the preset loop order, select the currently active parameter sub-block from the K parameter sub-blocks, set the currently active parameter sub-block to the trainable state, and set the other parameter sub-blocks to the frozen state.

[0013] Step 4: Create an optimizer for the currently active parameter sub-block and allocate optimizer state storage space. Perform forward and backward propagation on the neural network model using the training dataset. During backward propagation, calculate the parameter gradient of the currently active parameter sub-block and update the model parameters of the currently active parameter sub-block according to the parameter gradient. The parameter sub-blocks in the frozen state keep their model parameters unchanged. Do not allocate parameter gradient storage space or optimizer state storage space for the frozen parameter sub-blocks. Calculate the input gradient required to pass the loss function gradient to the currently active parameter sub-block according to the gradient propagation needs.

[0014] Step 5: Continuously train the current active parameter sub-block until the current active parameter sub-block meets the preset local convergence condition, release the optimizer state corresponding to the current active parameter sub-block, switch the current active parameter sub-block to the frozen state, and select the next parameter sub-block as the new current active parameter sub-block according to the preset loop order, recreate the optimizer for the new current active parameter sub-block and allocate optimizer state storage space.

[0015] Step 6: Repeat steps 3 to 5 to activate and update the K parameter sub-blocks in sequence until the neural network model meets the preset global convergence condition, and obtain the pre-trained neural network model.

[0016] The parameter sub-blocks are divided according to the hierarchical structure of the neural network model. Each parameter sub-block includes a network layer, multiple consecutive network layers, or a computation module based on network layers.

[0017] When the neural network model is a language model based on the Transformer structure, the parameter sub-block includes at least one of one or more Transformer layers, attention modules, feedforward neural network modules, embedding layers, and output layers.

[0018] When the neural network model is a convolutional neural network, the parameter sub-block includes at least one of one or more convolutional layers, residual modules, and fully connected layers.

[0019] The preset loop order is to select K parameter sub-blocks sequentially according to their numbers;

[0020] Let k be the number of times the current parameter sub-block has been switched, and let the number of the currently active parameter sub-block be k. The number of the current active parameter sub-block satisfy:

[0021] ;

[0022] Where mod represents the modulo operation;

[0023] After the current active parameter sub-block satisfies the local convergence condition, let And determine the next parameter sub-block based on the updated k.

[0024] The optimizer can be any one of the following: stochastic gradient descent optimizer, Adam optimizer, or AdamW optimizer.

[0025] During the training of the current active parameter sub-block, only the optimizer state corresponding to the current active parameter sub-block is saved;

[0026] When the current active parameter sub-block is switched to the frozen state, release the optimizer state corresponding to the current active parameter sub-block;

[0027] When a parameter sub-block in a frozen state is switched to a new currently active parameter sub-block, the corresponding optimizer state is reinitialized based on the model parameters of the new currently active parameter sub-block.

[0028] During backpropagation, for the parameter sub-block in the frozen state, the parameter gradients corresponding to the model parameters in the frozen parameter sub-block are not calculated, the model parameters in the frozen parameter sub-block are not updated, and the optimizer state corresponding to the frozen parameter sub-block is not saved.

[0029] For a frozen state parameter sub-block that is located after the current active parameter sub-block in the forward propagation direction, calculate the input gradient of the frozen state parameter sub-block and pass the input gradient step by step to the current active parameter sub-block along the back propagation direction.

[0030] For frozen parameter sub-blocks that are located before the currently active parameter sub-block in the forward propagation direction, backpropagation calculations for updating model parameters are not performed.

[0031] The number of parameter sub-blocks K is determined based on the available video memory capacity of the training device, the total number of parameters of the neural network model, and the peak video memory usage required for intermediate data during training.

[0032] When using FP32 precision to store model parameters and gradients, with the same number of parameters in each parameter sub-block, and the optimizer states including first-order moment states and second-order moment states, the peak memory usage of the neural network model during training using the block coordinate descent method is as follows: satisfy:

[0033] ;

[0034] And the number K of parameter sub-blocks satisfies:

[0035] ;

[0036] Peak memory usage during full parameter training It can be represented as:

[0037] .

[0038] in, This represents the total number of model parameters in a neural network model. This represents the storage space occupied by all model parameters. This represents the storage space occupied by the parameter gradients and optimizer state of the currently active parameter sub-block. This indicates the peak memory usage generated by intermediate data during the training process. This indicates the available video memory capacity of the training device, with each unit of video memory usage in bytes.

[0039] A complete parameter sub-block cyclic training process includes K parameter sub-block training phases. In each of the K parameter sub-block training phases, each parameter sub-block is set as the currently active parameter sub-block at least once.

[0040] After completing a full parameter sub-block iterative training process, all model parameters of the neural network model are updated at least once;

[0041] If the neural network model has not yet met the global convergence condition, continue to execute the next complete parameter sub-block loop training process.

[0042] Each parallel pipeline stage includes one or more consecutive parameter sub-blocks;

[0043] In the pipelined parallel phase that includes the current active parameter sub-block, the parameter gradient calculation, optimizer state processing, and model parameter update of the current active parameter sub-block are performed.

[0044] In the pipelined parallel phase that does not include the currently active parameter sub-block, the model parameters of the corresponding parameter sub-block remain unchanged, and the parameter gradient calculation, optimizer state processing and model parameter update for the frozen parameter sub-block are skipped. Only the complete forward propagation and the backpropagation calculation required to pass the loss function gradient to the currently active parameter sub-block are performed.

[0045] Among them, the neural network model is a large language model, which is pre-trained in a graphics processor cluster where the memory capacity of a single graphics processor is insufficient to store the complete training state.

[0046] Regarding training convergence and stability, this invention ensures global convergence based on block coordinate descent theory and maintains the stability of the training process through a phased update strategy. Based on this method, its significant cost-effectiveness has been verified in the following two low-cost scenarios:

[0047] When the number of parameter sub-blocks K equals 3, the memory usage of the model state in this invention is: Compared to training with complete parameters reduce Experimental results show that the method of this invention reduces peak memory usage by more than [percentage missing] on 1B, 2B, and 7B models. It can complete training even in configurations where full-parameter training results in memory overflow.

[0048] This invention has a lower cost with the same hardware: taking the 7B model as an example, in In a cluster, the average training cost is reduced to approximately With an RTX 4090 cluster, the cost is reduced to approximately For the 20B model, the theoretically calculated training cost is reduced to [amount missing]. and .

[0049] Cross-device migration capability: Using the method of this invention, migration capabilities that were originally limited to... Full-parameter pre-training implemented on the cluster can be seamlessly migrated to clusters with smaller memory but lower operating costs, such as RTX 4090, saving up to three-quarters of the computational overhead.

[0050] The model accuracy remains consistent: In any of the above scenarios, the BCD training method of this invention can achieve the same model accuracy as traditional full-parameter training, while the training process fluctuates less and the economic cost is significantly reduced. Attached Figure Description

[0051] Figure 1 This is a schematic diagram of the model divided into layers.

[0052] Figure 2 This is a graph showing the trend of perplexity (PPL) as a function of the number of training rounds under different training methods.

[0053] Figure 3 This chart compares GPU memory usage under different model sizes and training methods. Detailed Implementation

[0054] To make the technical problems to be solved, the technical solutions, and the beneficial effects of the present invention clearer, the present invention will be further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are merely illustrative of the present invention and are not intended to limit the present invention.

[0055] This embodiment employs block coordinate descent to perform full-parameter pre-training of the neural network model. First, the model parameters are divided into K parameter sub-blocks according to the network layer order or the boundaries of computational modules, where K is an integer greater than or equal to 2. Each parameter sub-block can include one network layer, multiple consecutive network layers, or a computational module composed of network layers. For language models based on the Transformer architecture, the sub-blocks can be divided according to one or more Transformer layers, or according to attention modules, feedforward neural network modules, embedding layers, or output layers; for ResNet models, a residual module can be used as a parameter sub-block, such as... Figure 1 As shown.

[0056] Based on the above parameter sub-block partitioning method, this embodiment uses the block coordinate descent method for model training. The specific training process is shown in the table below:

[0057]

[0058] At the start of training, a parameter sub-block is selected as the currently active parameter sub-block according to a preset loop order and set to the trainable state, while the remaining parameter sub-blocks are set to the frozen state. Let k be the number of times the current parameter sub-block is switched, and let the number of the currently active parameter sub-block be . ,but:

[0059] ;

[0060] Here, the initial value of k is 0, and mod represents the modulo operation.

[0061] After determining the current active parameter sub-block, an SGD, Adam, or AdamW optimizer is created for that sub-block, and the storage space required for the parameter gradients and optimizer states is allocated. The model performs a full forward propagation, and during the backpropagation, only the parameter gradients of the current active parameter sub-block are calculated, and its model parameters are updated based on these gradients.

[0062] Frozen parameter sub-blocks retain their model parameters unchanged; their parameter gradients are not calculated, and their corresponding optimizer states are not saved. When a frozen parameter sub-block follows the currently active parameter sub-block, its input gradient is still calculated to pass the loss function gradient to the currently active parameter sub-block. When a frozen parameter sub-block precedes the currently active parameter sub-block, backpropagation calculations for parameter updates are not performed.

[0063] After the current active parameter sub-block is trained to meet the preset local convergence condition, its optimizer state is released and it is switched to a frozen state; then let The next parameter sub-block is selected as the new active parameter sub-block, and the optimizer is recreated and its state is initialized. The parameter sub-blocks are trained in this manner until all parameter sub-blocks have completed one round of updates. If the model has not yet met the preset global convergence condition, the next round of parameter sub-block training continues; if the global convergence condition is met, training ends and the model is output.

[0064] The number K of parameter sub-blocks can be determined based on the number of model parameters and the available GPU memory capacity of the training device. Under the conditions of using FP32 precision, identical parameter numbers in each parameter sub-block, and activation recalculation, the peak GPU memory usage of the method of this invention can be expressed as:

[0065] ;

[0066] in, This indicates the total number of model parameters. This represents the total video memory used by all model parameters. This represents the memory space occupied by the parameter gradients and optimizer state of the currently active parameter sub-block. This represents the peak instantaneous activation memory usage generated during activation recalculation. The number of parameter sub-blocks, K, should ensure that the peak memory usage does not exceed the available memory capacity of the training device. When K equals 3, the memory usage of the model parameters, parameter gradients, and optimizer states corresponding to the model states is determined by the memory usage of the fully trained parameters. Reduced to .

[0067] In a single-GPU training scenario, as long as the GPU can store the complete model parameters, the number of parameters in the current active parameter sub-block can be determined based on the remaining GPU memory. Reducing the number of parameters in the active parameter sub-block can further reduce the GPU memory usage for parameter gradients and optimizer states, but it will increase the training time required to complete all parameter updates.

[0068] In multi-GPU training scenarios, the model is divided into multiple pipelined parallel stages according to the network layer order, and different pipelined parallel stages are deployed on different GPUs. The pipelined parallel stage containing the currently active parameter sub-block performs parameter gradient calculation, optimizer state processing, and model parameter updates; other pipelined parallel stages keep the model parameters unchanged and only perform model forward propagation and backpropagation calculations required to pass gradients to the currently active parameter sub-block. Since both parameter sub-blocks and pipelined parallel stages are divided according to network layers, parameter freezing, unfreezing, and switching can be completed at the boundaries of network layers or computation modules, reducing the impact on the original model computation structure.

[0069] The experimental environment consists of a server node equipped with 8 A100 servers:

[0070] One server node equipped with 8 A800 GPUs, and four server nodes equipped with 8 RTX 4090 GPUs each. All servers communicate via... High-speed network interconnection.

[0071] In the effectiveness experiment, the present invention respectively... and Training on dataset and The model was trained on the WikiText and WebText2 datasets using a parameter set with 0.15B. Models. In addition, larger-scale models, including those with 2B and 7B parameters, were trained. (abbreviation) , ) and LLaMA (abbreviated) , The datasets used include Wikipedia, Alpaca, and SlimPajama, all of which have been publicly released through the ModelScope community.

[0072] In the economic cost experiment, this invention demonstrates the actual cost of training 2B and 7B models; and by combining single-round training speed and The multiplier, calculated to account for the theoretical economic cost of training multiple models, includes:

[0073] GPT-2 series: 1.6B, 5.4B, 10B, 20B;

[0074] LLaMA series: 1.5B, 6B, 11B, 23B.

[0075] In all BCD experiments, the model was uniformly divided into three sub-models, meaning that only one-third of all parameters were updated in each round.

[0076] In terms of algorithm configuration, the experimental models included the ResNet series and 0.15B. Models, where the ResNet model is Training on the dataset, The model was trained on WikiText and WebText2. The optimizer settings were as follows: SGD learning rate of 0.1, momentum of 0.9, and weight decay. Adam's learning rate is Weight decay The remaining parameters are default; for 2B and 7B... Compared to the LLaMA model, Adam's learning rate is... Everything else remains the same. To fairly compare model convergence speeds, the training task uses only the minimum number of GPUs required to complete training. For example, full parameter training of the 7B model uses two A800 GPUs, the 2B model uses three RTX 4090 GPUs, and the 0.15B model uses one RTX 4090 GPU; while BCD training uses one A800, two RTX 4090s, and one RTX 4090 respectively. Regarding offload configuration, Stage 1 is enabled, offloading the optimizer state to the CPU (device set to cpu, ...). (If true), communication computation overlap and continuous gradient storage are enabled simultaneously. The 0.15B model can be fully loaded into GPU memory, so no Offload is required; the 2B model requires Offload to be enabled under dual-card conditions, but not under triple-card conditions, so the experiment chose the dual-card Offload scheme.

[0077] The effectiveness experiments show that, in small-scale model experiments, using Adam and SGD optimizers for training, and comparing BCD and full parameter update methods, from the perspective of the loss function, the BCD method outperforms full parameter training in most cases, with a more stable convergence process, thus verifying the conclusions in previous literature. Specific results are shown in Tables 2 and 3. Table 2 shows the accuracy and loss performance of ResNet and 0.15BGPT-2 models under different optimizers and training methods, while Table 3 shows the PPL and loss values ​​of various large-scale models on different datasets. Figure 3The study further demonstrated the trend of PPL change over training time. Under various training strategies, the BCD method showed a smoother convergence process and better results, verifying its training stability.

[0078] Table 2. Comparison of training results using different optimizers and training strategies on small models.

[0079]

[0080] Table 3. Performance comparison results of full parameter training and the BCD method of this invention on large models.

[0081]

[0082] Regarding memory usage, the method of this invention significantly reduces GPU memory consumption at different parameter scales such as 1B, 2B, and 7B. Compared to full-parameter training, peak memory usage is reduced by more than 45%, and the method of this invention can run stably even when memory overflow (OOM) occurs during full-parameter training, fully demonstrating its good resource adaptability.

[0083] Regarding the number of training rounds, the ratio of the number of rounds for the BCD method to the number of rounds for the full parameter method is defined as the BF multiple. Experiments show that when using the Adam optimizer, the average BF multiple is 1.39, and the worst case is 2.89. Although the number of rounds is slightly more, the overall training cost is still lower because less resources are used per round.

[0084] In the economic cost experiment, the training costs on the A100 ($1.20 / hour), A800 ($0.69 / hour), and RTX 4090 ($0.29 / hour) platforms were compared. As shown in Table 4, when training the L-7B model, the method of this invention reduced the cost from $43.74 to $31.22 on the A100 platform, from $26.30 to $18.74 on the A800 platform, and from $3.37 to $3.29 on the RTX 4090 platform. When training the G-7B model, the training cost on the Wiki dataset decreased from $70.08 to $41.36, and the training cost on the Alpaca dataset decreased from $18.49 to $12.21 and from $26.39 to $17.63, respectively, demonstrating a significant cost advantage.

[0085] Table 4. Comparison of actual training costs of the present invention and traditional training methods on different GPU platforms.

[0086]

[0087] Regarding theoretical costs, Tables 5 to 7 compare the training costs of using the BCD method on the RTX 4090 platform versus the full-parameter method on the RTX 4090, A100, and A800 platforms at different model sizes. Table 5 shows that, compared to the same platform (4090), the average cost saving exceeds 50%; Table 6 shows that, compared to training on the A100 platform, the average cost saving is as high as 80%; and Table 7 shows that, compared to training on the A800 platform, the average cost saving is approximately 67%. These results indicate that the BCD method is more advantageous in cluster environments with limited communication bandwidth and GPU memory, and is particularly suitable for training large models on consumer-grade GPUs.

[0088] Table 5. Training cost savings compared to full-parameter training on the same platform using the BCD method.

[0089]

[0090] Table 6. Training cost savings compared to full-parameter training of the A100 cluster using the BCD method.

[0091]

[0092] Table 7. Training cost savings compared to full-parameter training on the A800 cluster using the BCD method.

[0093]

[0094] In summary, the block coordinate descent-based model training method proposed in this invention effectively reduces memory usage and training costs while maintaining model performance, and improves training stability and hardware compatibility. It is suitable for efficient training of medium and large-scale neural network models under various computing resource conditions.

[0095] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions and improvements made within the spirit and principles of the present invention should be included within the scope of the present invention.

Claims

1. A low-cost model training method based on block coordinate descent, characterized in that, Includes the following steps: Step 1: Obtain the neural network model to be pre-trained, the training dataset, and the convergence criterion. Divide all the model parameters of the neural network model into K parameter sub-blocks according to the network layer order or the boundary of the computation module. Step 2: According to the network layer order of the neural network model, divide the K parameter sub-blocks into multiple pipelined parallel stages, and deploy different pipelined parallel stages on multiple graphics processors. Transmit the activation data generated by forward propagation and the gradient data generated by backpropagation between adjacent pipelined parallel stages. Step 3: Following the preset loop order, select the currently active parameter sub-block from the K parameter sub-blocks, set the currently active parameter sub-block to the trainable state, and set the other parameter sub-blocks to the frozen state. Step 4: Create an optimizer for the currently active parameter sub-block and allocate optimizer state storage space. Perform forward and backward propagation on the neural network model using the training dataset. During backward propagation, calculate the parameter gradient of the currently active parameter sub-block and update the model parameters of the currently active parameter sub-block according to the parameter gradient. The parameter sub-blocks in the frozen state keep their model parameters unchanged. Do not allocate parameter gradient storage space or optimizer state storage space for the frozen parameter sub-blocks. Calculate the input gradient required to pass the loss function gradient to the currently active parameter sub-block according to the gradient propagation needs. Step 5: Continuously train the current active parameter sub-block until the current active parameter sub-block meets the preset local convergence condition, release the optimizer state corresponding to the current active parameter sub-block, switch the current active parameter sub-block to the frozen state, and select the next parameter sub-block as the new current active parameter sub-block according to the preset loop order, recreate the optimizer for the new current active parameter sub-block and allocate optimizer state storage space. Step 6: Repeat steps 3 to 5 to activate and update the K parameter sub-blocks in sequence until the neural network model meets the preset global convergence condition, and obtain the pre-trained neural network model.

2. The low-cost model training method based on block coordinate descent according to claim 1, characterized in that, The parameter sub-blocks are divided according to the hierarchical structure of the neural network model. Each parameter sub-block includes a network layer, multiple consecutive network layers, or a computation module based on network layers. When the neural network model is a language model based on the Transformer structure, the parameter sub-block includes at least one of one or more Transformer layers, attention modules, feedforward neural network modules, embedding layers, and output layers. When the neural network model is a convolutional neural network, the parameter sub-block includes at least one of one or more convolutional layers, residual modules, and fully connected layers.

3. The low-cost model training method based on block coordinate descent according to claim 1, characterized in that, The preset loop order is to select the K parameter sub-blocks sequentially according to their numbers; Let k be the number of times the current parameter sub-block has been switched, and let the number of the currently active parameter sub-block be k. The number of the current active parameter sub-block satisfy: ； Where mod represents the modulo operation; After the current active parameter sub-block satisfies the local convergence condition, let And determine the next parameter sub-block based on the updated k.

4. The low-cost model training method based on block coordinate descent according to claim 1, characterized in that, The optimizer can be any one of the stochastic gradient descent optimizer, Adam optimizer, or AdamW optimizer. During the training of the current active parameter sub-block, only the optimizer state corresponding to the current active parameter sub-block is saved; When the current active parameter sub-block is switched to the frozen state, release the optimizer state corresponding to the current active parameter sub-block; When a parameter sub-block in a frozen state is switched to a new currently active parameter sub-block, the corresponding optimizer state is reinitialized based on the model parameters of the new currently active parameter sub-block.

5. The low-cost model training method based on block coordinate descent according to claim 1, characterized in that, During backpropagation, for the parameter sub-block in the frozen state, the parameter gradients corresponding to the model parameters in the frozen parameter sub-block are not calculated, the model parameters in the frozen parameter sub-block are not updated, and the optimizer state corresponding to the frozen parameter sub-block is not saved. For the frozen state parameter sub-block located after the current active parameter sub-block in the forward propagation direction, calculate the input gradient of the frozen state parameter sub-block and pass the input gradient step by step to the current active parameter sub-block along the back propagation direction; For frozen parameter sub-blocks located in the forward propagation direction before the currently active parameter sub-block, backpropagation calculations for updating model parameters are not performed.

6. The low-cost model training method based on block coordinate descent according to claim 1, characterized in that, The number of parameter sub-blocks K is determined based on the available video memory capacity of the training device, the total number of parameters of the neural network model, and the peak video memory usage required for intermediate data during training. Given that the model parameters and parameter gradients are stored, each parameter sub-block has the same number of parameters, and the optimizer states include first-order moment states and second-order moment states, the peak GPU memory usage M_BCD of the neural network model trained using the block coordinate descent method satisfies the following: ； And the number K of parameter sub-blocks satisfies: ； in, This represents the total number of model parameters in a neural network model. This represents the storage space occupied by all model parameters. This represents the storage space occupied by the parameter gradients and optimizer state of the currently active parameter sub-block. This indicates the peak memory usage generated by intermediate data during the training process. This indicates the available video memory capacity of the training device.

7. The low-cost model training method based on block coordinate descent according to claim 1, characterized in that, A complete parameter sub-block iterative training process includes K parameter sub-block training phases. In each of the K parameter sub-block training phases, each parameter sub-block is set as the currently active parameter sub-block at least once. After completing a full parameter sub-block iterative training process, all model parameters of the neural network model are updated at least once; If the neural network model has not yet met the global convergence condition, continue to execute the next complete parameter sub-block loop training process.

8. The low-cost model training method based on block coordinate descent according to claim 1, characterized in that, Each pipeline parallel stage includes one or more consecutive parameter sub-blocks; In the pipelined parallel phase that includes the current active parameter sub-block, the parameter gradient calculation, optimizer state processing, and model parameter update of the current active parameter sub-block are performed. In the pipelined parallel phase that does not include the currently active parameter sub-block, the model parameters of the corresponding parameter sub-block remain unchanged, and the parameter gradient calculation, optimizer state processing and model parameter update for the frozen parameter sub-block are skipped. Only the complete forward propagation is performed, and the backpropagation calculation required to pass the loss function gradient to the currently active parameter sub-block is performed.

9. A low-cost model training method based on block coordinate descent according to claim 1, characterized in that, The neural network model is a large language model, which is pre-trained in a graphics processing unit cluster that stores the complete training state of the large language model.

Images