Method and device for accelerating distributed training of ai large model based on adaptive mixed gradient compression

Abstract

This invention discloses a method and apparatus for accelerating distributed training of large AI models based on adaptive hybrid gradient compression. The method includes: in a master-slave node structure, a PS parameter server and computing nodes each generate gradients through local training iterations; compressing the gradients using an adaptive hybrid gradient compression method; sending the compressed gradients to the PS parameter server, followed by decompression, aggregation, and updating; sending the updated model parameters to each computing node; and obtaining a distributed machine learning model through multiple iterations of training. The adaptive hybrid gradient compression method includes: for a certain type of gradient, if it is a one-dimensional vector, no gradient compression is performed; if its gradient count is greater than or equal to a set threshold, a class-by-class gradient clustering mechanism is used for compression; otherwise, a Tensor-Train decomposition mechanism is used for compression. This invention reduces training time by reducing the communication required for distributed training, without sacrificing model accuracy or introducing large computational loads.

Description

Technical Field

[0001] This invention relates to the technical field of distributed training of large AI models, and in particular to a method and apparatus for accelerating distributed training of large AI models based on adaptive hybrid gradient compression. Background Technology

[0002] Distributed data parallel training is a popular and effective training method for deep neural networks. However, due to the frequent transmission of large amounts of gradient parameters in distributed training, communication bandwidth has become a bottleneck limiting training performance. This situation is exacerbated in unstable network environments such as edge computing. To address this issue, researchers have extensively explored gradient compression techniques, which can effectively reduce the amount of gradient data that needs to be transmitted, thereby reducing communication overhead. Gradient compression schemes compress gradients before sending them. However, existing methods are almost all relatively coarse, ignoring the differences between each class of gradients. This can increase training time and reduce model accuracy.

[0003] Different deep neural networks have different structures, and the number of parameters in each layer varies, resulting in different numbers of gradients across layers. This leads to variations in the number of gradients of each type within the neural network. Some types of gradients may have significantly more than others, potentially reaching nearly 600,000, while others may have only 10 values; these gradient types are typically one-dimensional vectors. This difference usually stems from the fact that some layers in the network structure are more sensitive to input data and therefore require more gradient information to adjust their parameters. In other layers, because they are less sensitive to input data and require fewer parameter updates, their gradient numbers are also smaller. This variation in the number of gradients of each type is a common phenomenon.

[0004] If the differences in the number of gradients of each class are not fully considered when compressing gradients, it will negatively impact the training performance of the model. For example, traditional gradient compression methods treat all gradients as the same and process them using the same compression method, such as Top-k gradient dilution. However, this coarse approach fails to account for the order-of-magnitude differences between the gradients of each class, leading to low training efficiency and decreased model accuracy. Therefore, proposing a more refined and efficient gradient compression method is an urgent problem to be solved. Summary of the Invention

[0005] The main objective of this invention is to overcome the shortcomings and deficiencies of the prior art and provide a method and apparatus for accelerating distributed training of large AI models based on adaptive hybrid gradient compression. This invention reduces the amount of communication required for training in distributed machine learning, thereby reducing training time without sacrificing model accuracy.

[0006] To achieve the above objectives, the present invention adopts the following technical solution:

[0007] In a first aspect, the present invention provides a method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression, comprising the following steps:

[0008] A distributed machine learning model is constructed, a training dataset is initialized, the training dataset is divided into n regions and sent to n computing nodes (workers); the distributed training of the distributed machine learning model is based on "PS (Parameter server) architecture + data parallelism", and has a master-slave node structure, including a PS parameter server, n computing nodes and connected communication lines.

[0009] In a master-slave node architecture, the PS parameter server initializes the model parameters w0 and pushes these parameters to all compute nodes. The PS parameter server and compute nodes simultaneously perform local training using the dataset, generating gradients after t iterations of training. and It is the gradient generated by the PS parameter server in the t-th iteration. It is the gradient generated by the nth computation node in the tth iteration;

[0010] Each computing node generates its own gradient Gradient compression is performed using an adaptive hybrid gradient compression method to generate compressed gradients. The adaptive hybrid gradient compression method specifically includes: gradient There are multiple classes. For a given class of gradients, if it is a one-dimensional vector, no compression is performed; if its number is greater than or equal to a set threshold, a class-by-class gradient clustering mechanism is used to cluster the gradients. Perform compression and set the compression mode identifier M=0; when its quantity is less than a set threshold, use the Tensor Train decomposition mechanism to process the gradient. Compression is performed, and the compression mode identifier M = 1 is set; the gradient after compression via the class-by-class gradient clustering mechanism This includes the centroid value and cluster identifier; the gradient compressed using the Tensor-Train decomposition mechanism. It includes a set of decomposed tensors and their corresponding singular values;

[0011] The compressed gradient Send to the PS parameter server; the PS parameter server processes the compressed gradient according to the compression mode identifier. Decompress the data to recover the original gradient, and then compare it with the original gradient. Aggregate the gradients to obtain the global gradient g. y ; Utilizing the global gradient gt Update the parameters of the distributed machine learning model to obtain the new model parameters.

[0012] New model parameters The new model parameters are sent to the computing nodes. After each computing node receives the new model parameters, it replaces its own old model parameters, and one iteration ends. After multiple iterations, the distributed training of the distributed machine learning model is completed.

[0013] As a preferred technical solution, the class-by-class gradient clustering mechanism specifically includes:

[0014] gradient There are multiple classes, and each class has multiple gradient values. A clustering algorithm is used to cluster each class of gradients into multiple clusters. Each cluster has a centroid value and a corresponding cluster identifier, resulting in multiple centroid values ​​and cluster identifiers. These multiple centroid values ​​and cluster identifiers represent the gradients compressed through the class-wise gradient clustering mechanism.

[0015] As a preferred technical solution, the Tensor Train decomposition mechanism is specifically as follows:

[0016] gradient There are multiple classes. One class of gradients is an N-order tensor X. Each class of gradients is decomposed into Tensor Train based on singular value decomposition. After N-1 singular value decompositions, N small tensors and corresponding singular spectra are obtained. Each singular spectrum contains some singular values. The number of non-zero singular values ​​in the singular spectrum is extracted to obtain the TT rank.

[0017] Taking a low-rank approximation of the TT rank, for a given tensor X, the Tensor-Train decomposition makes:

[0018]

[0019] At this point, the optimal TT low-rank approximation of tensor X is:

[0020]

[0021] Among them, a n s represents the nth geometric or auxiliary index. n This represents the nth open or physical index. This represents the Nth tensor, where the N tensors and their corresponding singular values ​​are the gradients compressed using the TT decomposition mechanism.

[0022] As a preferred technical solution, the PS parameter server processes the compressed gradient according to the compression mode identifier. Decompress, recover the original gradient and match it with the gradient. The aggregation process includes the following steps:

[0023] When the compression mode identifier M=1, a splicing-based tensor recovery mechanism is used to process the compressed gradient. Decompress, recover the original gradient and match it with the gradient. Perform aggregation;

[0024] When the compression mode identifier M = 0, a gradient recovery mechanism based on centroid replacement is used to recover the compressed gradient. Recover the original gradient and compare it with the gradient Perform aggregation.

[0025] As a preferred technical solution, the tensor recovery mechanism based on splicing specifically includes:

[0026] By reversing the order of the smaller tensors according to the Tensor-Train decomposition method, we obtain an N-order tensor, which is an approximate reconstruction of the original tensor.

[0027]

[0028] in, This represents the i-th element in the k-th small tensor obtained from Tensor-Train decomposition. k Dimension, and × k Indicates will Perform a Kronecker product with the previous result, where 1 ≤ k ≤ N.

[0029] As a preferred technical solution, the gradient recovery mechanism based on centroid replacement specifically includes:

[0030] The PS parameter server replaces all gradient values ​​within a cluster with the centroid value of that cluster, thereby approximately restoring the original gradient.

[0031] As a preferred technical solution, the method of utilizing the global gradient g t Update the parameters of the distributed machine learning model to obtain the new model parameters. Specifically:

[0032] Using the global gradient g t The parameters of the distributed machine learning model are updated according to the parameter update formula to obtain the new model parameters. The parameter update formula is as follows:

[0033]

[0034] in, η represents the model parameters of the PS parameter server at the t-th iteration, and η represents the learning rate at the time of update, 0≤η≤1.

[0035] Secondly, the present invention provides a system for accelerating distributed training of large AI models based on adaptive hybrid gradient compression, which is applied to the method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression, including a model building module, a partitioning module, a first gradient processing module, a second gradient processing module, and a model training module.

[0036] The model building module is used to build a distributed machine learning model, initialize the training dataset, partition the dataset and send it to n computing nodes (workers); the distributed training of the distributed machine learning model is based on "PS architecture + data parallelism", has a master-slave node structure, including a PS parameter server, n computing nodes and connected communication lines.

[0037] The partitioning module is used in a master-slave node structure to initialize model parameters w0 on the PS parameter server and push these parameters to all computing nodes. The PS parameter server and computing nodes simultaneously perform local training using the dataset, and generate gradients after completing t iterations of training. and It is the gradient generated by the PS parameter server in the t-th iteration. It is the gradient generated by the nth computation node in the tth iteration;

[0038] The first gradient processing module is used by each computing node to process the gradient it generates. Gradient compression is performed using an adaptive hybrid gradient compression method to generate compressed gradients. The adaptive hybrid gradient compression method specifically includes: gradient There are multiple classes. For a given class of gradients, if it is a one-dimensional vector, no compression is performed; if its number is greater than or equal to a set threshold, a class-by-class gradient clustering mechanism is used to cluster the gradients. Perform compression and set the compression mode identifier M=0; when its quantity is less than a set threshold, use the Tensor Train decomposition mechanism to process the gradient. Compression is performed, and the compression mode identifier M = 1 is set; the gradient after compression via the class-by-class gradient clustering mechanism This includes the centroid value and cluster identifier; the gradient compressed using the Tensor-Train decomposition mechanism. It includes a set of decomposed tensors and their corresponding singular values;

[0039] The first gradient processing module is used to process the compressed gradient. Send to the PS parameter server; the PS parameter server processes the compressed gradient according to the compression mode identifier. Decompress the data to recover the original gradient, and then compare it with the original gradient. Aggregate the gradients to obtain the global gradient g. t ; Utilizing the global gradient g t Update the parameters of the distributed machine learning model to obtain the new model parameters.

[0040] The model training module is used to train new model parameters. The new model parameters are sent to the computing nodes. After each computing node receives the new model parameters, it replaces its own old model parameters, and one iteration ends. After multiple iterations, the distributed training of the distributed machine learning model is completed.

[0041] Thirdly, the present invention provides an electronic device, the electronic device comprising:

[0042] At least one processor; and,

[0043] A memory communicatively connected to the at least one processor; wherein,

[0044] The memory stores computer program instructions that can be executed by the at least one processor, which enables the at least one processor to execute the method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression.

[0045] Fourthly, the present invention provides a computer-readable storage medium storing a program, which, when executed by a processor, implements the method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression.

[0046] Compared with the prior art, the present invention has the following advantages and beneficial effects:

[0047] (1) This invention proposes a class-by-class gradient clustering method, which uses the Minibatch-Kmeans algorithm to cluster each class of gradients respectively. Then, each computing node only sends all the centroids and cluster identifiers after clustering to the PS parameter server, which can reduce a lot of communication traffic.

[0048] (2) This invention proposes an adaptive hybrid gradient compression method. For gradients below a threshold, tensor decomposition is used, and for gradients above a threshold, clustering is used. The mechanism can be adaptively switched to compress the gradient, thereby maintaining the model accuracy without loss while ensuring training efficiency. Attached Figure Description

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

[0050] Figure 1 This is a flowchart of a method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression, according to an embodiment of the present invention.

[0051] Figure 2 This is a schematic diagram of the class-by-class gradient clustering mechanism in an embodiment of the present invention;

[0052] Figure 3 This is a schematic diagram of the Tensor Train decomposition mechanism in an embodiment of the present invention;

[0053] Figure 4 This is a schematic diagram of the system for accelerating distributed training of large AI models based on adaptive hybrid gradient compression, according to an embodiment of the present invention.

[0054] Figure 5 This is a structural diagram of the electronic device in an embodiment of the present invention. Detailed Implementation

[0055] To enable those skilled in the art to better understand the present application, the technical solutions in the embodiments of the present application will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are merely some embodiments of the present application, and not all embodiments. All other embodiments obtained by those skilled in the art based on the embodiments of the present application without creative effort are within the scope of protection of the present application.

[0056] In this application, the reference to "embodiment" means that a specific feature, structure, or characteristic described in connection with an embodiment may be included in at least one embodiment of this application. The appearance of this phrase in various places throughout the specification does not necessarily refer to the same embodiment, nor is it a mutually exclusive, independent, or alternative embodiment. It will be explicitly and implicitly understood by those skilled in the art that the embodiments described in this application can be combined with other embodiments.

[0057] Please see Figure 1 This embodiment describes a method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression, which includes the following steps:

[0058] S1. Construct a distributed machine learning model, initialize the dataset, partition the dataset, and send it to n computing nodes;

[0059] Optionally, the distributed machine learning model has a master-slave node structure, including a PS parameter server, n computing nodes, and connected communication lines.

[0060] S2. Using the dataset, perform local training on both the PS parameter server and the compute nodes simultaneously. After completing t rounds of iterative training, generate gradients respectively. and It is the gradient generated by the PS parameter server in the t-th iteration. It is the gradient generated by the nth computation node in the tth iteration;

[0061] Optionally, t rounds of iterative training, i.e., t%T=0;

[0062] S3. Combine the gradients from multiple computation nodes. Gradient compression is performed using an adaptive hybrid gradient compression method to generate compressed gradients.

[0063] Optionally, the adaptive hybrid gradient compression method specifically involves, when the gradient... If it is a one-dimensional vector, then no compression is performed; when the gradient... If the number is greater than or equal to a set threshold, gradient clustering is used to analyze the gradient. Perform compression and set the compression mode identifier; when the gradient... If the number is less than a set threshold, the gradient is decomposed using the Tensor Train mechanism. Perform compression and set the compression mode identifier;

[0064] Preferably, the gradient compressed by the class-by-class gradient clustering mechanism Including centroid value and cluster identifier;

[0065] Preferably, the class-by-class gradient clustering mechanism is as follows:

[0066] gradient There are multiple classes, and each class has multiple gradient values. A clustering algorithm is used to cluster each class of gradients into multiple clusters. Each cluster has a centroid value and a corresponding cluster identifier, resulting in multiple centroid values ​​and cluster identifiers. These multiple centroid values ​​and cluster identifiers represent the gradients compressed through the class-wise gradient clustering mechanism.

[0067] In a specific embodiment, the above-mentioned class-by-class gradient clustering mechanism is as follows:

[0068] gradient There are C classes, and each class has M gradient values. A clustering algorithm is used to cluster each class's gradient into K clusters. Each cluster has a centroid value and a corresponding cluster identifier, resulting in K centroid values ​​and cluster identifiers.

[0069] S31. Initialize the centroid values ​​of the K clusters. Then, use the Minibatch-Kmeans algorithm to calculate the total distance D between each gradient value and its nearest centroid. Iterate multiple times to minimize D until the membership of each cluster is stable, at which point the compression terminates. The K centroid values ​​and cluster identifiers are the gradients compressed by the class-wise gradient clustering mechanism.

[0070] S32. Specifically, for the v-th class of gradients, calculate the total distance D between each gradient value in that class and its nearest centroid. v for,

[0071]

[0072] in, It is the j-th value of the gradient of class v. Let represent the value of the centroid of the i-th cluster in class v, where 1 ≤ v ≤ C, 1 ≤ i ≤ K, and 1 ≤ j ≤ M.

[0073] Preferably, the Tensor Train decomposition mechanism, such as Figure 3 As shown, specifically,

[0074] S33. Treat the gradient as an Nth-order tensor X, and perform Tensor Train decomposition on the gradient through singular value decomposition. After N-1 Tensor Train decompositions, extract the number of non-zero singular values ​​in the singular spectrum to obtain the TT rank.

[0075] S34. Taking a low-rank approximation of the TT rank, for a given tensor X, the Tensor-Train decomposition makes:

[0076]

[0077] At this point, the optimal TT low-rank approximation of tensor X is:

[0078]

[0079] Where a n s represents the nth geometric or auxiliary index. n This represents the nth open or physical index. This represents the Nth tensor.

[0080] Compared with other tensor decomposition methods, Tensor-Train decomposition can handle sparse data better and can be computed efficiently, making it a decomposition method that performs well on high-dimensional data.

[0081] For example, such as Figure 3As shown, for a fourth-order tensor, it can first be regarded as the product of a second-order tensor and a third-order tensor. Perform singular value decomposition on this second-order tensor to obtain two small tensors and a singular spectrum (that is, some singular values). Then, regard these two small tensors and the original third-order tensor as the product of a new third-order tensor. Perform singular value decomposition on this third-order tensor to obtain three small tensors and some singular values. And so on, until four small tensors are obtained.

[0082] S4. Compress the gradient Send to the PS parameter server; the PS parameter server processes the compressed gradient according to the compression mode identifier. Recover the original gradient and compare it with the gradient Aggregate the gradients to obtain the global gradient g. t ; Utilizing the global gradient g t Update the parameters of the distributed machine learning model to obtain the new model parameters.

[0083] Optionally, the computing node only needs to send k centroid values ​​and cluster identifiers for each type of gradient to the PS parameter server. These centroid values ​​and cluster identifiers contain important information about the gradient.

[0084] Optionally, the PS parameter server adjusts the compressed gradient according to the compression mode identifier. Recover the original gradient and compare it with the gradient Aggregation involves the following steps:

[0085] When the compression mode identifier M=1, a splicing-based tensor recovery mechanism is used to process the compressed gradient. Recover the original gradient and compare it with the gradient Perform aggregation;

[0086] When the compression mode identifier M = 0, a gradient recovery mechanism based on centroid replacement is used to recover the compressed gradient. Recover the original gradient and compare it with the gradient Perform aggregation;

[0087] Preferably, the splicing-based tensor recovery mechanism specifically includes:

[0088] By concatenating the smaller tensors in reverse order according to their dimensions, we obtain an N-order tensor, which is an approximate reconstruction of the original tensor.

[0089]

[0090] in, This represents the i-th element in the k-th small tensor obtained from Tensor-Train decomposition. k Dimension, and × k Indicates will Perform a Kronecker product with the previous result, where 1 ≤ k ≤ N;

[0091] Preferably, the gradient recovery mechanism based on centroid replacement specifically includes:

[0092] The PS parameter server uses the centroid value of the cluster as the value of all gradients within that cluster to replace the original gradient, thereby approximately restoring the original gradient.

[0093] Preferably, the use of global gradient g t Update the parameters of the distributed machine learning model to obtain the new model parameters. Specifically,

[0094] Using the global gradient g t The parameters of the distributed machine learning model are updated according to the parameter update formula to obtain the new model parameters. The parameter update formula is as follows:

[0095]

[0096] in, The PS parameter server represents the model parameters at the t-th iteration, and η represents the learning rate during the update, which can be set to 0.9 for example.

[0097] For example, let's further illustrate the class-wise gradient clustering method with a specific embodiment. Suppose that the first class gradient in DownsizedVGGNet-11 has 16 values, and their distribution is as follows: Figure 2 As shown. After clustering the gradient into k=2 categories, the centroid of each cluster is used to approximate all gradient values ​​within the same cluster, where 0 (blue value) represents the category with a centroid of -0.19, and 1 (orange value) represents another category with a centroid of 1.48. The worker only needs to input the centroid and location information of the gradient (i.e., Figure 3 The cluster IDs of the workers are sent to the PS parameter server, thus reducing communication overhead. Therefore, for a Downsized VGGNet-11 with 66 gradients, only 66*2=132 centroids and their cluster IDs need to be sent to the PS parameter server. Once the PS receives these centroids and cluster IDs from the workers, it approximates the original gradients based on them. In other words, the PS parameter server uses -0.19 to replace 0 and 1.48 to replace 1, thereby approximating the original gradients.

[0098] Furthermore, after the PS parameter server receives the centroid and cluster identifier sent by the computing node, the PS parameter server approximately recovers the original gradient based on the centroid and cluster identifier. Specifically, the PS parameter server uses -0.19 of the centroid to replace 0 of the cluster identifier and uses 1.48 of the centroid to replace 1 of the cluster identifier, thereby achieving approximate recovery of the original gradient.

[0099] S5. Set the new model parameters Send the computing nodes and update their parameters. Through iterative training, a trained distributed machine learning model is obtained.

[0100] This invention reduces training time by minimizing the communication required for training in distributed machine learning, without sacrificing model accuracy. Using the CIFAR-10 and CIFAR-100 datasets for training, compared to the classic Bulk Synchronous Parallel (BSP) mechanism, this invention significantly reduces communication, thus avoiding communication bottlenecks and reducing training time. Furthermore, this invention achieves model accuracy close to that of BSP; that is, it suffers almost no loss in model accuracy, as shown in Tables 1 and 2.

[0101] Table 1 shows the experimental results of training Downsized VGGNet-11 and Downsized ResNet-18 on CIFAR-10.

[0102]

[0103] Table 2 shows the experimental results of training a downsized VGGNet-11 model on CIFAR-100.

[0104]

[0105] Based on the same idea as the method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression in the above embodiments, this invention also provides a system for accelerating distributed training of large AI models based on adaptive hybrid gradient compression. This system can be used to execute the above-described method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression. For ease of explanation, the structural diagram of the system embodiment for accelerating distributed training of large AI models based on adaptive hybrid gradient compression only shows the parts related to the embodiments of this invention. Those skilled in the art will understand that the illustrated structure does not constitute a limitation on the device, and may include more or fewer components than illustrated, or combine certain components, or have different component arrangements.

[0106] Please see Figure 4In another embodiment of this application, a system 10 for accelerating distributed training of large AI models based on adaptive hybrid gradient compression is provided. The system includes a model building module 11, a partitioning module 12, a first gradient processing module 13, a second gradient processing module 14, and a model training module 15.

[0107] The model building module 11 is used to build a distributed machine learning model, initialize the training dataset, partition the dataset and send it to n computing nodes (workers); the distributed training of the distributed machine learning model is based on "PS architecture + data parallelism", has a master-slave node structure, including a PS parameter server, n computing nodes and connected communication lines.

[0108] The partitioning module 12 is used to initialize model parameters w0 on the PS parameter server in a master-slave node structure and push these parameters to all computing nodes; the PS parameter server and computing nodes simultaneously perform local training using the dataset, and generate gradients after completing t iterations of training. and It is the gradient generated by the PS parameter server in the t-th iteration. It is the gradient generated by the nth computation node in the tth iteration;

[0109] The first gradient processing module 13 is used by each computing node to process the gradient it generates. Gradient compression is performed using an adaptive hybrid gradient compression method to generate compressed gradients. The adaptive hybrid gradient compression method specifically includes: gradient There are multiple classes. For a given class of gradients, if it is a one-dimensional vector, no compression is performed; if its number is greater than or equal to a set threshold, a class-by-class gradient clustering mechanism is used to cluster the gradients. Perform compression and set the compression mode identifier M=0; when its quantity is less than a set threshold, use the Tensor Train decomposition mechanism to process the gradient. Compression is performed, and the compression mode identifier M = 1 is set; the gradient after compression via the class-by-class gradient clustering mechanism This includes the centroid value and cluster identifier; the gradient compressed using the Tensor-Train decomposition mechanism. It includes a set of decomposed tensors and their corresponding singular values;

[0110] The first gradient processing module 14 is used to process the compressed gradient. Send to the PS parameter server; the PS parameter server processes the compressed gradient according to the compression mode identifier. Decompress the data to recover the original gradient, and then compare it with the original gradient. Aggregate the gradients to obtain the global gradient g. t ; Utilizing the global gradient g t Update the parameters of the distributed machine learning model to obtain the new model parameters.

[0111] The model training module 15 is used to train new model parameters. The new model parameters are sent to the computing nodes. After each computing node receives the new model parameters, it replaces its own old model parameters, and one iteration ends. After multiple iterations, the distributed training of the distributed machine learning model is completed.

[0112] It should be noted that the system for accelerating distributed training of large AI models based on adaptive hybrid gradient compression of the present invention corresponds one-to-one with the method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression of the present invention. The technical features and beneficial effects described in the embodiments of the method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression are applicable to the embodiments of the system for accelerating distributed training of large AI models based on adaptive hybrid gradient compression. For details, please refer to the description in the embodiments of the method of the present invention, which will not be repeated here.

[0113] Furthermore, in the above embodiments of the system for accelerating distributed training of large AI models based on adaptive hybrid gradient compression, the logical division of each program module is merely an example. In practical applications, the above functions can be assigned to different program modules as needed, for example, for the sake of corresponding hardware configuration requirements or the convenience of software implementation. That is, the internal structure of the system for accelerating distributed training of large AI models based on adaptive hybrid gradient compression can be divided into different program modules to complete all or part of the functions described above.

[0114] Please see Figure 5 In one embodiment, an electronic device is provided for implementing a method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression. The electronic device 20 may include a first processor 21, a first memory 22, and a bus, and may also include a computer program stored in the first memory 22 and executable on the first processor 21, such as an accelerated distributed training program for large AI models based on adaptive hybrid gradient compression 23.

[0115] The first memory 22 includes at least one type of readable storage medium, including flash memory, portable hard drive, multimedia card, card-type memory (e.g., SD or DX memory), magnetic memory, magnetic disk, optical disk, etc. In some embodiments, the first memory 22 can be an internal storage unit of the electronic device 20, such as the portable hard drive of the electronic device 20. In other embodiments, the first memory 22 can be an external storage device of the electronic device 20, such as a plug-in portable hard drive, smart media card (SMC), secure digital card (SD), flash card, etc., equipped on the electronic device 20. Furthermore, the first memory 22 can include both internal storage units and external storage devices of the electronic device 20. The first memory 22 can be used not only to store application software and various types of data installed on the electronic device 20, such as the code of the accelerated AI large model distributed training program 23 based on adaptive hybrid gradient compression, but also to temporarily store data that has been output or will be output.

[0116] In some embodiments, the first processor 21 may be composed of integrated circuits, such as a single packaged integrated circuit or multiple integrated circuits with the same or different functions, including combinations of one or more central processing units (CPUs), microprocessors, digital processing chips, graphics processors, and various control chips. The first processor 21 is the control unit of the electronic device, connecting various components of the entire electronic device through various interfaces and lines. It executes programs or modules stored in the first memory 22 and calls data stored in the first memory 22 to perform various functions of the electronic device 20 and process data.

[0117] Figure 5 Only electronic devices with components are shown; it will be understood by those skilled in the art that... Figure 5 The structure shown does not constitute a limitation on the electronic device 20, and may include fewer or more components than shown, or combine certain components, or have different component arrangements.

[0118] The distributed training program 23 for accelerating large AI models based on adaptive hybrid gradient compression, stored in the first memory 22 of the electronic device 20, is a combination of multiple instructions. When run in the first processor 21, it can achieve the following:

[0119] A distributed machine learning model is constructed, a training dataset is initialized, the training dataset is divided into n regions and sent to n computing nodes (workers); the distributed training of the distributed machine learning model is based on "PS architecture + data parallelism", and has a master-slave node structure, including a PS parameter server, n computing nodes and connected communication lines.

[0120] In a master-slave node architecture, the PS parameter server initializes the model parameters w0 and pushes these parameters to all compute nodes. The PS parameter server and compute nodes simultaneously perform local training using the dataset, generating gradients after t iterations of training. and It is the gradient generated by the PS parameter server in the t-th iteration. It is the gradient generated by the nth computation node in the tth iteration;

[0121] Each computing node generates its own gradient Gradient compression is performed using an adaptive hybrid gradient compression method to generate compressed gradients. The adaptive hybrid gradient compression method specifically includes: gradient There are multiple classes. For a given class of gradients, if it is a one-dimensional vector, no compression is performed; if its number is greater than or equal to a set threshold, a class-by-class gradient clustering mechanism is used to cluster the gradients. Perform compression and set the compression mode identifier M=0; when its quantity is less than a set threshold, use the Tensor Train decomposition mechanism to process the gradient. Compression is performed, and the compression mode identifier M = 1 is set; the gradient after compression via the class-by-class gradient clustering mechanism This includes the centroid value and cluster identifier; the gradient compressed using the Tensor-Train decomposition mechanism. It includes a set of decomposed tensors and their corresponding singular values;

[0122] The compressed gradient Send to the PS parameter server; the PS parameter server processes the compressed gradient according to the compression mode identifier. Decompress the data to recover the original gradient, and then compare it with the original gradient. Aggregate the gradients to obtain the global gradient g. t ; Utilizing the global gradient g t Update the parameters of the distributed machine learning model to obtain the new model parameters.

[0123] New model parameters The new model parameters are sent to the computing nodes. After each computing node receives the new model parameters, it replaces its own old model parameters, and one iteration ends. After multiple iterations, the distributed training of the distributed machine learning model is completed.

[0124] Furthermore, if the modules / units integrated in the electronic device 20 are implemented as software functional units and sold or used as independent products, they can be stored in a non-volatile computer-readable storage medium. The computer-readable medium may include: any entity or device capable of carrying the computer program code, a recording medium, a USB flash drive, a portable hard drive, a magnetic disk, an optical disk, a computer memory, or a read-only memory (ROM).

[0125] Those skilled in the art will understand that all or part of the processes in the methods of the above embodiments can be implemented by a computer program instructing related hardware. The program can be stored in a non-volatile computer-readable storage medium, and when executed, it can include the processes of the embodiments of the above methods. Any references to memory, storage, databases, or other media used in the embodiments provided in this application can include non-volatile and / or volatile memory. Non-volatile memory can include read-only memory (ROM), programmable ROM (PROM), electrically programmable ROM (EPROM), electrically erasable programmable ROM (EEPROM), or flash memory. Volatile memory can include random access memory (RAM) or external cache memory. By way of illustration and not limitation, RAM is available in various forms, such as static RAM (SRAM), dynamic RAM (DRAM), synchronous DRAM (SDRAM), dual data rate SDRAM (DDRSDRAM), enhanced SDRAM (ESDRAM), synchronous link DRAM (SLDRAM), RAMbus direct RAM (RDRAM), direct memory bus dynamic RAM (DRDRAM), and RAMbus dynamic RAM (RDRAM), etc.

[0126] The technical features of the above embodiments can be combined in any way. For the sake of brevity, not all possible combinations of the technical features in the above embodiments are described. However, as long as there is no contradiction in the combination of these technical features, they should be considered to be within the scope of this specification.

[0127] The above embodiments are preferred embodiments of the present invention, but the embodiments of the present invention are not limited to the above embodiments. Any changes, modifications, substitutions, combinations, or simplifications made without departing from the spirit and principle of the present invention shall be considered equivalent substitutions and shall be included within the protection scope of the present invention.

Claims

1. A method for accelerating distributed training of an AI large model based on adaptive mixed gradient compression, characterized in that, Includes the following steps: A distributed machine learning model is constructed, a training dataset is initialized, the training dataset is divided into n regions and sent to n computing nodes (workers); the distributed training of the distributed machine learning model is based on "PS architecture + data parallelism", and has a master-slave node structure, including a PS parameter server, n computing nodes and connected communication lines. In a master-slave node structure, the PS parameter server initializes the model parameters w0 and pushes these parameters to all computing nodes; The PS parameter server and compute nodes simultaneously perform local training using the dataset, generating gradients after t iterations of training. and It is the gradient generated by the PS parameter server in the t-th iteration. It is the gradient generated by the nth computation node in the tth iteration; Each computing node generates its own gradient Gradient compression is performed using an adaptive hybrid gradient compression method to generate compressed gradients. The adaptive hybrid gradient compression method specifically includes: gradient There are multiple classes. For a given class of gradients, if it is a one-dimensional vector, no compression is performed; if its number is greater than or equal to a set threshold, a class-by-class gradient clustering mechanism is used to cluster the gradients. Perform compression and set the compression mode identifier M=0; when its quantity is less than a set threshold, use the Tensor Train decomposition mechanism to process the gradient. Perform compression and set the compression mode identifier M=1; Gradient compressed by the class-by-class gradient clustering mechanism Including centroid value and cluster identifier; Gradient compressed by Tensor-Train decomposition mechanism It includes a set of decomposed tensors and their corresponding singular values; The compressed gradient Send to the PS parameter server; the PS parameter server processes the compressed gradient according to the compression mode identifier. Decompress the data to recover the original gradient, and then compare it with the original gradient. Aggregate the gradients to obtain the global gradient g. t ; Utilizing the global gradient g t Update the parameters of the distributed machine learning model to obtain the new model parameters. New model parameters The new model parameters are sent to the computing nodes. After each computing node receives the new model parameters, it replaces its own old model parameters, and one iteration ends. After multiple iterations, the distributed training of the distributed machine learning model is completed.

2. The method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression according to claim 1, characterized in that, The class-by-class gradient clustering mechanism is as follows: gradient There are multiple classes, and each class has multiple gradient values. A clustering algorithm is used to cluster each class of gradients into multiple clusters. Each cluster has a centroid value and a corresponding cluster identifier, resulting in multiple centroid values ​​and cluster identifiers. These multiple centroid values ​​and cluster identifiers represent the gradients compressed through the class-wise gradient clustering mechanism.

3. The method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression according to claim 1, characterized in that, The Tensor Train decomposition mechanism is as follows: gradient There are multiple classes. One class of gradients is an N-order tensor X. Each class of gradients is decomposed into Tensor Train based on singular value decomposition. After N-1 singular value decompositions, N small tensors and corresponding singular spectra are obtained. Each singular spectrum contains some singular values. The number of non-zero singular values ​​in the singular spectrum is extracted to obtain the TT rank. Taking a low-rank approximation of the TT rank, for a given tensor X, the Tensor-Train decomposition makes: At this point, the optimal TT low-rank approximation of tensor X is: Among them, a n s represents the nth geometric or auxiliary index. N Indicates the Nth open or physical index. This represents the Nth tensor, where the N tensors and their corresponding singular values ​​are the gradients compressed using the TT decomposition mechanism.

4. The method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression according to claim 1, characterized in that, The PS parameter server processes the compressed gradient based on the compression mode identifier. Decompress, recover the original gradient and compare it with the gradient. The aggregation process includes the following steps: When the compression mode identifier M=1, a splicing-based tensor recovery mechanism is used to process the compressed gradient. Decompress, recover the original gradient and compare it with the gradient. Perform aggregation; When the compression mode identifier M = 0, a gradient recovery mechanism based on centroid replacement is used to recover the compressed gradient. Recover the original gradient and compare it with the gradient Perform aggregation.

5. The method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression according to claim 4, characterized in that, The splicing-based tensor recovery mechanism is specifically as follows: By reversing the order of the smaller tensors according to the Tensor-Train decomposition method, we obtain an N-order tensor, which is an approximate reconstruction of the original tensor. in, This represents the i-th element in the k-th small tensor obtained from Tensor-Train decomposition. k Dimension, and × k Indicates will Perform a Kronecker product with the previous result, where 1 ≤ k ≤ N.

6. The method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression according to claim 4, characterized in that, The gradient recovery mechanism based on centroid replacement is specifically as follows: The PS parameter server replaces all gradient values ​​within a cluster with the centroid value of that cluster, thereby approximately restoring the original gradient.

7. The method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression according to claim 1, characterized in that, The use of global gradient g t Update the parameters of the distributed machine learning model to obtain the new model parameters. Specifically: Using the global gradient g t The parameters of the distributed machine learning model are updated according to the parameter update formula to obtain the new model parameters. The parameter update formula is as follows: in, η represents the model parameters of the PS parameter server at the t-th iteration, and η represents the learning rate at the time of update, 0≤η≤1.

8. A system for accelerating distributed training of large AI models based on adaptive hybrid gradient compression, characterized in that, The method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression, as described in any one of claims 1-7, includes a model building module, a partitioning module, a first gradient processing module, a second gradient processing module, and a model training module. The model building module is used to build a distributed machine learning model, initialize the training dataset, partition the dataset and send it to n computing nodes (workers); the distributed training of the distributed machine learning model is based on "PS architecture + data parallelism" and has a master-slave node structure, including a PS parameter server, n computing nodes and connected communication lines. The partitioning module is used to initialize model parameters w0 on the PS parameter server in a master-slave node structure and push these parameters to all computing nodes. The PS parameter server and compute nodes simultaneously perform local training using the dataset, generating gradients after t iterations of training. and It is the gradient generated by the PS parameter server in the t-th iteration. It is the gradient generated by the nth computation node in the tth iteration; The first gradient processing module is used by each computing node to process the gradient it generates. Gradient compression is performed using an adaptive hybrid gradient compression method to generate compressed gradients. The adaptive hybrid gradient compression method specifically includes: gradient There are multiple classes. For a given class of gradients, if it is a one-dimensional vector, no compression is performed; if its number is greater than or equal to a set threshold, a class-by-class gradient clustering mechanism is used to cluster the gradients. Perform compression and set the compression mode identifier M=0; when its quantity is less than a set threshold, use the Tensor Train decomposition mechanism to process the gradient. Perform compression and set the compression mode identifier M=1; Gradient compressed by the class-by-class gradient clustering mechanism Including centroid value and cluster identifier; Gradient compressed by TensorTrain decomposition mechanism It includes a set of decomposed tensors and their corresponding singular values; The first gradient processing module is used to process the compressed gradient. Send to the PS parameter server; The PS parameter server processes the compressed gradient based on the compression mode identifier. Decompress the data to recover the original gradient, and then compare it with the original gradient. Aggregate the gradients to obtain the global gradient g. t ; Utilizing the global gradient g t Update the parameters of the distributed machine learning model to obtain the new model parameters. The model training module is used to train new model parameters. The new model parameters are sent to the computing nodes. After each computing node receives the new model parameters, it replaces its own old model parameters, and one iteration ends. After multiple iterations, the distributed training of the distributed machine learning model is completed.

9. An electronic device, characterized in that, The electronic device includes: At least one processor; and, A memory communicatively connected to the at least one processor; wherein, The memory stores computer program instructions that can be executed by the at least one processor to enable the at least one processor to perform the method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression as described in any one of claims 1-7.

10. A computer-readable storage medium storing a program, characterized in that, When the program is executed by the processor, it implements the method for accelerating distributed training of large AI models based on adaptive hybrid gradient compression as described in any one of claims 1-7.

Images