A distributed training method for hybrid expert models based on dynamic load balancing

Abstract

The present invention discloses a distributed training method for a hybrid expert model based on dynamic load balancing. The method first captures token features, expert parameters, memory and computing resource information, and communication status through a sampler to form sampled data. Secondly, a MoE training load model constructed based on the sampled data is used to obtain the total memory resources and computing resources of each node. Load balancing is then achieved through expert scheduling to reduce all-to-all communication delays. Finally, an expert capacity model is constructed to set capacity values ​​for different experts. The present invention realizes dynamic adjustment of expert capacity, improves the resource utilization of equipment, reduces communication overhead, and resolves the contradiction between training accuracy and resource consumption.

Description

Technical Field

[0001] The present invention belongs to the technical field of neural network, and in particular relates to a hybrid expert model distributed training method based on dynamic load balancing. Background Art

[0002] In recent years, deep neural networks based on the transformer architecture have achieved great success. Increasing the model size to improve the performance of artificial intelligence models has become one of the mainstream trends in the development of artificial intelligence. With the development of models based on the transformer architecture, progress has been made in many fields, such as computer vision and natural language processing (e.g., VIT, BERT), and the scale of model parameters has also increased from hundreds of billions to more than trillions (e.g., llama2, GPT-3, M6). However, these dense models usually require a lot of computing resources and training time. The Mixture-of-Experts (MoE) model uses multiple sub-models (experts) instead of dense Feed-Forward Network (FNN) layers and is a sparse activation model architecture. It has been proven that large models can be designed with this sparse architecture to achieve the goal of enhancing model performance without significantly increasing computing resources.

[0003] The MoE architecture uses many sub-models, called experts, to process data or tasks, rather than a single, dense, large model. Each expert handles a specific type of data or task. During MoE model training, different experts are assigned to different GPU accelerators, and a gating network is used to intelligently distribute input tokens to specific experts. The MoE architecture results in a significant amount of "all-to-all" communication between GPU accelerators, including distributing tokens to specific experts and collecting computation results after token calculations are completed. If an expert is broken down and deployed across different cards, known as expert parallelism, the cost of all-to-all communication increases. Furthermore, differences in data distribution and the types of data processed by experts can cause load imbalances between different experts.

[0004] Existing methods are mainly implemented using Pytorch, relying on the dynamic shape function of dynamic graphs to enable each expert to adapt to the differences in data or tasks in different iterations. The dynamic shape function of dynamic graphs can support changing the input shape in different iterations of the model. Static graphs are another mainstream computational graph used by machine learning frameworks for deep learning model training. Compared with dynamic graphs, they have two major advantages: high computing performance and direct deployment. In order to take advantage of the advantages of static graphs to train the MoE model, the MoE model was trained and analyzed based on MindSpore, and the following problems were discovered: 1) In the process of implementing expert parallel training using a distributed method, popular experts receive more tokens than other experts, which makes the device where they are located too heavy, while the utilization rate of other devices is low, resulting in uneven load. 2) Since expert parallel training requires cross-node transmission of tokens, it brings a lot of communication time and reduces the efficiency of MoE end-to-end training.

[0005] 3) Static graphs only support static shape functions, which requires setting a fixed shape size for the input in advance, and no changes are allowed during the iteration process. When using static graphs to implement MoE model training, it is necessary to set a fixed token capacity for each expert in advance, that is, the expert capacity. When an expert with a small amount of data or token load needs to process less data than the set expert capacity, it needs to fill in 0 to reach the set expert capacity, which will result in a large waste of computing and storage resources because it needs to fill in invalid data to maintain a fixed shape. If the amount of data to be processed exceeds the expert capacity, the excess data will be discarded, which will affect the accuracy of the model. Summary of the Invention

[0006] The present invention addresses three problems of the MoE model based on static graph training: 1) load imbalance between GPU accelerators caused by different popularity of experts; 2) communication delay introduced by All-to-All;

[0007] 3) The conflict between model accuracy and resource utilization. This paper proposes a distributed training method for hybrid expert models based on dynamic load balancing, implemented using a static computation graph and the MindSpore framework. First, this method constructs a sampler to collect features such as model structure, experts, and data distribution. Based on the sampled data, a load prediction model is proposed to estimate and predict expert load. Second, experts are classified into popular and unpopular experts based on their load. Furthermore, a dynamic expert scheduling method is employed to balance the load across GPU accelerators. This method generates replica experts for popular experts and dynamically schedules them to other GPU accelerators based on their real-time load and node idle resources. Replica experts directly process local tokens sent to popular experts, reducing the frequency of all-to-all communication and improving model training efficiency. Finally, an expert capacity model is designed to set appropriate expert capacities for popular and unpopular experts. The present invention sets larger capacities for popular experts and their replicas to reduce token discard, and smaller capacities for unpopular experts to minimize resource waste. This method improves the training efficiency of MoE models and reduces training costs.

[0008] The present invention provides a distributed training method for a hybrid expert model based on dynamic load balancing. The analysis method is as follows:

[0009] Step 1: To evaluate and optimize the training performance of the MoE model, this paper analyzes the structure of the MoE model and the key factors affecting its training performance, and uses a sampler to capture important information to form sampled data. The collected information includes the following parts:

[0010] Token features: refers to the features or attributes of each representation (such as words or subwords) in the dataset.

[0011] Expert parameters: These include weights and biases associated with each expert, as well as intermediate data such as activations and gradients that need to be stored during forward and backward passes.

[0012] Memory and computing resource information: The computing and memory resources of the nodes in the cluster and the GPUs of different nodes.

[0013] Communication: The rate of data transfer between nodes (inter-node communication) and within nodes (intra-node communication) in a distributed computing environment.

[0014] Step 2: In order to effectively achieve load balancing and communication optimization, the present invention designs a MoE training load model. The MoE training load model constructed based on the sampled data obtains the memory and computing resources of each node. First, using the sampler, the mapping relationship between each expert and the token can be obtained, and a statistical method is used to count the number of tokens processed by different experts in different layers in different iterations. On this basis, the tokens are sorted in descending order, and the top p% of the data are selected as valid data. This step aims to exclude extreme outliers and achieve a reliable assessment of the expert workload. At the same time, the concept of load prediction cycle is introduced, m iterations are regarded as a cycle, and the average load of the cycle is used to predict the load of the i-th expert in the k-th layer in the future cycle, which is recorded as The sampler then samples each expert's parameters to determine the load on each GPU accelerator. The sampler is then used again to obtain the memory and computing resources of each node. This ensures that the computing and memory resource requirements of the GPU accelerator devices are met during the dynamic scheduling of experts.

[0015] Step 3: Based on the MoE training load prediction model, an expert scheduling method is designed to achieve load balancing and reduce All-to-All communication latency. It aims to dynamically schedule experts and effectively manage the load distribution in the MoE model. The main idea is to evaluate the token load of different experts within a load prediction cycle and classify them according to their load factors. Experts with a load factor greater than a threshold q are classified as popular experts, while others are considered unpopular experts. In this invention, the value of q is set to 60%. The popular experts are then scheduled to the GPU accelerator that meets the requirements, balancing the load between GPUs, and localizing token processing to reduce the frequency of All-to-All communication.

[0016] Step 4: Due to the limitations of static graph design, all experts have the same expert capacity. If the popular expert has a small expert capacity, a large number of tokens will be discarded, resulting in reduced training accuracy. Unpopular experts need to fill too many 0 vectors due to their large expert capacity, resulting in a waste of hardware resources. Therefore, an expert capacity model is constructed to set capacity values ​​for different experts. First, the baseline capacity B is obtained by counting the tokens of each expert in the sampling period of m iterations. Secondly, the capacity coefficient r is obtained by obtaining the average number of tokens that expert i needs to process in m iterations through the sampler. Finally, the number of tokens processed by the sampler is T. i, to obtain the capacity factor F, thereby achieving dynamic adjustment of expert capacity. This can avoid large amounts of token discard or filling, and improve the resource utilization of the GPU accelerator by allowing multiple experts to fully share resources on the same GPU accelerator.

[0017] Through the methods in steps 3 and 4, the cost of MoE distributed training is reduced in three key ways: First, by replicating popular experts to the remaining devices, the amount of data required to be processed by each device is balanced, balancing the memory and computing resource overhead between devices. Second, by localizing token processing, the frequency of all-to-all communication is reduced, lowering communication overhead. Third, by dynamically adjusting the expert capacity, the number of zero vectors required to supplement unpopular experts is reduced, reducing computing resource overhead.

[0018] Beneficial effects of the present invention

[0019] A sampler is designed to sample token features, expert parameters, computing and memory resources, and communication conditions during the distributed training of the MoE model. Based on this information, a MoE load prediction model is established, laying the foundation for subsequent expert scheduling and dynamic expert capacity adjustment.

[0020] Based on the MoE load prediction model, an expert scheduling method was designed to address load imbalance and all-to-all communication latency. A load evaluation model estimates each expert's token load and classifies experts into popular and unpopular based on this coefficient. If the current GPU accelerator's load meets the requirements of a popular expert, a replica expert is dispatched to that GPU accelerator, sharing resources with the unpopular expert. This achieves load balancing and localized token processing, reducing all-to-all communication frequency.

[0021] Finally, based on the load prediction model, an expert capacity model was constructed to set appropriate capacity values ​​for different experts, realize dynamic adjustment of expert capacity, avoid large amounts of zero vector padding, improve device resource utilization, and resolve the contradiction between training accuracy and resource consumption. BRIEF DESCRIPTION OF THE DRAWINGS

[0022] Figure 1 It is the overall architectural design diagram. DETAILED DESCRIPTION

[0023] The present invention will be further described below with reference to the accompanying drawings and specific implementation steps:

[0024] Step 1: Use samplers to capture important information. The collected information is shown in Table 1.

[0025] Table 1: Sampling information

[0026]

[0027] During MoE model training, token characteristics play a crucial role in determining which experts are selected for processing. The MoE gating mechanism determines token distribution based on these characteristics, ensuring that each expert processes the most relevant token data. Expert parameters include weights and biases associated with each expert, as well as intermediate data such as activations and gradients that need to be stored during forward and backward passes. Different experts may have different computational and memory requirements based on their scale and complexity. Efficient memory management and computational resource management across cluster nodes and GPUs on different nodes are crucial for maximizing GPU utilization and preventing bottlenecks. In MoE models, efficient communication is crucial for synchronizing and sharing parameters or data between GPU accelerators. Inter-node bandwidth is relatively low, while intra-node bandwidth is relatively high. Different bandwidth frequencies will affect scalability across multiple machines. Optimizing communication bandwidth is crucial for reducing latency and improving overall training and inference speed in MoE models.

[0028] Step 2: After obtaining the mapping relationship between token and expert by using token features obtained by the sampler, assume is the number of tokens assigned to the jth expert in the kth layer in the i-th iteration. A statistical method is used to count the number of tokens processed by different experts in different layers during different iterations. The token counts processed by different experts on different layers in m iterations can be expressed as On this basis, the tokens are sorted in descending order and the first p% of the data are selected as valid data. At the same time, the concept of load prediction cycle is introduced, m iterations are regarded as a cycle, and the average load of the cycle is used to predict the expert's load in the future cycle. The load of each expert i is represented by k. As shown in Formula 1:

[0029]

[0030] in, represents the jth maximum token count after ranking expert i at layer k. The value of m is determined by the difference in token counts in the iteration. When the token count gap exceeds the threshold ω, the expert's load is recalculated, which is determined by the impact of historical data on the expert's load in different load prediction cycles. In addition to the individual expert load, the load of each GPU accelerator needs to be constructed, as shown in the following example.

[0031] As shown in formula 2:

[0032]

[0033] Among them, T i represents the number of tokens processed by expert i, Compute(T i ) and Memory(T i ) is a computational function of the resources required for computation and storage, D j Represents the load of GPU accelerator j. j Finally, the sampler is used to obtain the expert parameters and the overall load requirements, and the sampler is used to obtain the computing and memory resources of each device. This is mainly used to ensure that during the expert dynamic scheduling process, the computing resource requirements and memory resource requirements on the GPU accelerator device will be met.

[0034] Step 3: Within a load prediction cycle, the load evaluation model is used to estimate the token load of each expert, and the experts are divided into hot and unpopular according to the coefficient. The current load conditions are evaluated, including the computing and storage resources of different GPU accelerators. If the current GPU accelerator's load conditions meet the requirements of the hot expert, the replica expert is scheduled to the GPU accelerator to share resources with the unpopular expert, thereby achieving localized token processing, which means that the tokens originally intended for the hot expert will be processed by the local replica expert. The remaining tokens will also be distributed to unpopular experts as needed to minimize the frequency of All-to-All communication. In each prediction cycle, the expert load is dynamically evaluated and the scheduling of the hot expert is adjusted accordingly. The details of this scheduling strategy are in Algorithm 1:

[0035]

[0036] Step 4: An expert capacity model is constructed to set appropriate capacity values ​​for different experts.

[0037] The specific expert capacity model is as follows. Assuming B is the baseline capacity and F is the capacity factor used to dynamically adjust the expert capacity, the capacity C of each expert i is i It can be expressed as the following formula:

[0038] C i =B+F*r (3)

[0039] Where B is constructed by statistically analyzing the token counts of each expert within a sampling period of m iterations. Extreme values ​​are discarded and the average value of m iterations is used as the token load of each expert. The token loads of all experts are sorted and the values ​​covering the top p% of the number of tokens are selected as the baseline capacity b. The factor r is defined as:

[0040]

[0041] in, represents the average number of tokens that expert i needs to process in m iterations. The capacity factor F is calculated as follows:

[0042]

[0043] Among them, T i represents the number of tokens of expert i determined by the load evaluation model. For unpopular experts, T i -B is negative, which leads to a decrease in Ci. On the contrary, F is favorable for popular experts, which leads to C i Increase Total_token to the total number of tokens. In addition, you must ensure that the total memory requirement M of all experts on the same GPU accelerator is cost is smaller than the total memory of the GPU accelerator. The storage requirement of expert i is given by:

[0044] M coat =M p +M a +M t (6)

[0045] Among them, M p 、M a and M t They represent the memory overhead of expert parameters, intermediate activations, and tokens, respectively. Expert parameters and intermediate activations will be obtained through the sampler, and the memory overhead of tokens will be calculated based on C obtained from Formula 3. i The storage requirements of each expert will be compared with the total memory resources. If they do not meet the requirements, C will be reduced on the original basis. i size so that the final result meets the requirements.

[0046] The specific approach is shown in Algorithm 2: Algorithm 2 dynamically adjusts the capacity of experts in the MoE model based on the number of tokens they process. This method balances the computational load across devices by allocating more capacity to popular experts and less capacity to unpopular experts. This approach mitigates accuracy issues caused by token truncation and resource waste caused by zero padding. Furthermore, memory cost checks ensure that total memory usage remains within hardware limits, promoting efficient and scalable training of large models.

[0047]

[0048] The GPT-MoE model was used to evaluate the method of the present invention, and Fastermoe and SwitchTransformers were used as baselines. In terms of end-to-end performance, compared with Switch Transformers and Fastermoe, the present invention achieved 30% and 33% acceleration respectively; in terms of communication optimization, compared with Fastermoe and SwitchTransformers, the present invention reduced the communication time by 18% and 20% respectively. In addition, the advantages of dynamic expert capacity were also verified. First, the accuracy and PPL values ​​were tested. The results confirmed that the dynamic expert capacity adjustment method did not damage the accuracy and performance of the model. After the expert capacity was adjusted, the performance of the present invention was compared with Fastermoe. The present method improved the training efficiency by up to 35%, indicating that the supplement of the zero vector of the token was reduced and the consumption of computing resources was reduced.

Claims

1. A distributed training method for hybrid expert models based on dynamic load balancing, characterized in that: The following steps are involved: Step 1: Capture token features, expert parameters, memory and computing resource information, and communication status through the sampler to form sampled data; Step 2: Based on the MoE training load model built with sampled data, obtain the memory and computing resources of each node; Step 3: Implement load balancing through expert scheduling to reduce all-to-all communication delays; Step 4: Construct an expert capacity model and set capacity values ​​for different experts. The specific implementation process is as follows: first, obtain the baseline capacity B by counting the tokens of each expert in the sampling period of m iterations; second, obtain the capacity coefficient r by using the sampler to obtain the average number of tokens that expert i needs to process in m iterations; finally, obtain the number of tokens processed by the sampler T i , we get the capacity factor F, and realize the dynamic adjustment of expert capacity. By allowing multiple experts to fully share resources on the same GPU accelerator, we can improve the resource utilization of the GPU accelerator. B is the baseline capacity, and F is the capacity factor, which is used to dynamically adjust the expert capacity. Then the capacity C of each expert i is i It is expressed as the following formula: C i =B+F*r Where B is constructed by statistically analyzing the token counts of each expert within a sampling period of m iterations. Extreme values ​​are discarded and the average value of m iterations is used as the token load of each expert. The token loads of all experts are sorted and the values ​​covering the top p% of the number of tokens are selected as the baseline capacity B. The factor r is defined as: in represents the average number of tokens that expert i needs to process in m iterations; the capacity coefficient F is calculated as follows: Where T i represents the number of tokens of expert i determined by the load evaluation model. For unpopular experts, T i -B is negative, leading to C i Decreases; on the contrary, F is favorable to popular experts, leading to C i Increase; Total_token is the total number of tokens; in addition, the total memory requirement M of all experts on the same GPU accelerator cost is smaller than the total memory of the GPU accelerator, the storage requirement of expert i is given by: M cost =M p +M a +M t Among them, M p 、M a and M t They represent the memory overhead of expert parameters, intermediate activations, and tokens, respectively. Expert parameters and intermediate activations will be obtained through the sampler, and the memory overhead of tokens will be calculated based on C i Get, compare the storage requirements of each expert with the total memory resources, and when they do not meet the requirements, reduce C on the original basis. i size so that the final result meets the requirements.

2. The distributed training method for hybrid expert models based on dynamic load balancing according to claim 1, characterized in that: The token features refer to the features or attributes of each representation in the dataset; The expert parameters include weights and biases associated with each expert, and intermediate data stored during forward and backward passes; The memory and computing resource information: computing resources and memory resources of the nodes in the cluster and the GPUs of different nodes; The communication situation: the data transmission rate between nodes and within nodes in a distributed computing environment.

3. The distributed training method for hybrid expert models based on dynamic load balancing according to claim 2, characterized in that: The specific implementation process of step 2 is as follows: Using the sampler, we get the mapping relationship between each expert and token, and use statistical methods to count the number of tokens processed by different experts in different layers during different iterations; Sort the tokens in descending order and select the first p% of the data as valid data; Consider m iterations as a cycle and use the average load of this cycle to predict the load of experts in the future cycle to represent the load of each expert i on layer k The expert parameters of each expert are sampled using a sampler to obtain the load of each GPU accelerator; Use samplers to obtain the memory and computing resources of each node.

4. The distributed training method for hybrid expert models based on dynamic load balancing according to claim 3 is characterized in that: The load As shown in the following formula: in, represents the jth maximum token count after ranking expert i at layer k. The value of m is determined by the difference in token counts in the iteration. When the token count difference exceeds the threshold ω, the expert's load is recalculated. It is determined based on the impact of historical data on the expert's load in different load prediction cycles. The load of each GPU accelerator is shown in the following formula: Where T i represents the number of tokens processed by expert i, Compute(T i ) and Memory(T i ) is a computational function of the resources required for computation and storage, D j represents the load of GPU accelerator j; In getting D j Finally, the sampler is used to obtain the expert parameters and the overall load requirements, and the sampler is used to obtain the computing and memory resources of each device.

5. The distributed training method for hybrid expert models based on dynamic load balancing according to claim 4 is characterized in that: The specific implementation process of step 3 is as follows: the token loads of different experts are evaluated within a load prediction cycle and classified according to their load factors; experts with load factors greater than a threshold q are classified as popular experts, while others are considered unpopular experts. Popular experts are dispatched to GPU accelerators that meet the requirements, balancing the load between GPUs, localizing token processing, and reducing the frequency of all-to-all communication.

Images