A large model parallel training communication optimization method and device

Abstract

The application discloses a large model parallel training communication optimization method, which is applied to a system composed of multiple computing devices, each computing device is configured with a GPU, the GPU contains at least one process, the process is used for a task partition in large model parallel training, and the method is used for optimization of a ZeRO-1 stage of ZeRO data, the method comprises the following steps: performing forward calculation and reverse calculation of a model to determine gradient information; based on the gradient information, updating momentum in a local optimizer state; processing the updated momentum through 1-bit Allreduce communication, and updating errors generated in the compression process; and updating the model parameter partition stored on each computing device by using the aggregated momentum and variance. The method can reduce communication overhead and improve the efficiency of parallel training.

Description

Technical Field

[0001] This invention relates to the field of large model training technology, and in particular to a communication optimization method for parallel training of large models. Background Technology

[0002] Traditional data parallelism statically maintains all model states required throughout the training process across all data parallelization steps, leading to memory redundancy. ZeRO data parallelism eliminates memory redundancy in data parallelization by partitioning the model states instead of copying them. This results in memory consumption on each computing device being inversely proportional to the data parallelism dimension, making ZeRO data parallelism popular in both academia and industry. ZeRO has three main optimization phases, storing the optimizer state, gradients, and model parameters separately on multiple computing devices, reducing the memory required to store these states on each computing device in traditional data parallelism. However, this introduces more communication overhead during training. Summary of the Invention

[0003] To address the problems existing in the prior art, embodiments of this application provide a method, apparatus, computing device, computer storage medium, and product containing a computer program for optimizing communication during parallel training of large models, which can reduce communication overhead and improve the efficiency of parallel training.

[0004] In a first aspect, embodiments of this application provide a communication optimization method for large-scale parallel training of a model, applied to a system composed of multiple computing devices, each computing device being configured with a GPU, the GPU containing at least one process, the process being used for a task partition in large-scale parallel training of a model, the method being used for optimization of the ZeRO-1 stage of ZeRO data, the method comprising: performing forward and backward computation of the model to determine gradient information; updating the momentum in the local optimizer state based on the gradient information; processing the updated momentum through 1-bit Allreduce communication and updating the error generated during compression; and updating the model parameter partition stored on each computing device using the aggregated momentum and variance.

[0005] In some possible implementations, 1-bit Allreduce communication includes: sending compressed momentum data blocks to other computing devices using the All-to-All communication primitive and receiving compressed momentum data blocks from other computing devices; performing reduction operations on all received momentum data blocks; and collecting momentum data blocks from all computing devices using the Allgather communication primitive to complete global momentum aggregation.

[0006] In some possible implementations, the error generated during compression is updated, including recording the error generated in each compression and using it to compensate for the momentum after compression in the next iteration.

[0007] In some possible implementations, the method is applied to the optimization of the ZeRO-2 stage of the ZeRO data. After determining the gradient information, the method also includes: using a 1-bit Reduce-Scatter for communication to determine the aggregated momentum; and updating the model parameters based on the aggregated momentum and a fixed variance.

[0008] Among some possible implementations, the method also includes communicating model parameters via Allgather.

[0009] In some possible implementations, the method is applied to the optimization of the ZeRO-3 stage of ZeRO data. The method includes: each process performing forward computation of the model, using Allgather communication to collect complete model information on all computing devices; after completing the forward computation, each computing device discards the parts that do not belong to the current parameter partition; performing backward computation of the model, determining gradient information, and using Allgather communication to collect gradient information on all computing devices; performing 1-bit Reduce-Scatter communication and updating the compression error; the computing devices again discard parameter information that is not in the current partition, retaining only the necessary data for the current partition; and using the aggregated gradient information and variance, calculating the update of each process's assigned model parameter partition independently.

[0010] Secondly, embodiments of this application provide a communication optimization device for large-scale parallel training of a model, deployed in a system composed of multiple computing devices. Each computing device is equipped with a GPU, and the GPU contains at least one process. The process is used for a task partition in the large-scale parallel training of the model. The device is used for optimizing the ZeRO-1 stage of ZeRO data. The device includes: a gradient information acquisition module, used to perform forward and backward computation of the model to determine gradient information; a momentum update module, used to update the momentum in the local optimizer state based on the gradient information; a processing module, used to process the updated momentum through 1-bit Allreduce communication and update the error generated during the compression process; the processing module is also used to update the model parameter partition stored on each computing device using the aggregated momentum and variance.

[0011] Thirdly, embodiments of this application provide a computer-readable storage medium including computer-readable instructions that, when read and executed by a computer, cause the computer to perform the method as described in any of the first aspects.

[0012] Fourthly, embodiments of this application provide a computing device, including a processor and a memory, wherein the memory stores computer program instructions, which, when executed by the processor, perform the method as described in any of the first aspects.

[0013] Fifthly, embodiments of this application provide a product comprising a computer program that, when the computer program product is run on a processor, causes the processor to perform the method as described in any of the first aspects. Attached Figure Description

[0014] To more clearly illustrate the technical solutions of the embodiments of the present invention, the drawings used in the following description of the embodiments will be briefly introduced. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0015] Figure 1 This is a schematic diagram of a 1-bit Allreduce communication method provided in an embodiment of this application;

[0016] Figure 2 This is a flowchart illustrating a method for optimizing communication during parallel training of a large model based on 1-bit compression, as provided in an embodiment of this application.

[0017] Figure 3 This is a schematic diagram of an optimized process for ZeRO-1 provided in an embodiment of this application;

[0018] Figure 4 This is a schematic diagram of an optimized process for ZeRO-2 provided in an embodiment of this application;

[0019] Figure 5 This is a schematic diagram of an optimized process for ZeRO-3 provided in an embodiment of this application;

[0020] Figure 6 This is a schematic diagram of a large model parallel training communication optimization device based on 1-bit compression provided in an embodiment of this application. Detailed Implementation

[0021] To make the objectives, technical solutions, and advantages of the embodiments of the present invention clearer, the technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0022] In this article, the term "and / or" describes the relationship between related objects, indicating that three relationships can exist. For example, A and / or B can represent: A existing alone, A and B existing simultaneously, or B existing alone. The symbol " / " in this article indicates that the related objects are in an "or" relationship; for example, A / B means A or B.

[0023] The terms "first" and "second," etc., used in the specification and claims herein are used to distinguish different objects, not to describe a specific order of objects. For example, "first response message" and "second response message," etc., are used to distinguish different response messages, not to describe a specific order of response messages.

[0024] In the embodiments of this application, the terms "exemplary" or "for example" are used to indicate that something is an example, illustration, or description. Any embodiment or design that is described as "exemplary" or "for example" in the embodiments of this application should not be construed as being more preferred or advantageous than other embodiments or design. Specifically, the use of the terms "exemplary" or "for example" is intended to present the relevant concepts in a specific manner.

[0025] In the description of the embodiments of this application, unless otherwise stated, "multiple" means two or more, for example, multiple processing units means two or more processing units, multiple elements means two or more elements, etc.

[0026] To facilitate understanding of the embodiments of this application, the following will provide further explanation and description with reference to the accompanying drawings and specific embodiments. These embodiments do not constitute a limitation on the embodiments of the present invention.

[0027] First, the technical terms involved in the embodiments of this application will be introduced:

[0028] 1. Communication primitives are fundamental operations used in distributed computing to describe communication within a group of nodes. These operations must not be interrupted during execution to ensure correctness and system stability. Communication primitives are typically implemented as program segments consisting of several instructions, used to achieve a specific function such as broadcasting, collecting, or reducing. These primitives form the basis of aggregated communication, and their combinations can enable more complex communication patterns.

[0029] 2. All-to-All: This is a many-to-many communication primitive in which each node in the cluster sends data to all other nodes.

[0030] 3. All-Gather collects data from each GPU onto all GPUs, so that all GPUs have a complete dataset.

[0031] 4. Reduce communication is a widely used communication primitive in distributed computing. It involves reducing multiple data sources, such as summation and finding the maximum value, to reduce the amount of data or extract useful information.

[0032] 5. Scatter communication is used to send different data from one GPU to different GPUs.

[0033] Next, the technical solutions provided in the embodiments of this application will be introduced.

[0034] ZeRO, short for Zero Redundancy Optimizer, is an optimization technique for large-scale distributed model training. This technique primarily addresses the memory consumption issue in distributed training by reducing redundancy through slicing, achieving a time-for-space tradeoff and allowing the training of larger models on limited hardware resources. ZeRO has three levels: ZeRO-1, ZeRO-2, and ZeRO-3. ZeRO-1 splits the optimizer state, reducing memory consumption by 4 times, but with the same communication overhead as data parallelism. ZeRO-2 further splits the gradients, reducing memory consumption by 8 times, while maintaining the same communication overhead. ZeRO-3, building on ZeRO-2, also splits the model parameters, making the model's memory consumption inversely proportional to the parallelism of data parallelism, but increasing communication overhead. By reducing the number of complete model states required to store per GPU, ZeRO technology reduces memory consumption, making it possible to train large models. For example, for a model with 7.5B parameters, using ZeRO-1 can reduce memory usage from 120GB to 31.4GB, while ZeRO-2 can further reduce it to 16.6GB. ZeRO-3 is further optimized, allowing larger models to be trained on more GPUs, for example, training a model with over 1 trillion parameters on 1024 GPUs. However, this significantly increases communication overhead.

[0035] In view of this, this application provides a communication optimization method for large-scale parallel training of models based on 1-bit compression. 1-bit compression is a data compression technique that significantly reduces the required storage space or communication bandwidth by representing data in 1-bit form. In the context of deep learning model training, especially in distributed training, 1-bit compression can be used to optimize communication efficiency. Error-compensated 1-bit compression involves first compressing the data into a 1-bit representation, then calculating and storing the compression error, and finally summing it with the error when using the data for calculation in the next iteration. In the Adam optimizer, compressing the optimizer state (momentum) into 1 bit improves communication efficiency, and this algorithm has achieved good results in deep learning model training using traditional data parallel methods with parameter server architectures. However, in current large-scale language model parallel training, parameter server architectures can no longer meet the computational efficiency requirements of increasingly large-scale model training, and the Allreduce architecture is gradually becoming the mainstream method in large-scale language model training. The communication optimization method for large model parallel training based on 1-bit compression provided in this application aims to improve parallel training efficiency by reducing the amount and number of communications in ZeRO data parallel training. Specifically, this invention develops a 1-bit Allreduce communication method suitable for the Allreduce data parallel architecture based on a 1-bit compression algorithm, which significantly improves the communication efficiency in large model data parallel training. The 1-bit Allreduce method improves the workflow of ZeRO data parallelism, reduces the number and amount of communications in parallel training, and greatly improves the computational efficiency and training throughput of ZeRO data parallel training. At the same time, combined with the characteristics of the Adam optimizer, it ensures the convergence and accuracy of model training.

[0036] One-bit compression is a data compression technique that significantly reduces required storage space or communication bandwidth by representing data in one-bit form. In the context of deep learning model training, especially in distributed training, one-bit compression can be used to optimize communication efficiency. The implementation of the one-bit Allreduce communication method is based on NCCL (NVIDIA Collective Communications Library) and uses the All-to-All and Allgather communication primitives. NCCL is a communication library for high-performance parallel computing that provides various low-level communication primitives to support efficient data parallel operations on GPUs. Multiple processes can run on the GPU, participating in model optimization. The method execution steps include the following three steps:

[0037] S11: For the momentum to be communicated, after error compensation using the error obtained from the previous iteration, compress it using the compression operator. Then, update the error of the first communication in the second line, and divide the momentum into N blocks in the third line, and use All-to-All for communication.

[0038] In this embodiment, the momentum in the optimizer state is compressed into a 1-bit representation. Since 1-bit compression loses numerical information, the error introduced during compression needs to be calculated and stored. In subsequent iterations, this error value is used to compensate for the compressed data, reducing the impact of compression. In each iteration, the compressed data and the stored error are used together for calculation. Before communication or update operations, the error is added to the 1-bit representation to recover some information from the original data, completing error compensation. In the current step, the compression error calculated in the previous iteration is first used to compensate for the current momentum. After error compensation, the momentum is compressed into a 1-bit representation using a compression operator. After 1-bit compression, the new error introduced by this compression needs to be calculated and stored for use in the next iteration. The compressed momentum data is divided into N blocks, where N is the number of computing devices participating in the communication. An All-to-All communication mode is used, where each computing device sends its own data block to all other devices. In All-to-All communication, each i-th worker process delivers blocks other than the i-th block to its corresponding process and receives blocks delivered by processes other than the i-th worker process. For example, Figure 1 A schematic diagram of a 1-bit Allreduce communication method provided in an embodiment of this application is shown. Figure 1 As shown, in All-to-All communication, taking three processes as an example, each process contains three blocks. The first process passes the second block to the second process and the third block to the third process. Simultaneously, it receives the first block from the second process and the first block from the third process. Similarly, the second process passes the first block to the first process and the third block to the third process, while simultaneously receiving the second block from both the first and third processes. The third process passes the first block to the first process and the second block to the second process, while simultaneously receiving the third block from both the first and second processes.

[0039] S12: After communication, each i-th working process obtains the i-th block of all processes and performs a reduction operation on it. In this embodiment, the reduction operation is performed by taking the average value.

[0040] In this embodiment, after completing All-to-All communication, each process contains the momentum of other processes; that is, each i-th worker process obtains the i-th block from all processes. At this point, a reduction operation is performed on the i-th process. A reduction operation is an operation that combines multiple data values ​​into a single value. Common reduction operations include summation, averaging, finding the maximum or minimum value, etc. In this embodiment, the reduction operation chosen is averaging. Please continue to refer to... Figure 1 By averaging the results for each process, we can obtain the result after the reduction operation. At this point, the processes involved in the communication have completed 1-bit Reduce-Scatter communication, which is called 1-bit Reduce-Scatter communication.

[0041] S13: After obtaining the average value, compress and update the second compression communication error. Use Allgather to collect different blocks from each process to obtain the final aggregated momentum.

[0042] In this embodiment, after the reduction operation, the integrated data is obtained, which is a local average. After obtaining the local average, it can be further compressed. Because 1-bit compression loses all information except the sign bit, this step requires calculating and storing the error caused by the compression for use in subsequent iterations to compensate for the impact of compression. Next, each process uses Allgather to obtain the local average of other processes. Through Allgather, each process can obtain the global momentum. At this point, the processes participating in the communication complete 1-bit Allreduce communication.

[0043] Based on the implementation of the above-described 1-bit Allreduce communication method, this application provides a method for optimizing communication during large-scale parallel training of models based on 1-bit compression. For example, Figure 2 This diagram illustrates a flowchart of a large-model parallel training communication optimization method based on 1-bit compression, provided in an embodiment of this application. This method is applied to the ZeRO-1 stage (level) of ZeRO data. Figure 2 As shown, the communication optimization method for parallel training of large models can include the following steps:

[0044] S21: Perform forward and backward computation of the model to determine gradient information.

[0045] In this embodiment, the model is trained jointly using multiple GPUs, with at least one process on each GPU dedicated to model training. This embodiment uses an example where each computing device has one GPU and each GPU uses one process. The model and training data are divided into n parts, the same number as the total number of processes. Each process is used to execute a portion of the overall model training. During training, forward computation of the model is performed first to generate the predicted output. Then, backward computation of the model is performed to generate gradient information.

[0046] S22: Update the momentum in the local optimizer state based on gradient information.

[0047] In the ZeRO-1 stage of the original ZeRO data, after determining the gradient information, Allreduce communication is performed immediately. In this embodiment, gradient communication is not performed immediately after the forward and backward computations are completed. To save memory and reduce data transmission, the gradient storage space is reused to calculate and update the momentum partition of the local computing device with the corresponding partition in the gradient, thereby achieving local momentum update. At this time, the content of the gradient storage space contains the updated momentum of the local computing device and the gradients required by other computing devices to update their momentum.

[0048] S23: Process the updated momentum via 1-bit Allreduce communication.

[0049] In this embodiment, by means of Figure 1 The 1-bit Allreduce shown communicates with this new momentum tensor to aggregate global momentum, and each computing device can obtain the complete aggregated momentum.

[0050] S24: Errors generated during the update compression process.

[0051] In this embodiment, during the 1-bit compression process in 1-bit Allreduce, the original multi-precision floating-point momentum values ​​are compressed into a 1-bit representation. This compression method significantly reduces data precision, thus introducing errors. When compressing momentum values, the system records the error generated during this compression. In subsequent iterations, the recorded error is used to compensate for the compressed momentum values. In each iteration, the momentum values ​​are compressed again, and each compression may introduce new errors. Therefore, the recorded error values ​​need to be continuously updated for more accurate compensation in the next iteration. In this way, although the data is compressed and decompressed in each iteration, by continuously compensating for compression errors, the impact of compression on the training process can be reduced, ensuring the accuracy and convergence of the training results.

[0052] S25: Use the aggregated momentum and the updated compression error to update the partitions of model parameters stored on each computing device.

[0053] In this embodiment, the complete aggregated momentum is obtained, and this momentum is used in conjunction with the variance to complete the model parameter update. At this point, it is no longer necessary to use Allgather to communicate the model parameters; after the parameters are updated, only the momentum of the corresponding partition of this process is retained.

[0054] This completes one round of model training. Repeat steps S21-S25 in a loop until the model training is complete.

[0055] Figure 3 This is a schematic diagram of the overall optimization process of ZeRO-1 provided in the embodiments of this application, as shown below. Figure 3 As shown in the diagram, the example described uses the division of the model and data into two parts. For each process, the assigned model and data are processed. After forward and backward computation, gradient information is determined. Local momentum updates are performed, followed by 1-bit Allreduce communication. After completion, compression error updates are performed. Based on the momentum and variance information obtained from the 1-bit Allreduce communication, the model parameters are updated, thus completing one iterative cycle.

[0056] Based on the optimization of ZeRO-1 data, this application also provides a method for optimizing large-scale parallel training communication using 1-bit compression, specifically for optimizing ZeRO-2 data. Since ZeRO-2 further splits the gradients based on ZeRO-1, it is a method built upon ZeRO-1. Its steps are similar to... Figure 2 The optimized portion of the ZeRO-1 shown is the same. For an example, please refer to... Figure 4 After the forward and backward computations are completed, the momentum is updated and stored in the corresponding partition of the gradient storage space. Next, communication is performed using 1-bit Reduce-Scatter (steps S11 and S12) to obtain the aggregated momentum. The aggregated momentum, along with a fixed variance, is used to update the model parameters. At this point, only a portion of the model parameters are updated on each computing device. Finally, Allgather is used to communicate the model parameters to obtain the complete model parameters.

[0057] Based on the optimization of ZeRO-2, this application also provides a method for optimizing parallel training communication of large models using 1-bit compression, specifically for optimizing ZeRO-3 data. Since ZeRO-3 is based on ZeRO-2 and also splits the model parameters, a 1-bit Reduce-Scatter can be used to replace the original Reduce-Scatter in the ZeRO-3 optimization for compressed gradient communication. For an example, please refer to... Figure 5 , Figure 5 This is a schematic diagram illustrating the optimization process of ZeRO-3 provided in an embodiment of this application. For example... Figure 5 As shown, the optimized communication method for large model parallel training based on 1-bit compression applied to ZeRO-3 specifically includes: allocating the model and data to each process. Each process performs the forward computation of the model. During this process, the Allgather communication operation is used to collect complete model information on all computing devices to ensure that each device can perform complete forward computation. After the forward computation is completed, each computing device no longer needs to store complete model parameters, so each process only retains the parameter partition it is responsible for, discarding parts that do not belong to that parameter partition to reduce storage space. Next, the process performs the backward computation of the model to determine gradient information. During this process, the Allgather communication operation is used to collect gradient information on all computing devices to prepare for subsequent gradient compression and updates. 1-bit Reduce-Scatter communication is performed, and the compression error is updated. Then, the computing device again discards parameter information that is not in its own partition, retaining only the necessary data. Using the aggregated gradient information and variance, each process independently updates its assigned model parameter partition, thus completing one iteration. After multiple iterations, the model is trained.

[0058] The above describes the communication optimization for large-scale parallel training based on 1-bit compression provided in the embodiments of this application, including communication optimization for the ZeRO-1, ZeRO-2, and ZeRO-3 stages of ZeRO data. By adopting this method, the communication overhead during parallel training of ZeRO data can be reduced, thereby improving the efficiency of parallel training.

[0059] It is understood that the sequence number of each step in the above embodiments does not imply the order of execution. The execution order of each process should be determined by its function and internal logic, and should not constitute any limitation on the implementation process of the embodiments of this application. Furthermore, in some possible implementations, each step in the above embodiments may be selectively executed according to actual circumstances; it may be partially or fully executed, without limitation here. All or part of any feature of any embodiment of this application can be freely and arbitrarily combined without contradiction. The combined technical solutions are also within the scope of this application.

[0060] Based on the methods in the above embodiments, this application also provides a large model parallel training communication optimization device based on 1-bit compression.

[0061] For example, Figure 6 An optimized communication device for large-scale parallel training of models based on 1-bit compression is shown. For example... Figure 6 As shown, the large model parallel training communication optimization device 600 based on 1-bit compression may include: a gradient information acquisition module 601, a momentum update module 602, and a processing module 603.

[0062] The gradient information acquisition module 601 is used to perform forward and backward calculations of the model to determine gradient information.

[0063] The momentum update module 602 is used to update the momentum in the local optimizer state based on the gradient information.

[0064] The processing module 603 is used to process the updated momentum through 1-bit Allreduce communication and update the errors generated during the compression process.

[0065] The processing module 603 is also used to update the model parameter partitions stored on each computing device using the aggregated momentum and variance.

[0066] It should be understood that the above-described device is used to execute the methods in the above embodiments. The implementation principle and technical effect of the corresponding program modules in the device are similar to those described in the above methods. The working process of the device can be referred to the corresponding process in the above methods, and will not be repeated here.

[0067] Based on the methods in the above embodiments, this application provides a computer-readable storage medium storing a computer program that, when run on a processor, causes the processor to execute the methods in the above embodiments.

[0068] Based on the methods in the above embodiments, this application provides a computer program product that, when run on a processor, causes the processor to execute the methods in the above embodiments.

[0069] It is understood that the processor in the embodiments of this application may be a central processing unit (CPU), or other general-purpose processors, digital signal processors (DSPs), application-specific integrated circuits (ASICs), field-programmable gate arrays (FPGAs), or other programmable logic devices, transistor logic devices, hardware components, or any combination thereof. A general-purpose processor may be a microprocessor or any conventional processor.

[0070] The method steps in the embodiments of this application can be implemented in hardware or by a processor executing software instructions. The software instructions can consist of corresponding software modules, which can be stored in random access memory (RAM), flash memory, read-only memory (ROM), programmable read-only memory (PROM), erasable programmable read-only memory (EPROM), electrically erasable programmable read-only memory (EEPROM), registers, hard disks, portable hard disks, CD-ROMs, or any other form of storage medium known in the art. An exemplary storage medium is coupled to the processor, enabling the processor to read information from and write information to the storage medium. Of course, the storage medium can also be a component of the processor. The processor and the storage medium can reside in an ASIC.

[0071] In the above embodiments, implementation can be achieved entirely or partially through software, hardware, firmware, or any combination thereof. When implemented using software, it can be implemented entirely or partially in the form of a computer program product. The computer program product includes one or more computer instructions. When the computer program instructions are loaded and executed on a computer, all or part of the processes or functions described in the embodiments of this application are generated. The computer can be a general-purpose computer, a special-purpose computer, a computer network, or other programmable device. The computer instructions can be stored in a computer-readable storage medium or transmitted through the computer-readable storage medium. The computer instructions can be transmitted from one website, computer, server, or data center to another website, computer, server, or data center via wired (e.g., coaxial cable, fiber optic, digital subscriber line (DSL)) or wireless (e.g., infrared, wireless, microwave, etc.) means. The computer-readable storage medium can be any available medium that a computer can access or a data storage device such as a server or data center that integrates one or more available media. The available medium can be a magnetic medium (e.g., floppy disk, hard disk, magnetic tape), an optical medium (e.g., DVD), or a semiconductor medium (e.g., solid-state disk (SSD)).

[0072] It is understood that the various numerical designations used in the embodiments of this application are merely for descriptive convenience and are not intended to limit the scope of the embodiments of this application.

Claims

1. A communication optimization method for parallel training of large models, characterized in that, A method for optimizing the ZeRO-1 stage of ZeRO data is applied to a system consisting of multiple computing devices, each of which is configured with a GPU, the GPU containing at least one process for a task partition in parallel training of a large model. The method includes: Perform forward and backward computations of the model to determine gradient information; Based on the gradient information, the momentum in the local optimizer state is updated; wherein, the momentum partition of the computing device and the corresponding partition in the gradient are calculated and updated to generate new momentum by reusing the gradient storage space to realize local momentum update. The momentum value is compressed into a 1-bit representation, and the error generated during the compression process is recorded through an error compensation mechanism. The error compensation mechanism includes: before each compression, the current momentum is compensated using the error of the previous iteration, and the error value of the current compression is updated after compression. The compensated and updated momentum is processed through 1-bit Allreduce communication, and the errors generated during compression are updated. Specifically, the All-to-All communication primitive is used to send compressed momentum data blocks to other computing devices and receive compressed momentum data blocks from other computing devices. A reduction operation is performed on all received momentum data blocks. The Allgather communication primitive is used to collect momentum data blocks from all computing devices to complete the global aggregation of momentum. The momentum and variance after global aggregation are used to update the partition of model parameters stored on each computing device.

2. The method according to claim 1, characterized in that, The error generated during the update compression process includes: recording the error generated in each compression and using it to compensate for the momentum after compression in the next iteration.

3. The method according to claim 1, characterized in that, The method is applied to the optimization of the ZeRO-2 stage of ZeRO data. After determining the gradient information, the method further includes: Use 1-bit Reduce-Scatter for communication to determine the aggregated momentum; Based on the aggregated momentum and fixed variance, the model parameters are updated.

4. The method according to claim 3, characterized in that, The method also includes communicating model parameters via Allgather.

5. The method according to claim 1, characterized in that, The method is applied to the optimization of the ZeRO-3 stage of ZeRO data, and the method includes: Each process performs forward computation of the model and uses the Allgather communication operation to collect complete model information from all computing devices. After completing the forward computation, each computing device will discard any parts that do not belong to the current parameter partition. Perform reverse computation of the model to determine gradient information, and use the Allgather communication operation to collect gradient information from all computing devices; Perform 1-bit Reduce-Scatter communication and update compression errors; The computing device discards parameter information from non-current partitions again, retaining only the necessary data from the current partition; Using the aggregated gradient information and variance, each process independently updates its assigned model parameter partition.

Images