Hybrid parallel and dynamic scheduling method of hybrid expert model based on 3D near-memory processing

Abstract

The invention provides a hybrid parallel and dynamic scheduling method of a hybrid expert model based on 3D near-memory processing. The method comprises the following steps: establishing a joint performance analysis model; performing off-line linear programming optimization expert distribution; performing Bayesian optimization on physical node mapping; performing online reasoning; carrying out online dynamic priority detection; and an expert pre-broadcast and communication-friendly lexical element distribution strategy with optimal efficiency is provided. According to the method, node balancing optimization is realized through offline linear programming, and the problem of load imbalance of 3D NMP calculation is remarkably improved; in combination with a Bayesian optimization mapping strategy of link balance, the communication speed-up ratio is increased, and NoC link congestion is reduced; a dynamic scheduling strategy adapts to dynamic changes of expert activation in real-time reasoning through calculation load prediction and a pre-broadcast mechanism. Through cooperation of the offline automatic hybrid parallel mapping algorithm and the online dynamic scheduling strategy, the calculation load and the communication overhead are effectively balanced, and the reasoning efficiency of the hybrid expert model MoE on the 3D near-memory processing architecture is remarkably improved.

Description

Technical Field

[0001] The present invention relates to the field of deep learning and large language model (LLM) optimization, and specifically to an efficient reasoning method for a mixture of experts (MoE), which is particularly suitable for computing resource scheduling and communication optimization under a 3D near-memory processing (3D NMP) architecture. Background Art

[0002] In recent years, with the widespread application of LLM in natural language processing tasks, the high amount of model computation has become a major challenge. MoE has become the core solution to improve computing efficiency through the sparse activation mechanism. This model can significantly reduce the number of parameters actually involved in the calculation while maintaining the scale of model parameters. This design enables models with hundreds of billions of parameters to be trained and inferred under limited computing power. At the same time, through expert division of labor, the model's generalization ability for multiple tasks (such as translation, reasoning, and code generation) is improved, becoming a key architecture for breaking through performance and efficiency bottlenecks in the current LLM field. However, this mechanism exacerbates the pressure on memory bandwidth. Especially in memory-constrained scenarios such as edge devices, frequent expert parameter access leads to the "memory wall" problem, which becomes a key bottleneck restricting the deployment of MoE.

[0003] The 3D near-memory processing (3D NMP) architecture that has emerged in recent years provides high-bandwidth, low-latency data access capabilities by vertically stacking dynamic random access memory (DRAM) and logic units. Its parallel memory access design can support efficient data throughput and is considered an ideal solution to the memory-constrained problem. However, the distributed memory and network-on-chip (NoC) architecture of 3D NMP introduces new mapping challenges: traditional parallel strategies have difficulty balancing computational load and communication overhead. Specifically, tensor parallelism splits the expert parameter tensor into multiple computing unit nodes. Although it achieves load balancing, its all-reduce communication needs to go through the router (Router, i.e., on the NoC) on the NoC. Figure 2 、 Figure 3 The symbol “R” in the figure) aggregates partial sums on different computing unit nodes, which has a significant cost in a distributed architecture; expert parallelism distributes complete experts to independent nodes. Although it reduces the communication volume of the NoC, the uneven frequency of expert activation leads to a serious imbalance in node load and reduces the utilization of computing resources.

[0004] To further complicate matters, the dynamic routing mechanism of MoE makes expert activation patterns highly uncertain. Traditional static allocation strategies may not be able to adapt to this dynamic nature in real time, further amplifying the shortcomings of this distributed architecture. Furthermore, the memory capacity limitations of edge inference scenarios make the load balancing strategy of expert replication in traditional graphics processing unit (GPU) clusters unfeasible. Therefore, a hybrid optimization solution that adapts to the characteristics of distributed architectures is urgently needed.

[0005] While existing research attempts to mitigate dynamic challenges by leveraging activation similarity through expert prefetching, it fails to fully integrate the hardware topology characteristics of 3D NMPs, making it difficult to simultaneously optimize both computational load and link communication. Therefore, designing a parallel strategy and scheduling method that balances computational balance, communication efficiency, and dynamic adaptability has become a core technical challenge for efficient MoE model inference in 3D NMP architectures. Summary of the Invention

[0006] To address the distributed inference and dynamic load imbalance issues of the Hybrid Model of Experts (MoE) under the 3D Near Memory Processing (3D NMP) architecture, a collaborative optimization method combining offline hybrid parallel mapping and online dynamic scheduling is proposed to balance the computing load, reduce communication overhead, and improve the MoE inference efficiency of edge devices.

[0007] The technical solutions provided by the present invention are as follows:

[0008] A hybrid parallel and dynamic scheduling method based on a hybrid expert model for 3D near-memory processing, characterized by comprising the following steps:

[0009] Step 1: Establish a joint performance analysis model and build a unified analysis framework that includes computing and communication;

[0010] Given an autoregressive hybrid expert model M, consisting of L layers, each layer has E experts, and each word token activates e experts; the hidden layer dimension of the expert is h, and the intermediate layer dimension is IS; the computing power of each computing unit node is comp, the interconnection bandwidth of the on-chip network NoC between different computing unit nodes is BW, and there are D computing unit nodes in total; the batch size of inference is B; use f i Denotes the historical average activation frequency of expert i, with f g P represents the historical average activation frequency of expert group g. ic It represents the proportion of the i-th expert deployed on the c-th computing unit node; the computational overhead is calculated by the formula Quantification reflects the maximum value of node computing load; communication overhead is approximated by linear model Estimate, the two are combined to form the node-level inference delay objective function t node_overhead =t comp +2γt comm , where γ is the scaling factor determined by linear regression between the linear approximation model and the actual communication simulation time;

[0011] Step 2: Offline linear programming optimization expert allocation

[0012] Solve the expert parameter allocation ratio P through linear programming ic (0≤P ic ≤1), to minimize the node-level inference latency objective function t node_overhead As the goal, the solution automatically implements a hybrid parallel strategy;

[0013] Step 3: Bayesian optimization of physical node mapping

[0014] P based on linear programming solution ic ,Using Bayesian optimization, we optimize the NoC communication efficiency by minimizing link congestion, and find the optimal mapping strategy for mapping logical computing unit nodes to actual physical nodes for 3D near-memory processing;

[0015] Step 4: Online Inference

[0016] Deploy the expert to the 3D near-memory processing according to the mapping strategy in step 3, and start online inference of the 3D near-memory processing. Each autoregressive iteration of inference starts from the first layer of the hybrid expert model.

[0017] Step 5: Online dynamic priority detection

[0018] The hybrid expert model M contains L layers. During online inference, the expert load of the l+1 layer is predicted in real time at the lth layer, and the expert priority score of the l+1 layer is calculated. in is the expert activation frequency predicted based on the intermediate activation of layer l;

[0019] Step 6: Expert pre-broadcast for optimal efficiency

[0020] After obtaining the priority scores, we pre-broadcast as many experts with expected heavy loads in layer l+1 as possible to all memory banks within the time it takes to infer one layer. We also determine the optimal size of the pre-broadcast data block based on the α-β communication model of the network-on-chip (NoC).

[0021] Step 7: Communication-friendly word token distribution strategy

[0022] After the expert pre-broadcast, each token can be routed to any node in the l+1 layer that holds a copy of its activated expert. Nodes that already have the routed expert are selected as routing candidates. Among these candidate nodes, the token is dispatched to the node with the lowest current computational load.

[0023] Step 8: After receiving the token, the corresponding computing unit node of layer l starts the online inference of layer l+1, and repeats steps 5-7 for layer l+2, repeating the cycle until the online dynamic scheduling of layer L of the hybrid expert model is completed, ending one iteration of the inference process.

[0024] Furthermore, the constraints of the linear programming in step 2 include:

[0025] (1) Expert parameter allocation completeness: ∑ c P ic =1;

[0026] (2) Node load balancing: in

[0027] Furthermore, the hybrid parallel strategy implemented in step 2 is:

[0028] High-frequency experts tend to be distributed across multiple nodes, which is a tendency toward tensor parallelism; whereas low-frequency experts tend to be concentrated on a single node, which is a tendency toward expert parallelism.

[0029] Furthermore, the optimal size c of the pre-broadcast data block in step 6 satisfies: Among them, α and β are the communication model parameters of the on-chip network, and k is the number of experts that can be pre-broadcasted within the inference time of one layer.

[0030] The technical effects of the present invention are as follows:

[0031] (1) Node balancing optimization is achieved through offline linear programming, which significantly improves the problem of 3D NMP computational load imbalance. Compared with expert parallelism (EP), the computational speedup is about 2 times. Under different configurations, the speed is increased by 1.0 to 3.0 times compared with tensor parallelism (TP), expert parallelism (EP) and hybrid baseline (only considering expert tensor parallelism with computational load balancing), effectively alleviating the load imbalance of the MoE model.

[0032] (2) Combining the Bayesian optimization mapping strategy with link balancing, a communication speedup ratio of approximately 1.2 times is achieved, which reduces the NoC link congestion and makes the communication link load distribution more balanced, further improving the end-to-end performance compared to the node balancing scheme alone.

[0033] (3) The dynamic scheduling strategy adapts to the dynamic changes of expert activation in real-time reasoning through computational load prediction and pre-broadcast mechanism. In complex tasks such as mathematics and coding, it increases the speed by an average of 1.15 to 1.25 times compared with static deployment, significantly reducing the MoE layer latency.

[0034] (4) Overall, this patented method effectively balances computational load and communication overhead through the collaboration of an offline automatic hybrid parallel mapping algorithm and an online dynamic scheduling strategy, significantly improving the inference efficiency of the MoE model on a 3D near-memory processing architecture. Experiments show that compared to traditional tensor parallelism (TP), expert parallelism (EP), and hybrid baseline strategies, this patented method achieves an end-to-end speed improvement of 1.1 to 1.8 times under different hardware configurations and model sizes, providing a solution that combines scalability and resource utilization for efficient MoE inference on edge devices. BRIEF DESCRIPTION OF THE DRAWINGS

[0035] Figure 1 is a flow chart of the method of the present invention;

[0036] Figure 2 This is an example diagram of the offline hybrid parallel mapping strategy of the present invention, where (a) tensor parallelism; (b) expert parallelism; (c) offline linear programming optimization expert allocation; (d) Bayesian optimization physical node mapping;

[0037] Figure 3 Schematic diagram of the online dynamic scheduling mechanism of the present invention, where (a) is online dynamic priority detection; (b) is expert pre-broadcast and communication-friendly token allocation strategy. DETAILED DESCRIPTION

[0038] The present invention will be further described below through specific embodiments in conjunction with the accompanying drawings.

[0039] The present invention aims at a collaborative optimization method of offline hybrid parallel mapping and online dynamic scheduling of hybrid expert model (MoE) under 3D near memory processing (3D NMP) architecture. Figure 1 Flowchart of the method of the present invention; taking the operation of the DeepSeek-V2-Lite-Chat model on a hybrid bonding architecture (Hybrid Bonding DRAM) as an example, the method includes the following steps:

[0040] Step 1: Establish a joint performance analysis model and build a unified analysis framework that includes computing and communication. Given an autoregressive hybrid expert model DeepSeek-V2-Lite-Chat, it contains 26 layers, each layer has 64 experts, and each token activates 6 experts. The hidden layer dimension of the expert is 2048, and the middle layer dimension is 1408. The computing power of each computing unit node of Hybrid BondingDRAM is 1.25TFLPOS, and the interconnection bandwidth of the on-chip network NoC between different computing unit nodes is 75GB / s. There are a total of 16 computing unit nodes. The batch size of inference is assumed to be 16. In addition, f i Denotes the average activation frequency of expert i, with f g P represents the historical average activation frequency of expert group g. ic represents the proportion of the i-th expert deployed on the c-th computing unit node. The computational overhead can be calculated by the formula Quantification reflects the maximum value of node computing load; communication overhead is approximated by linear model It is estimated that the two are combined to form the node-level inference delay objective function. Assume that after simulation, it is known that the scaling factor γ = 2 is determined by linear regression of the linear approximation model and the actual communication simulation time. From the regression analysis, it can be seen that the two have high linearity. Therefore, the node-level inference delay objective function t node_overhead =t comp +4t comm . Figure 2 (a) and Figure 2 (b) shows a partial analysis result of applying the above joint performance analysis model to the tensor parallel strategy and the expert parallel strategy. When the tensor parallel strategy is adopted, Figure 2 As shown in (a), since the experts (E1, E2, E3, E4) are split into all computing unit nodes, a large communication overhead is required; when the expert parallel strategy is adopted, as shown in Figure 2 As shown in (b), due to the imbalance of the load of each expert, the reasoning performance is limited by the computing unit node where the heaviest loaded expert (E3) is located (i.e. Figure 2 Node 3 in (b).

[0041] Step 2: Offline linear programming optimization expert allocation. Solve the expert parameter allocation ratio P through linear programming (LP) ic (0≤P ic ≤1), to minimize t node_overhad As the goal, the constraints include expert parameter allocation completeness (∑ c P ic =1) and node load balancing ( in ). The solution automatically implements a hybrid parallel strategy: high-frequency experts tend to be distributed across multiple nodes due to their heavy computational burden (i.e., they tend to favor tensor parallelism), while low-frequency experts tend to be concentrated on a single node due to their light computational burden (i.e., they tend to favor expert parallelism). In addition, since the objective function takes both computational and communication overheads into account, the expert will not be split across too many computational unit nodes when solving the parallel strategy, in order to achieve a balance between computational load and communication volume. Some of the results obtained are as follows: Figure 2 As shown in (c), since expert 3 (E3) has a heavy load and needs to process more tokens, while expert 4 (E4) has a light load and needs to process fewer tokens, part of the weight of expert 3 (E3) can be placed on node 4 where expert 4 (E4) is located to make full use of the excess computing power of node 4 and prevent node 3 from being overloaded. According to the performance analysis results, the inference performance (latency) of this strategy is significantly better than Figure 2 Tensor parallelism in (a) and Figure 2 The expert parallelism in (b) has a higher expected speedup.

[0042] Step 3: Bayesian optimization of physical node mapping. Based on LP solution P ic , mapping the logical computing unit nodes to the actual physical nodes of the 3D NMP. Here, Bayesian optimization is mainly used to minimize link congestion, optimize NoC communication efficiency, and find the optimal mapping strategy. Some of the results obtained are as follows Figure 2 As shown in (d), the exchange of the positions of nodes 2 and 3 can alleviate Figure 2 Frequent communication between nodes 3 and 4 in (c) reduces link congestion, optimizes NoC communication efficiency, and achieves a higher expected speedup ratio.

[0043] Step 4: Online Inference. Deploy the expert to the Hybrid Bonding DRAM according to the mapping strategy in Step 3, and start online inference on the Hybrid Bonding DRAM. Each autoregressive iteration of inference begins with online inference at the first layer of the hybrid expert model, DeepSeek-V2-Lite-Chat.

[0044] Step 5: Online dynamic priority detection. Offline optimization can only produce static results and cannot handle dynamic and changing reasoning scenarios. The hybrid expert model DeepSeek-V2-Lite-Chat contains 26 layers. During online reasoning, the expert load of the second layer is predicted in real time at the first layer of MoE. Specifically, the expert priority score of the second layer can be calculated. in is the expert activation frequency predicted based on the intermediate activation of layer 1. Some of the test results are as follows Figure 3 As shown in (a), during online reasoning, the expert 1 (E1) in the second layer is heavily loaded, and the node 1 where it is located will have a large reasoning delay, which may seriously affect the reasoning efficiency.

[0045] Step 6: Pre-broadcast of experts with optimal efficiency. After obtaining the priority scores, pre-broadcast as many experts with higher expected load (i.e., higher priority) as possible in the second layer to all storage banks within the time it takes to infer one layer. The optimal pre-broadcast data block size is obtained based on the α-β communication model of NoC. Specifically, assuming that the 3D NMP has a total of 16 computing unit nodes and α = 10 -7 s, The inference time of one layer is sufficient to pre-broadcast 5 experts, so the optimal data block size c≈23.53KB.

[0046] Step 7: Communication-friendly token distribution strategy. After the expert pre-broadcast, each token can be routed to any node holding a copy of its activated expert. To avoid incurring additional communication overhead, routing candidates are limited to nodes that already have the expert being routed. Among these candidate nodes, the token is dispatched to the node with the lowest current computational load. This minimizes load imbalance without incurring additional communication overhead.

[0047] Some of the effects of scheduling are as follows Figure 3 As shown in (b), during Layer 1 reasoning, the expected overload expert 1 (E1) for Layer 2 is pre-broadcasted to all nodes, and the corresponding token is sent to the appropriate node for computation. This approach significantly reduces the latency of Node 1 without incurring additional communication costs.

[0048] Step 8: After receiving the token, the corresponding computing unit node of the first layer starts the online reasoning of the second layer, and then repeats steps 5-7 for the third layer, and so on, until the online dynamic scheduling of the 26th layer of the hybrid expert model is completed, ending one iteration of the reasoning process.

[0049] In actual application, multiple inference iterations may be performed, and each iteration repeats the process from step 5 to step 8 until the hybrid expert model outputs <eos>The token ends or the maximum output length is reached.

[0050] Finally, it should be noted that the purpose of disclosing the embodiments is to facilitate a further understanding of the present invention. However, those skilled in the art will appreciate that various substitutions and modifications are possible without departing from the spirit and scope of the present invention and the appended claims. Therefore, the present invention should not be limited to the contents disclosed in the embodiments; the scope of protection claimed by the present invention shall be determined by the scope defined in the claims.< / eos>

Claims

1. A hybrid parallel and dynamic scheduling method based on a hybrid expert model for 3D near-memory processing, characterized in that: The steps include: Step 1: Establish a joint performance analysis model and build a unified analysis framework that includes computing and communication Given an autoregressive hybrid expert model M, consisting of L layers, each layer has E experts, and each word token activates e experts; the hidden layer dimension of the expert is h, and the intermediate layer dimension is IS; the computing power of each computing unit node is comp, the interconnection bandwidth of the network on chip (NoC) between different computing unit nodes is BW, and there are D computing unit nodes in total; the batch size of inference is B; use f i Denotes the historical average activation frequency of expert i, with f g P represents the historical average activation frequency of expert group g. ic represents the proportion of the i-th expert deployed on the c-th computing unit node; The cost is calculated by the formula Quantification, reflecting the maximum value of node computing load; Communication overhead is modeled by linear approximation Estimate, the two are combined to form the node-level inference delay objective function t node_overhead =t comp +2γt comm , where γ is the scaling factor determined by linear regression between the linear approximation model and the actual communication simulation time; Step 2: Offline linear programming optimization expert allocation Solve the expert parameter allocation ratio P through linear programming ic (0≤P ic ≤1), to minimize the node-level inference latency objective function t node_overhead As the goal, the solution automatically implements a hybrid parallel strategy; Step 3: Bayesian optimization of physical node mapping P based on linear programming solution ic ,Using Bayesian optimization, we optimize the NoC communication efficiency by minimizing link congestion, and find the optimal mapping strategy for mapping logical computing unit nodes to actual physical nodes for 3D near memory processing; Step 4: Online Inference Deploy the expert to the 3D near-memory processing according to the mapping strategy in step 3, and start online inference of the 3D near-memory processing. Each autoregressive iteration of inference starts from the first layer of the hybrid expert model. Step 5: Online dynamic priority detection The hybrid expert model M contains L layers. During online inference, the expert load of the l+1 layer is predicted in real time at the lth layer, and the expert priority score of the l+1 layer is calculated. in is the expert activation frequency predicted based on the intermediate activation of layer l; Step 6: Expert pre-broadcast for optimal efficiency After obtaining the priority scores, pre-broadcast as many experts with heavier expected loads in layer l+1 as possible to all storage banks within the time it takes to infer one layer; The optimal size of pre-broadcast data block is obtained based on the α-β communication model of the NoC. Step 7: Communication-friendly word token distribution strategy After the expert pre-broadcast, each token can be routed to any node in the l+1 layer that holds a copy of its activated expert. Nodes that already have the routed expert are selected as routing candidates. Among these candidate nodes, the token is dispatched to the node with the lowest current computational load. Step 8: After receiving the token, the corresponding computing unit node of layer l starts the online inference of layer l+1, and repeats steps 5-7 for layer l+2, repeating the cycle until the online dynamic scheduling of layer L of the hybrid expert model is completed, ending one iteration of the inference process.

2. The method according to claim 1, wherein The constraints of the linear programming in step 2 include: (1) Expert parameter allocation completeness: ∑ c P ic =1; (2) Node load balancing: in 3. The method according to claim 1, wherein The hybrid parallel strategy implemented in step 2 is: High-frequency experts tend to be distributed across multiple nodes, which is a tendency toward tensor parallelism; whereas low-frequency experts tend to be concentrated on a single node, which is a tendency toward expert parallelism.

4. The method according to claim 1, wherein The optimal size c of the pre-broadcast data block in step 6 satisfies: Among them, α and β are the communication model parameters of the on-chip network, and k is the number of experts that can be pre-broadcasted within the inference time of one layer.

5. The method according to claim 1, wherein The hybrid expert model adopts the DeepSeek-V2-Lite-Chat model.

6. The method according to claim 1, wherein The 3D near memory process adopts a hybrid bonding architecture HybridBonding DRAM.

Images