Tensor segmentation and mapping method of neural network operator on wafer chip

Abstract

The invention provides a tensor segmentation and mapping method of a neural network operator on a wafer chip, and the method comprises the steps: expanding a forward calculation graph into a training calculation graph containing forward propagation, back propagation and gradient updating, and employing different operators in the calculation of these stages; carrying out dimension segmentation on the input tensor of each operator in the training calculation graph to generate a plurality of global candidate segmentation strategies; aiming at each candidate segmentation strategy, enumerating feasible resource mapping strategies of all operators, and combining to generate candidate design points; and screening an optimal design point through cost evaluation, and outputting a corresponding candidate segmentation strategy and a resource mapping strategy. According to the method, forward and backward operators in a neural network training process are decoupled through expansion of a calculation graph, and a segmentation scheme for avoiding tensor copying in forward and backward calculation in the training process and a resource mapping strategy corresponding to the segmentation scheme are explored, so that tensor copying is avoided, global resource occupation is reduced, and the training efficiency is improved. Therefore, the same wafer-level chip can support larger model training.

Description

Technical Field

[0001] The present application relates to the field of neural network training technology, and in particular to a method for tensor segmentation and mapping of a neural network operator on a wafer chip. Background Art

[0002] In recent years, with the widespread application of deep neural networks, especially large-scale pre-trained models (such as the Transformer series) in fields such as natural language processing, computer vision, and multimodal understanding, the scale of neural network models has grown exponentially. These models often contain billions or even tens of billions of parameters. While achieving significant performance improvements, they also pose unprecedented challenges to computing resources. Especially during the training phase, the graphics memory space, data transmission bandwidth, and computing power required by large models far exceed the capacity of a single device.

[0003] To address this issue, distributed parallel computing has become a mainstream solution. By dividing and distributing neural network parameters and computational tasks across multiple physical devices or compute nodes, distributed computing can significantly alleviate resource bottlenecks and improve training efficiency. However, in actual deployment, how to properly partition the tensors in the neural network and map them to hardware resources to minimize cross-device communication overhead and maximize hardware resource utilization is a key issue in distributed system design.

[0004] Related technologies maintain the same operator splitting scheme at each stage of training. However, the splitting dimensions at specific stages can lead to tensor duplication, which in turn leads to additional storage resources on wafer-scale chips, limiting the ability of wafer-scale chips to train larger models. Summary of the Invention

[0005] In view of this, the present application provides a tensor segmentation and mapping method for a neural network operator on a wafer chip to solve the above technical problems.

[0006] In a first aspect of the present application, a method for tensor segmentation and mapping of a neural network operator on a wafer chip is provided, the method comprising:

[0007] Expanding the input forward propagation computation graph of the target neural network to obtain a training computation graph including forward propagation, backpropagation, and gradient update of the target neural network, wherein the forward propagation, backpropagation, and gradient update calculations each use a different operator;

[0008] Performing dimension splitting on the input tensor of each operator in the training computation graph to generate a plurality of candidate splitting strategies, wherein each candidate splitting strategy includes a tensor splitting method for all operators in the training computation graph;

[0009] For each candidate splitting strategy, based on the tensor splitting method of each operator, all feasible resource mapping strategies of the operator are enumerated, and different resource mapping strategies of all operators are combined to generate multiple candidate design points;

[0010] Cost evaluation is performed on all candidate design points, and the optimal design point is screened out based on the cost evaluation result, and the resource mapping strategy and candidate segmentation strategy corresponding to the optimal design point are output as the optimal solution.

[0011] According to one embodiment of the present application, the forward propagation computation graph of the input target neural network is expanded to obtain a training computation graph including forward propagation, back propagation, and gradient update of the target neural network, including:

[0012] Decoupling each forward computation operator in the forward propagation computation graph into independent forward and reverse computation operators;

[0013] The gradient propagation path is explicitly constructed, and the corresponding loss calculation operator and gradient update operator are added to obtain a training computation graph including forward propagation, back propagation and gradient update of the target neural network.

[0014] According to one embodiment of the present application, the input tensor of each operator in the training computation graph is dimensionally split to generate multiple candidate splitting strategies, including:

[0015] Performing dimensionality analysis on the input tensor of each operator in the training computation graph to determine the input tensor dimension of the operator;

[0016] Based on the sub-computation granularity, the input tensor dimension of the operator is split to generate multiple tensor splitting methods of the operator;

[0017] The tensor splitting methods of the operators in the training computation graph are combined, and all feasible candidate splitting strategies are enumerated, wherein each candidate splitting strategy includes the tensor splitting methods of all operators in the training computation graph.

[0018] According to one embodiment of the present application, the input tensor of the operator includes a first input tensor and a second input tensor, and the input tensor dimensions of the operator include a row dimension M of the first input tensor, an intermediate dimension N shared by the first input tensor and the second input tensor, and a column dimension K of the second input tensor;

[0019] The input tensor dimension of the operator is split, including: splitting the input tensor of the operator in M, N and / or K dimensions.

[0020] According to one embodiment of the present application, different computing cores and storage cores on the wafer chip are identified by unique numbers, and the tensor segmentation method based on each operator is used to enumerate all feasible resource mapping strategies for the operator, including:

[0021] For each operator in the candidate splitting strategy, based on the tensor splitting mode of the operator, determine the subtasks that the operator needs to execute in parallel, and enumerate all feasible mapping combinations of computing core particles and storage core particles for each subtask;

[0022] The mapping combinations of different computing core particles and storage core particles of all subtasks are combined to obtain all feasible resource mapping strategies of the operator.

[0023] According to one embodiment of the present application, performing cost evaluation on all candidate design points includes:

[0024] For each operator in each candidate design point, evaluating the communication, computation, and storage costs incurred when the operator is mapped to corresponding hardware resources;

[0025] The communication, computation, and storage costs of all operators in the candidate design point are comprehensively considered to obtain a cost evaluation result of the candidate design point.

[0026] In a second aspect of the present application, a tensor splitting and mapping device for a neural network operator on a wafer chip is provided, the device comprising:

[0027] An expansion unit is used to expand the forward propagation calculation graph of the input target neural network to obtain a training calculation graph including the forward propagation, backpropagation and gradient update of the target neural network, wherein the forward propagation, backpropagation and gradient update calculations respectively use different operators;

[0028] A splitting unit, configured to perform dimension splitting on the input tensor of each operator in the training computation graph and generate a plurality of candidate splitting strategies, wherein each candidate splitting strategy includes a tensor splitting method for all operators in the training computation graph;

[0029] A mapping unit is configured to enumerate all feasible resource mapping strategies for each candidate slicing strategy based on the tensor slicing mode of each operator, and combine different resource mapping strategies of all operators to generate multiple candidate design points;

[0030] The evaluation unit is used to perform cost evaluation on all candidate design points, screen out the optimal design point based on the cost evaluation result, and output the resource mapping strategy and candidate segmentation strategy corresponding to the optimal design point as the optimal solution.

[0031] In a third aspect of the present application, an electronic device is provided, comprising a processor and a memory, wherein the memory stores machine-executable instructions that can be executed by the processor, and the processor is used to execute the machine-executable instructions to implement the steps of the method proposed in the above embodiment.

[0032] In a fourth aspect of the present application, a machine-readable storage medium is provided, wherein the machine-readable storage medium stores machine-executable instructions, and when the machine-executable instructions are executed by a processor, the steps of the method proposed in the above embodiment are implemented.

[0033] In a fifth aspect of the present application, a computer program product is provided, comprising a computer program / instruction, which implements the steps of the method proposed in the above embodiment when executed by a processor.

[0034] It can be seen from the above technical solution that by expanding the forward propagation calculation graph of the input target neural network, a training calculation graph including the forward propagation, back propagation and gradient update of the target neural network is obtained, wherein different operators are used for the calculation of forward propagation, back propagation and gradient update respectively; the input tensor of each operator in the training calculation graph is dimensionally split to generate multiple candidate splitting strategies, wherein each candidate splitting strategy includes the tensor splitting method of all operators in the training calculation graph; for each candidate splitting strategy, based on the tensor splitting method of each operator, all feasible resource mapping strategies of the operator are enumerated, and the different resource mapping strategies of all operators are combined to generate multiple candidate design points; all candidate design points are cost evaluated, and the optimal design point is screened out based on the cost evaluation results, and the resource mapping strategy corresponding to the optimal design point and the candidate splitting strategy are output as the optimal solution. This application decouples the forward and backward operators in the neural network training process through the expansion of the computational graph, and independently explores operator splitting. That is, by decoupling the forward and backward operators, a tensor splitting scheme is explored that avoids tensor copying in both forward and backward calculations during training, thereby avoiding tensor copying and reducing global hardware resource usage, so that the same wafer-level chip can support larger model training.

[0035] It should be understood that the foregoing general description and the following detailed description are exemplary and explanatory only and are not restrictive of the present application. BRIEF DESCRIPTION OF THE DRAWINGS

[0036] Figure 1 This is a flow chart of a tensor segmentation and mapping method for a neural network operator on a wafer chip provided in an embodiment of the present application;

[0037] Figure 2 This is a schematic diagram of the expansion of the calculation graph provided in the embodiment of the present application;

[0038] Figure 3This is an operator-centered tensor segmentation coding table provided in an embodiment of the present application;

[0039] Figure 4 is a schematic diagram of candidate design points provided in an embodiment of the present application;

[0040] Figure 5 This is a flow chart of a tensor segmentation and mapping method for a neural network operator on a wafer chip provided by another embodiment of the present application;

[0041] Figure 6 This is a flow chart of a tensor segmentation and mapping method for a neural network operator on a wafer chip provided by another embodiment of the present application;

[0042] Figure 7 1 is a schematic structural diagram of a tensor splitting and mapping device for a neural network operator on a wafer chip provided in an embodiment of the present application;

[0043] Figure 8 It is a schematic diagram of the hardware structure of an electronic device shown in an exemplary embodiment of the present application. DETAILED DESCRIPTION

[0044] Exemplary embodiments will be described in detail herein, with examples illustrated in the accompanying drawings. In the following description, when referring to the drawings, identical numerals in different figures represent identical or similar elements, unless otherwise indicated. The embodiments described in the following exemplary embodiments are not intended to represent all embodiments consistent with the present application. Rather, they are merely examples of apparatus and methods consistent with certain aspects of the present application, as detailed in the appended claims.

[0045] The terms used in this application are for the purpose of describing particular embodiments only and are not intended to limit this application. As used in this application and the appended claims, the singular forms "a," "an," "the," and "the" are intended to include the plural forms as well, unless the context clearly indicates otherwise.

[0046] In order to enable those skilled in the art to better understand the technical solutions provided by the embodiments of the present application, and to make the above-mentioned purposes, features and advantages of the embodiments of the present application more obvious and easy to understand, the technical solutions in the embodiments of the present application are further described in detail below with reference to the accompanying drawings.

[0047] In recent years, with the widespread application of deep neural networks, especially large-scale pre-trained models (such as the Transformer series) in fields such as natural language processing, computer vision, and multimodal understanding, the scale of neural network models has grown exponentially. These models often contain billions or even tens of billions of parameters. While achieving significant performance improvements, they also pose unprecedented challenges to computing resources. Especially during the training phase, the graphics memory space, data transmission bandwidth, and computing power required by large models far exceed the capacity of a single device.

[0048] To address this issue, distributed parallel computing has become a mainstream solution. By dividing and distributing neural network parameters and computational tasks across multiple physical devices or compute nodes, distributed computing can significantly alleviate resource bottlenecks and improve training efficiency. However, in actual deployment, how to properly partition the tensors in the neural network and map them to hardware resources to minimize cross-device communication overhead and maximize hardware resource utilization is a key issue in distributed system design.

[0049] Related technologies maintain the same operator splitting scheme at each stage of training. However, the splitting dimensions at specific stages can lead to tensor duplication, which in turn leads to additional storage resources on wafer-scale chips, limiting the ability of wafer-scale chips to train larger models.

[0050] In view of this, an embodiment of the present application discloses a tensor splitting and mapping method for a neural network operator on a wafer chip to solve the above technical problems.

[0051] like Figure 1 As shown, Figure 1 This is a flow chart of a tensor splitting and mapping method for a neural network operator on a wafer chip provided in an embodiment of the present application.

[0052] The tensor segmentation and mapping method of the neural network operator on the wafer chip may include the following steps:

[0053] S101: Expand the forward propagation calculation graph of the input target neural network to obtain a training calculation graph including forward propagation, back propagation and gradient update of the target neural network, wherein the forward propagation, back propagation and gradient update calculations respectively use different operators.

[0054] The forward propagation computational graph of the input target neural network is the original computational graph that only contains the forward propagation computational process. The corresponding original computational graph can be constructed according to the structure of the target neural network model (such as Transformer, CNN, etc.). The graph can use operators as nodes and data dependencies as edges to express the computational flow of the target neural network model during the forward propagation process.

[0055] The original computational graph that only contains the forward propagation calculation process is expanded to a complete training computational graph that includes the forward propagation, backpropagation and gradient update of the target neural network.

[0056] It should be noted that in the expanded training computational graph, different operators are used for forward propagation, backpropagation, and gradient update calculations. Backpropagation refers to the process of calculating the gradient information of the weight parameters of each layer by passing the error signal back from the output layer to the input layer layer by layer using the chain rule based on the loss function calculated by the output layer during neural network training, and then using it to update the model parameters to minimize the loss function. Gradient update refers to the process of numerically updating the weight parameters of the current layer using a certain form of optimization rules (such as stochastic gradient descent and momentum optimization) based on the gradient information calculated by the backpropagation algorithm during neural network training.

[0057] In some embodiments, each forward computation operator in the input forward propagation computation graph may be decoupled into independent forward and backward computation operators;

[0058] Explicitly construct the gradient propagation path and add the corresponding loss calculation operator and gradient update operator to obtain the training computation graph that includes the forward propagation, backpropagation and gradient update of the target neural network.

[0059] Specifically, if Figure 2 As shown, Figure 2 This is a schematic diagram of the expansion of the calculation graph provided in the embodiment of this application. Figure 2 In , nodes represent operators and edges represent tensors.

[0060] for Figure 2 The original computational graph (a) in Figure 1 only includes the forward propagation process, which involves two operators: operator 1 (denoted as Op1) and operator 2 (denoted as Op2). Operator 1's input tensors include X and W1, and its output tensor includes Op1_o; operator 2's input tensors include Op1_o and W2, and its output tensor includes Op2_o.

[0061] Based on the characteristics of neural network training, the original computation graph (a) is expanded into the training computation graph (b). Exemplarily, the expansion process includes, but is not limited to, adding necessary control dependency edges, explicitly modeling gradient flow paths, and data reuse channels to support subsequent segmentation and scheduling analysis. Furthermore, the expansion process can be customized based on the characteristics of different training stages.

[0062] for Figure 2In the training computation graph (b), each forward computation operator in the original computation graph (a) is decoupled into independent forward and backward computation operators in the training computation graph (b). For example, operator 1 in the original computation graph (a) is decomposed into two stages: forward propagation and backward propagation, corresponding to operators 1 and 6 in the training computation graph (b), respectively; operator 2 in the original computation graph (a) is also decomposed into two stages: forward propagation and backward propagation, corresponding to operators 2 and 4 in the training computation graph (b), respectively. By explicitly constructing the gradient propagation path, adding the corresponding gradient update operators, namely operators 5 and 7, and deriving the corresponding loss calculation operator, namely operator 3, a complete training computation graph including forward propagation, backward propagation, and gradient update of the target neural network is obtained.

[0063] It should be noted that in Figure 2 In the computation graph expansion diagram shown, the two forward operators in the original computation graph (a) are only used for illustration. The actual application scope is not limited to a specific number of operators and can be adapted to any number of operators.

[0064] In an embodiment of the present application, by decoupling the forward operator in the original computation graph into forward and backward operators, and using the forward and backward operators as independent tensor splitting units, it is possible to explore a tensor splitting scheme that is more suitable for wafer-level chips, that is, a global tensor splitting scheme that avoids tensor copying.

[0065] S102: Perform dimension splitting on the input tensor of each operator in the training computation graph to generate a plurality of candidate splitting strategies, wherein each candidate splitting strategy includes a tensor splitting method for all operators in the training computation graph.

[0066] Taking each operator in the expanded training computation graph as the core, its input tensor is dimensionally split, that is, the large tensor is divided into multiple small blocks, and multiple candidate splitting strategies are generated. Each candidate splitting strategy contains the tensor splitting method of all operators in the training computation graph. In an embodiment of the present application, the original computation graph is expanded in the manner described in S101, and the forward calculation and reverse calculation of each operator are decoupled. Therefore, for the same operator, its forward and backward propagation can adopt different tensor splitting schemes, giving these different operators different splitting space explorations to achieve better performance.

[0067] In some embodiments, the input tensor of each operator in the training computation graph is dimensionally split to generate multiple candidate splitting strategies, including:

[0068] S1021: Performing dimensional analysis on the input tensor of each operator in the training computation graph to determine the input tensor dimension of the operator;

[0069] In the Transformer architecture, linear matrix multiplication is one of its most core computational operations and is widely present in sub-modules such as the attention mechanism and feedforward network.

[0070] In some embodiments, an operator in a training computation graph can be used to perform a matrix multiplication operation, and the input tensor of the operator includes a first input tensor and a second input tensor. The input tensor of each operator in the training computation graph can be dimensionally analyzed to determine the row dimension M of the first input tensor of the operator, the intermediate dimension N shared by the first input tensor and the second input tensor, and the column dimension K of the second input tensor. Exemplarily, the first / second input tensors include but are not limited to input activations, weights, output feature maps, etc.

[0071] For example, Figure 2 Taking operator 1 in the training calculation diagram (b) as an example, its forward calculation process can be formalized as shown in formula (1).

[0072] OP1_o [M,K] =X [M,N] ×W1 [N,K] (Formula 1)

[0073] In formula (1), X represents the input activation tensor, which is an M×N two-dimensional tensor; W1 represents the weight matrix tensor, which is an N×K two-dimensional tensor; OP1_o represents the output tensor, which is an M×K two-dimensional tensor. The input tensor of each operator in the training computation graph is dimensionally analyzed to determine the row dimension M of the input activation tensor of the operator, the intermediate dimension N shared by the input data tensor and the weight matrix tensor, and the column dimension K of the weight matrix tensor.

[0074] S1022: Splitting the input tensor dimension of the operator based on the sub-computation granularity to generate multiple tensor splitting methods of the operator;

[0075] Subcomputation granularity refers to the number of independent computational sub-blocks into which an operator's input tensor is split along a specific dimension. For example, if the subcomputation granularity is N, the operator's input tensor can be split along one or more dimensions, each into N parts.

[0076] In some embodiments, the input tensor of an operator may be split in M, N, and / or K dimensions.

[0077] In some embodiments, the various tensor splitting methods for the generated operators can be represented in the form of split codes. For example, for each operator, CUT(m,n,k) records the splitting status of the operator in the three dimensions M, N, and K. For example, CUT(2,1,1) indicates that the input tensor of the operator is split in two along the M dimension, and no split is performed in the N and K dimensions.

[0078] For each operator, multiple tensor splitting methods can be generated. Assuming the sub-computation granularity is 2, the generated splitting codes for the operator can include CUT(2,1,1), CUT(1,2,1), CUT(1,1,2), and so on.

[0079] S1023: Combine the tensor splitting methods of the operators in the training computation graph, and enumerate all feasible candidate splitting strategies, wherein each candidate splitting strategy includes the tensor splitting methods of all operators in the training computation graph.

[0080] Enumerates all possible tensor splitting combinations for operators in the training computation graph.

[0081] Assume that the training computation graph contains two operators, namely operator 1 and operator 2. According to the method described in S1022, the tensor splitting methods generated for operator 1 and operator 2 are (method 1, method 2) and (method 3, method 4) respectively. By combining the tensor splitting methods of each operator in the training computation graph, we can obtain: {(operator 1: method 1, operator 2: method 3), (operator 1: method 1, operator 2: method 4), (operator 1: method 2, operator 2: method 3), (operator 1: method 2, operator 2: method 4)}, that is, a total of 2+2=4 feasible candidate splitting strategies can be enumerated.

[0082] In some embodiments, each candidate segmentation strategy can be represented by a tensor segmentation coding table. Figure 3 As shown, Figure 3 This is an operator-centered tensor splitting coding table provided in an embodiment of the present application. With the operator as the basic unit, CUT(m,n,k) is used to record the splitting status of each operator in the training calculation graph for the three dimensions M, N, and K. For example, for operator 1, CUT(1,1,2) means that the input tensor of the operator is divided into two parts in the K dimension, and no division is performed in the M and N dimensions. Similarly, by maintaining a tensor splitting coding table containing all operators in the training calculation graph, all design spaces after flexible operator splitting can be obtained.

[0083] S103: For each candidate segmentation strategy, based on the tensor segmentation mode of each operator, enumerate all feasible resource mapping strategies of the operator, and combine different resource mapping strategies of all operators to generate multiple candidate design points.

[0084] After enumerating all feasible candidate splitting strategies as described in S102, all possible resource mapping strategies are enumerated for the tensor splitting method of each operator in each candidate splitting strategy. A resource mapping strategy refers to the execution plan for allocating the split sub-computation blocks to specific hardware resources. The resource mapping strategies of all operators in the candidate splitting strategy are then combined to generate multiple candidate design points corresponding to the candidate splitting strategy.

[0085] For example, a wafer-level chip (wafer chip) refers to an integrated circuit system that uses an entire wafer as a single logic chip construction platform. The wafer-level chip completes the interconnection and system construction of multiple chips at the wafer level. A wafer-level chip usually integrates dozens of computing chips and has extremely high computing density, bandwidth utilization and system integration.

[0086] Wafer-level chips contain computing cores and storage cores. Computing cores are basic hardware units integrated on wafer-level chips for computing, while storage cores are basic hardware units integrated on wafer-level chips for data storage and caching.

[0087] Different computing cores and storage cores on the wafer-level chip are identified by unique numbers. For example, computing cores include computing core 0, computing core 1, computing core 2, etc.; storage cores include storage core 0, storage core 1, storage core 2, etc.

[0088] In some embodiments, based on the tensor partitioning method of each operator, all feasible resource mapping strategies for the operator are enumerated, including:

[0089] S1031: For each operator in the candidate splitting strategy, based on the tensor splitting mode of the operator, determine the subtasks that the operator needs to execute in parallel, and enumerate all feasible mapping combinations of computing core particles and storage core particles for each subtask;

[0090] For example, suppose the training computation graph contains two operators, namely operator 1 and operator 2, and the candidate splitting strategies are: operator 1, CUT (1,1,2); operator 2, CUT (1,2,1). For operator 1 in the candidate splitting strategy, the input matrix of operator 1 includes the input activation matrix X and the weight matrix tensor W, where the weight matrix tensor W is split into two parts in the K dimension. Then, all feasible mapping combinations of computation core particles and storage core particles are enumerated for sub-computation blocks 1 and sub-computation blocks 2 that need to be executed in parallel. For example, the possible mapping combinations of computation core particles and storage core particles for sub-computation block 1 are {X op1 : Calculate core particle 0 and store core particle 0; Calculate core 0, store core 0}, {X op1 : Calculate core particle 1, store core particle 1; Computation core particle 1, storage core particle 1} and so on; the possible mapping combinations of computing core particles and storage core particles of sub-computing block 2 are {X op1 : Calculate core particle 1, store core particle 1; Calculate core particle 1, store core particle 1}, {X op1 : Calculate core particle 2 and store core particle 2; Compute core 2, store core 2} and so on.

[0091] For operator 2 in the candidate splitting strategy, the input activation matrix X and weight matrix tensor W are both split into two parts in N dimensions. Then, all feasible mapping combinations of computing core particles and storage core particles are enumerated for sub-computation blocks 3 and 4 that need to be executed in parallel. For example, a possible mapping combination of computing core particles and storage core particles for sub-computation block 3 is { Calculate core particle 2 and store core particle 2; Computation core particle 2, storage core particle 2}; Sub-computation block 4 A possible mapping combination of computing core particles and storage core particles is { Calculate core particle 3 and store core particle 3; Compute core particle 3, store core particle 3}.

[0092] S1032: Combine the mapping combinations of different computing core particles and storage core particles of all subtasks to obtain all feasible resource mapping strategies of the operator.

[0093] Combine the mapping combinations of different computing core particles and storage core particles of all subtasks of an operator to obtain all feasible resource mapping strategies of the operator. For example, for operator 1 in the previous example, all feasible resource mapping strategies of the operator can be {X op1 : Calculate core 0 / 1 and store core 0 / 1; Calculate core 0 and store core 0; Calculate core particle 1, store core particle 1}, {X op1 :Calculate core particle 1 / 2, store core particle 1 / 2; Calculate core particle 1 and store core particle 1; Compute core 2, store core 2} and so on.

[0094] The resource mapping strategies of all operators in the candidate segmentation strategy are combined to generate multiple candidate design points corresponding to the candidate segmentation strategy. That is, each candidate segmentation strategy will derive multiple resource mapping possibilities, and the combination of these possibilities constitutes a candidate design point. Figure 4 As shown, Figure 4 Schematic diagram of candidate design points provided in an embodiment of the present application, which include resource mapping strategies for all operators in the training computation graph.

[0095] S104: Cost evaluation is performed on all candidate design points, and an optimal design point is screened out based on the cost evaluation result, and the resource mapping strategy and candidate segmentation strategy corresponding to the optimal design point are output as the optimal solution.

[0096] Each candidate partitioning strategy can generate multiple candidate design points, representing different hardware resource mapping solutions for each candidate partitioning strategy. A cost evaluation is performed on each candidate design point for each candidate partitioning strategy. Based on the cost evaluation results, methods such as enumeration, heuristic search, greedy strategies, or reinforcement learning are used to select the candidate design point with the best performance / cost trade-off from the design space as the optimal design point. The resource mapping strategy and candidate partitioning strategy corresponding to the optimal design point are then output as the optimal solution.

[0097] In some embodiments, a cost evaluation is performed on all candidate design points, including:

[0098] S1041: For each operator in each candidate design point, evaluate the communication, computation, and storage costs incurred when the operator is mapped to corresponding hardware resources;

[0099] For each candidate design point, a comprehensive evaluation of the computation, communication, and storage costs of each operator in the candidate design point is performed, specifically including indicators such as cross-device bandwidth usage, on-chip tensor reuse efficiency, and load balancing, so as to construct a complete multi-dimensional performance cost model.

[0100] S1042: Integrate the communication, computation, and storage costs of all operators in the candidate design point to obtain a cost evaluation result of the candidate design point.

[0101] For each candidate design point, the communication, computation, and storage costs of all operators in the candidate design point are integrated to obtain the cost evaluation result of the candidate design point.

[0102] In an embodiment of the present application, by expanding the forward propagation calculation graph of the input target neural network, a training calculation graph including forward propagation, back propagation and gradient update of the target neural network is obtained, wherein different operators are used for the calculation of forward propagation, back propagation and gradient update respectively; the input tensor of each operator in the training calculation graph is dimensionally split to generate multiple candidate splitting strategies, wherein each candidate splitting strategy includes the tensor splitting method of all operators in the training calculation graph; for each candidate splitting strategy, based on the tensor splitting method of each operator, all feasible resource mapping strategies of the operator are enumerated, and the different resource mapping strategies of all operators are combined to generate multiple candidate design points; cost evaluation is performed on all candidate design points, and the optimal design point is screened out based on the cost evaluation result, and the resource mapping strategy and candidate splitting strategy corresponding to the optimal design point are output as the optimal solution. This application decouples the forward and backward operators in the neural network training process through the expansion of the computational graph, and independently explores operator splitting. That is, by decoupling the forward and backward operators, a tensor splitting scheme is explored that avoids tensor copying in both forward and backward calculations during training, thereby avoiding tensor copying and reducing global hardware resource usage, so that the same wafer-level chip can support larger model training.

[0103] like Figure 5 As shown, Figure 5 This is a flow chart of a tensor splitting and mapping method for a neural network operator on a wafer chip provided in another embodiment of the present application, wherein different computing cores and storage cores on the wafer-level chip are identified by unique numbers.

[0104] Input the forward propagation computational graph of the target neural network, that is, input the original computational graph that only contains the forward propagation computation process.

[0105] The original computation graph is expanded into a training computation graph, including adding necessary control dependency edges, explicitly modeling gradient flow paths, and data reuse channels to support subsequent segmentation and scheduling analysis. This expansion process allows for customized modeling based on the characteristics of different training phases. In the expanded training computation graph, different operators are used for forward propagation, backpropagation, and gradient updates.

[0106] Taking each operator in the training graph as the core, we perform dimension splitting on its input tensor, dividing the large tensor into multiple smaller pieces, and generate multiple split encoding strategies for that operator. We then combine the split encoding strategies of all operators in the training graph and enumerate all feasible candidate split strategies. Each candidate split strategy can be represented by a tensor split encoding table centered around the operator. Each tensor split encoding table contains the tensor splitting methods for all operators in the training graph.

[0107] For each operator in each candidate partitioning strategy, the subtasks that need to be executed in parallel are determined based on the operator's tensor partitioning method. For each subtask, all feasible mapping combinations of compute coreparticles and storage coreparticles are enumerated. The mapping combinations of compute coreparticles and storage coreparticles for all subtasks are combined to obtain all feasible resource mapping strategies for the operator. The resource mapping strategies of all operators in the candidate partitioning strategy are combined to generate multiple candidate design points for the candidate partitioning strategy.

[0108] For each operator in each candidate design point, the communication, computation, and storage costs incurred by the operator when mapped to the corresponding hardware resources are evaluated; the communication, computation, and storage costs of all operators in the candidate design point are combined to obtain the cost evaluation result of the candidate design point.

[0109] Based on the cost evaluation results, methods such as enumeration, heuristic search, greedy strategies, or reinforcement learning can be used to select candidate design points from the design space that offer the best performance / cost trade-off as the optimal design point. Furthermore, if pre-defined constraints are not met, the design point enumeration step can be returned to generate new candidate design points, forming a closed optimization loop.

[0110] The resource mapping strategy and candidate segmentation strategy corresponding to the optimal design point are output as the optimal solution.

[0111] In this embodiment, by expanding the computational graph, the forward and backward operators in the neural network training process are decoupled and independent operator partitioning exploration is performed. Specifically, by decoupling the forward and backward operators, a tensor partitioning scheme is explored that avoids tensor copying in both forward and backward computations during training. This flexible tensor partitioning exploration provides a larger design space, enabling a global tensor partitioning scheme that reduces tensor copying at the expense of increased communication traffic. The high-bandwidth architecture of wafer-scale chips is well-suited to this increased communication-heavy tensor partitioning and resource mapping scheme.

[0112] like Figure 6 As shown, Figure 6 1 is a flow chart of a tensor segmentation and mapping method for a neural network operator on a wafer chip, provided by another embodiment of the present application. In particular, different computing cores and storage cores on the wafer-level chip are identified by unique numbers.

[0113] Step 1: Input the original computation graph

[0114] The input consists only of the original computational graph for the forward propagation process. Nodes in the original computational graph represent operators, such as operator 1 (denoted as Op1) and operator 2 (denoted as Op2). Edges represent tensors, such as X and W1, which represent the input tensors of operator 1, and Op1_o, which represents the output tensor of operator 1.

[0115] Exemplarily, W1 and W2 represent weight matrix tensors of the target neural network model.

[0116] Step 2: Expanding the computational graph

[0117] During the training phase of the target neural network model, its loss calculation, backpropagation, and gradient update paths are expanded. Each forward computation operator in the original computation graph is decoupled into independent forward and backward computation operators in the training computation graph. For example, operator 1 in the original computation graph is decomposed into two phases: forward propagation and backward propagation, corresponding to operators 1 and 6 in the training computation graph, respectively. The corresponding gradient update operators, namely operators 5 and 7, are added, and the corresponding loss calculation operator, namely operator 3, is derived, thus obtaining a complete training computation graph that includes the forward propagation, backward propagation, and gradient update of the target neural network.

[0118] Step 3: Operator-centered tensor segmentation encoding table

[0119] For each operator in the training computation graph, its input tensor is dimensionally split, dividing the large tensor into multiple smaller pieces. This generates multiple split encoding strategies for that operator. The split encoding strategies for all operators in the training computation graph are combined, and all feasible candidate split strategies are enumerated. Each candidate split strategy is represented by a tensor split encoding table centered around the operator. Each tensor split encoding table contains the tensor splitting methods for all operators in the training computation graph. For example, CUT(1,1,2) corresponding to OP1 indicates that the input tensor of operator 1 is split into 1, 1, and 2 in the (M,N,K) dimensions, respectively. The purpose of splitting is to distribute the data to multiple computation cores for parallel computation.

[0120] Step 4: Tensor Split Space

[0121] For different tensor partitioning encoding tables, this table shows how to spatially partition the operator's input tensor into multiple parts that need to be executed in parallel. For example, Partition 1 shows that CUT(1,1,2) partitions the input tensor in the K dimension into two subtasks that need to be executed in parallel.

[0122] Step 5: Enumerate design points and estimate costs

[0123] For each operator in each candidate splitting strategy, enumerate all feasible mapping combinations of compute coreparticles and storage coreparticles for each subtask of the operator. Combine the different mapping combinations of compute coreparticles and storage coreparticles for all subtasks of the operator to obtain all feasible resource mapping strategies for the operator. Combine the resource mapping strategies of all operators in the candidate splitting strategy to generate multiple candidate design points for the candidate splitting strategy.

[0124] This embodiment implements hardware resource allocation by maintaining a mapping table that maintains the numbers of operators and the compute and storage cores on wafer-level devices, thereby obtaining a mapping relationship between operators and wafer-level cores. By assigning globally unique numbers to different compute and storage cores on the wafer-level chip, the subtasks generated by different operators are assigned to different compute and storage cores for execution.

[0125] For each candidate design point, a comprehensive evaluation of the computation, communication, and storage costs of each operator in the candidate design point is performed. For example, the evaluation of computation, communication, and storage costs includes, but is not limited to, computational core usage (e.g., how many cores are used), storage core requirements, and communication volume for forward and backward propagation.

[0126] The communication, computation, and storage costs of all operators in the candidate design point are integrated to obtain the cost evaluation result of the candidate design point, which is used for subsequent optimal mapping decisions.

[0127] Step 6: On-chip mapping at the grain level (optimal solution exploration)

[0128] According to the mapping table in step 5, all operators in the candidate design points are mapped to the corresponding core particles on the wafer-level chip. Different core particle groups can perform tensor parallelism (TP) and pipeline parallelism (PP). It also includes multi-time step scheduling. This is because all operations in the mapping process may not be completed within a single time step and need to be performed across multiple time steps. For example, Op1, Op2, and Op3 are executed in Time Step 0, and Op4 and Op6 are continued in Time Step 1, as well as Op5 and Op7 (not shown).

[0129] Based on the cost evaluation results, methods such as enumeration, heuristic search, greedy strategy or reinforcement learning can be used to select the candidate design point with the best performance / cost trade-off from the design space as the optimal design point, and the resource mapping strategy and candidate segmentation strategy corresponding to the optimal design point are output as the optimal solution.

[0130] In this embodiment, by expanding the computational graph, the forward and backward operators in the neural network training process are decoupled and independent operator partitioning exploration is performed. Specifically, by decoupling the forward and backward operators, a tensor partitioning scheme is explored that avoids tensor copying in both forward and backward computations during training. This flexible tensor partitioning exploration provides a larger design space, enabling a global tensor partitioning scheme that reduces tensor copying at the expense of increased communication traffic. The high-bandwidth architecture of wafer-scale chips is well-suited to this increased communication-heavy tensor partitioning and resource mapping scheme.

[0131] The above content describes the method provided by this application. The following describes the device provided by this application:

[0132] See Figure 7 , which is a structural diagram of a tensor splitting and mapping device for a neural network operator on a wafer chip provided in an embodiment of the present application.

[0133] like Figure 7 As shown, the device may include:

[0134] An expansion unit 710 is configured to expand the input forward propagation computation graph of the target neural network to obtain a training computation graph including forward propagation, backpropagation, and gradient update of the target neural network, wherein the forward propagation, backpropagation, and gradient update calculations each use a different operator;

[0135] A splitting unit 720 is configured to perform dimension splitting on the input tensor of each operator in the training computation graph and generate a plurality of candidate splitting strategies, wherein each candidate splitting strategy includes a tensor splitting method for all operators in the training computation graph;

[0136] A mapping unit 730 is configured to enumerate all feasible resource mapping strategies for each candidate partitioning strategy based on the tensor partitioning mode of each operator, and combine different resource mapping strategies of all operators to generate multiple candidate design points;

[0137] The evaluation unit 740 is configured to perform cost evaluation on all candidate design points, select an optimal design point based on the cost evaluation result, and output the resource mapping strategy and candidate segmentation strategy corresponding to the optimal design point as the optimal solution.

[0138] Optionally, the expansion unit 710 is specifically configured to:

[0139] Decoupling each forward computation operator in the forward propagation computation graph into independent forward and reverse computation operators;

[0140] The gradient propagation path is explicitly constructed, and the corresponding loss calculation operator and gradient update operator are added to obtain a training computation graph including forward propagation, back propagation and gradient update of the target neural network.

[0141] Optionally, the segmentation unit 720 is specifically configured to:

[0142] Performing dimensionality analysis on the input tensor of each operator in the training computation graph to determine the input tensor dimension of the operator;

[0143] Based on the sub-computation granularity, the input tensor dimension of the operator is split to generate multiple tensor splitting methods of the operator;

[0144] The tensor splitting methods of the operators in the training computation graph are combined, and all feasible candidate splitting strategies are enumerated, wherein each candidate splitting strategy includes the tensor splitting methods of all operators in the training computation graph.

[0145] The input tensor of the operator includes a first input tensor and a second input tensor, and the input tensor dimensions of the operator include a row dimension M of the first input tensor, an intermediate dimension N shared by the first input tensor and the second input tensor, and a column dimension K of the second input tensor;

[0146] The segmentation unit 720 is specifically configured to:

[0147] Split the input tensor of the operator into M, N and / or K dimensions.

[0148] Optionally, different computing cores and storage cores on the wafer chip are all identified by unique numbers;

[0149] The mapping unit 730 is specifically configured to:

[0150] For each operator in the candidate splitting strategy, based on the tensor splitting mode of the operator, determine the subtasks that the operator needs to execute in parallel, and enumerate all feasible mapping combinations of computing core particles and storage core particles for each subtask;

[0151] The mapping combinations of different computing core particles and storage core particles of all subtasks are combined to obtain all feasible resource mapping strategies of the operator.

[0152] Optionally, the evaluation unit 740 is specifically configured to:

[0153] For each operator in each candidate design point, evaluating the communication, computation, and storage costs incurred when the operator is mapped to corresponding hardware resources;

[0154] The communication, computation, and storage costs of all operators in the candidate design point are comprehensively considered to obtain a cost evaluation result of the candidate design point.

[0155] The implementation process of the functions and effects of each unit in the above-mentioned device is specifically described in the implementation process of the corresponding steps in the above-mentioned method, and will not be repeated here.

[0156] The embodiment of the present application also provides a hardware structure. Figure 8 , Figure 8 This is a structural diagram of an electronic device provided in an embodiment of the present application. Figure 8As shown, the hardware structure may include: a processor and a machine-readable storage medium, the machine-readable storage medium storing machine-executable instructions that can be executed by the processor; the processor is used to execute the machine-executable instructions to implement the method disclosed in the above example of this application.

[0157] Based on the same application concept as the above method, an embodiment of the present application also provides a machine-readable storage medium, on which a number of computer instructions are stored. When the computer instructions are executed by a processor, the method disclosed in the above example of the present application can be implemented.

[0158] Exemplarily, the machine-readable storage medium may be any electronic, magnetic, optical, or other physical storage device that may contain or store information, such as executable instructions, data, and the like. For example, the machine-readable storage medium may be: RAM (Random Access Memory), volatile memory, non-volatile memory, flash memory, a storage drive (such as a hard disk drive), a solid-state drive, any type of storage disk (such as a CD, DVD, etc.), or similar storage media, or a combination thereof.

[0159] It should be noted that, in this document, relational terms such as target and objective are used only to distinguish one entity or operation from another entity or operation, and do not necessarily require or imply any actual relationship or order between these entities or operations. Moreover, the terms "comprise", "include" or any other variants thereof are intended to cover non-exclusive inclusion, so that a process, method, article or device comprising a series of elements includes not only those elements, but also other elements not explicitly listed, or elements inherent to such process, method, article or device. In the absence of further limitations, an element defined by the phrase "comprising a ..." does not exclude the presence of other identical elements in the process, method, article or device comprising the element.

[0160] The above description is only a preferred embodiment of the present application and is not intended to limit the present application. Any modifications, equivalent replacements, improvements, etc. made within the spirit and principles of the present application shall be included in the scope of protection of the present application.

Claims

1. A method for tensor segmentation and mapping of a neural network operator on a wafer chip, characterized in that: The method includes: Expanding the input forward propagation computation graph of the target neural network to obtain a training computation graph including forward propagation, backpropagation, and gradient update of the target neural network, wherein the forward propagation, backpropagation, and gradient update calculations each use a different operator; Performing dimension splitting on the input tensor of each operator in the training computation graph to generate a plurality of candidate splitting strategies, wherein each candidate splitting strategy includes a tensor splitting method for all operators in the training computation graph; For each candidate splitting strategy, based on the tensor splitting method of each operator, all feasible resource mapping strategies of the operator are enumerated, and different resource mapping strategies of all operators are combined to generate multiple candidate design points; Cost evaluation is performed on all candidate design points, and the optimal design point is screened out based on the cost evaluation result, and the resource mapping strategy and candidate segmentation strategy corresponding to the optimal design point are output as the optimal solution.

2. The method according to claim 1, characterized in that The forward propagation computation graph of the input target neural network is expanded to obtain a training computation graph including forward propagation, back propagation, and gradient update of the target neural network, including: Decoupling each forward computation operator in the forward propagation computation graph into independent forward and reverse computation operators; The gradient propagation path is explicitly constructed, and the corresponding loss calculation operator and gradient update operator are added to obtain a training computation graph including forward propagation, back propagation and gradient update of the target neural network.

3. The method according to claim 1, characterized in that The input tensor of each operator in the training computation graph is dimensionally split to generate multiple candidate splitting strategies, including: Performing dimensionality analysis on the input tensor of each operator in the training computation graph to determine the input tensor dimension of the operator; Based on the sub-computation granularity, the input tensor dimension of the operator is split to generate multiple tensor splitting methods of the operator; The tensor splitting methods of the operators in the training computation graph are combined, and all feasible candidate splitting strategies are enumerated, wherein each candidate splitting strategy includes the tensor splitting methods of all operators in the training computation graph.

4. The method according to claim 3, characterized in that The input tensor of the operator includes a first input tensor and a second input tensor, and the input tensor dimensions of the operator include a row dimension M of the first input tensor, an intermediate dimension N shared by the first input tensor and the second input tensor, and a column dimension K of the second input tensor; The input tensor dimension of the operator is split, including: splitting the input tensor of the operator in M, N and / or K dimensions.

5. The method according to claim 1, wherein Different computing cores and storage cores on the wafer chip are identified by unique numbers. Based on the tensor splitting method of each operator, all feasible resource mapping strategies of the operator are enumerated, including: For each operator in the candidate splitting strategy, based on the tensor splitting mode of the operator, determine the subtasks that the operator needs to execute in parallel, and enumerate all feasible mapping combinations of computing core particles and storage core particles for each subtask; The mapping combinations of different computing core particles and storage core particles of all subtasks are combined to obtain all feasible resource mapping strategies of the operator.

6. The method according to claim 1, characterized in that The cost evaluation of all candidate design points includes: For each operator in each candidate design point, evaluating the communication, computation, and storage costs incurred when the operator is mapped to corresponding hardware resources; The communication, computation, and storage costs of all operators in the candidate design point are comprehensively considered to obtain a cost evaluation result of the candidate design point.

7. A tensor segmentation and mapping device for a neural network operator on a wafer chip, characterized in that: The device includes: An expansion unit is used to expand the forward propagation calculation graph of the input target neural network to obtain a training calculation graph including the forward propagation, backpropagation and gradient update of the target neural network, wherein the forward propagation, backpropagation and gradient update calculations respectively use different operators; A splitting unit, configured to perform dimension splitting on the input tensor of each operator in the training computation graph and generate a plurality of candidate splitting strategies, wherein each candidate splitting strategy includes a tensor splitting method for all operators in the training computation graph; A mapping unit is configured to enumerate all feasible resource mapping strategies for each candidate slicing strategy based on the tensor slicing mode of each operator, and combine different resource mapping strategies of all operators to generate multiple candidate design points; The evaluation unit is used to perform cost evaluation on all candidate design points, screen out the optimal design point based on the cost evaluation result, and output the resource mapping strategy and candidate segmentation strategy corresponding to the optimal design point as the optimal solution.

8. An electronic device, characterized in that: The system comprises a processor and a memory, wherein the memory stores machine-executable instructions that can be executed by the processor, and the processor is configured to execute the machine-executable instructions to implement the method according to any one of claims 1 to 6.

9. A machine-readable storage medium, characterized in that The machine-readable storage medium stores machine-executable instructions, and when the machine-executable instructions are executed by a processor, the method according to any one of claims 1 to 6 is implemented.

10. A computer program product comprising a computer program / instructions, characterized in that When the computer program / instructions are executed by a processor, the steps of the method according to any one of claims 1 to 6 are implemented.

Images