A method to improve the speed of inference and training of deep operator network models

Abstract

The present invention belongs to the field of artificial intelligence and high-performance computing technology, and particularly relates to a method for improving the speed of reasoning, solving and training of deep operator network models. The present invention performs multi-level optimization on the PI-DeepONet model to improve its reasoning, solving and training speeds. The present invention customizes the SGEMM operator, and within the model size range commonly used by PI-DeepONet, the SGEMM operator can achieve a speedup ratio of up to 1.5 times compared to the best single-precision matrix multiplication in the basic linear algebra library of the CUDA computing platform. The fusion operator is designed for the special structure of the modified MLP used in the model, and the end-to-end reasoning speed of the model obtains a speedup ratio of more than 2 times, and the graphics processor calculation time can be shortened by up to 5.7 times. The present invention uses data parallelism to perform distributed training on the model, and uses a communication masking strategy, and the scalability is maintained at 83%-85%.

Description

Technical Field

[0001] The present invention belongs to the field of artificial intelligence and high performance computing technology, and in particular relates to a method for improving the speed of reasoning, solving and training of deep operator network models. Background Art

[0002] Solving partial differential equations plays a vital role in many scientific research fields such as engineering technology, biomathematics and image processing. For example, in the field of petroleum engineering, by solving fluid mechanics equations such as Burgers' equation and Darcy's flow equation, the velocity and pressure distribution of the fluid in the reservoir can be obtained, and a numerical model of the real oil field can be constructed, which in turn plays a guiding role in decision-making in reserve assessment, production capacity prediction and oil field development plan prediction, helping to improve the efficiency and economic benefits of oil exploration and development.

[0003] Solving partial differential equations has long been a very important and challenging task. On the one hand, the mathematical modeling of many physical systems is constructed through partial differential equations. Solving partial differential equations is an important means of conducting scientific research in these fields and is also the basis of many industrial software and physical calculations. Classical solution methods include finite difference method and finite element method, which usually perform complex meshing on a given domain, divide the continuous problem domain into discrete grids, and use simple functions to approximate the results of the division, and finally obtain approximate solutions to the equations. On the other hand, there are several difficulties in using traditional numerical methods to solve partial differential equations: first, these methods all rely on meshing, and the process of designing a suitable mesh for a specific solution problem is challenging in itself; then, these methods are accompanied by a large number of iterative calculations at each step of the calculation process, which often means a high computational cost; in addition, the process of using numerical methods to solve partial differential equations may vary greatly depending on the problem.

[0004] Recently, deep learning-based methods have become another effective means of solving partial differential equations, and are widely used in many fields such as fluid mechanics, heat conduction equations, and materials science. This type of method can already achieve good accuracy; at the same time, compared with traditional numerical methods, trained deep learning models have huge advantages in solving speed.

[0005] Since the 1990s, there have been several studies devoted to using neural networks to solve partial differential equations. Among them, Dissanayake and Phan Thien transformed the problem of solving partial differential equations into an unconstrained minimization problem, proposed an algorithm for solving partial differential equations using neural networks, and used their algorithm to solve the Poisson equation and the nonlinear heat conduction equation. This work is the beginning of the research direction of using neural networks to solve partial differential equations. However, due to the limitations of the computing power of computers at the time and the neural network models used, these methods can often only be used to solve simple partial differential equations, so they did not receive widespread attention and sufficient attention at the time.

[0006] In 2019, Raissi et al. proposed a groundbreaking work PINN. They used the control equations, initial values ​​and boundary value conditions of partial differential equations in deep neural networks to construct the loss function to constrain the convergence direction of the network model. Compared with simply using labeled data for training, the model trained under this idea is more in line with the conditions of the equation itself. After that, many variants of PINN were proposed. UQPINN combines PINN and adversarial learning strategies to deal with various partial differential equations. fPINN combines PINN and time discretization technology to solve time fractional partial differential equations. PPINN reduces the training cost of the algorithm and handles long-term problems more efficiently by decomposing a long-term problem into multiple independent short-term problems. But no matter what degree of improvement has been made, the core idea is to add the conditions of the partial differential equation itself to the deep neural network model to constrain the training of the model, so they are all deep neural networks based on physical information.

[0007] Another method is to use operator neural networks with unique structures to solve partial differential equations. The object of the network model approximation of this type of method is no longer the solution function of the partial differential equation, but the operator that maps the basis function to the solution function. This type of method often has better generalization. The two most representative works are the deep operator network (DeepONet) proposed by Lu Lu et al. and the Fourier neural operator (FNO) proposed by ZongYi Li et al. Afterwards, improvements and applications based on these two operator neural networks have also been proposed one after another. S-DeepONet introduces a temporal network structure into the branch network of DeepONet to solve problems with temporal input data. Goswami S et al. proposed a method to improve the speed of reasoning, solving and training of deep operator network models. V-DeepONet based on physical principles is used for brittle fracture analysis. GegeWen et al. proposed U-FNO for solving multiphase flow problems. Meer Mehran Rashid et al. used FNO to learn the mechanical response of two-dimensional composite materials.

[0008] Although the method of solving partial differential equations based on operator neural networks has the advantage of strong versatility, it cannot cope with the situation where there is less or even no labeled data due to its purely data-driven approach. In real problems, real data in some fields may be difficult to obtain in large quantities or the cost of obtaining is very high, and there is only a very small amount of labeled data. Therefore, some work has tried to combine the idea of ​​operator neural networks with that of neural networks based on physical information, and add partial differential equation conditions as constraints to the operator neural network model. PI-DeepONet is a method to improve the inference and training speed of deep operator network models. It is an advanced deep learning model for solving partial differential equations. Based on the DeepONet network model, PI-DeepONet uses the control function and initial boundary value conditions of partial differential equations to construct the loss function, which can train the network model to convergence in the absence of labeled data, and achieves good accuracy.

[0009] However, when using PI-DeepONet to solve partial differential equations, the training and inference speed of the model will be limited due to some structural characteristics of the model: the narrower fully connected layer width and larger batch size will lead to inefficient matrix multiplication calculations, and the frequent activation of many element-by-element operators will also bring additional overhead, which will limit the inference speed of the model; on the other hand, due to the use of the PINN structure, a large number of derivative calculations are introduced into the training process, and the time overhead of model training has become difficult to ignore. Summary of the invention

[0010] The present invention proposes a method for improving the speed of reasoning, solving and training of deep operator network models, so as to solve the problems that the deep operator network models take a long time to train, have limited scalability and slow reasoning speed in practical applications.

[0011] The technical solution of the present invention is achieved in this way:

[0012] A method for improving the speed of inference solving and training of a deep operator network model comprises the following steps:

[0013] S1. Customized SGEMM operator for PI-DeepONet neural network model: designed three-level block strategies of thread block, thread bundle and thread to achieve high computation-to-memory ratio, and optimized memory access mode to reduce memory access time overhead;

[0014] S2. Design an operator fusion strategy; fuse a large number of scattered element-by-element calculations in the PI-DeepONet neural network model with matrix multiplication. The time consumed by these element-by-element calculations mainly comes from frequent operator startup and frequent access to global memory for reading or writing large-size tensors. Operator fusion effectively eliminates these additional overheads.

[0015] S3. Design a distributed parallel training strategy; reduce the time required for PI-DeepONet neural network model training by overlapping the time overhead of computing and communication.

[0016] Optionally, in step S1, a double buffer mechanism is used to cover up the delay caused by each level of memory access. In the graphics processor, the complete matrix multiplication calculation process can be divided into the following three steps:

[0017] D1: data movement from global memory to shared memory;

[0018] D2: data movement from shared memory to registers;

[0019] D3: Use the data in the register to perform floating-point multiplication and accumulation instructions;

[0020] Among them, there is a dependency relationship between D1, D2 and D3. D1 will be executed at the beginning of each large loop, and several small loops will be executed after D1 ends. D2 needs to be executed at the beginning of each small loop, and D3 can be executed only after D2 ends.

[0021] Optionally, in step S1, D1, D2 and D3 can be pipelined by using double buffer space; specifically, D1 and D2 are executed before the start of the first large loop, and the data required for the first large loop is pre-fetched into the shared memory, and the data required for the first small loop is pre-fetched into the register; thereafter, at the start of each large loop, the data required for the next large loop is pre-fetched, and the memory access overhead of the global memory is overlapped with all subsequent small loop steps; at the start of each small loop, the data required for the next small loop is pre-fetched into the register, and the memory access overhead of the shared memory is overlapped with the floating-point multiplication and accumulation instructions.

[0022] Optionally, in step S1, the multi-level storage of the graphics processor is fully utilized by vectorizing memory access and eliminating conflicts within the unit. For shared memory access, in each small loop, the A matrix elements with a size of m*1 and the B matrix elements with a size of 1*n are moved from the shared memory to the registers of the corresponding threads. When the tiles of the A matrix are moved from the global memory to the shared memory, they need to be stored in the shared memory in the transposed order. Memory padding is added to make the memory access locations of every 8 threads in the same thread bundle fall into different shared memory units. In the process of moving data from the shared memory to the register, each thread no longer calculates a continuous block of data, but calculates four separated blocks of data.

[0023] Optionally, in step S2, the bias addition, tanh activation function and SGEMM operator are fused to obtain the Linear-tanh operator, and the upper limit of the number of registers that can be used by each thread is limited to 128. At the same time, the PTX instruction is used to implement the tanh calculation to reduce the use of registers, and all the calculation steps in the residual connection of the modified MLP are fused with the Linear-tanh operator. In this step, the registers storing the elements of the A matrix and the B matrix are reused to store the elements of the U matrix and the V matrix. Each thread uses the vectorized memory access instruction to obtain the required data according to the tile of the Linear-tanh output result. For the back-propagation calculation process, the unfused state is maintained, and the customized SGEMM operator is used to replace the original operator;

[0024] An additional memory access step is added during the forward propagation of the fusion operator to write the output of the tanh activation function back to the global memory. During the backward propagation, the output of the fusion operator is used to recalculate the output value of the tanh activation function through the inverse process of the residual connection calculation process. During the backward propagation process, the S matrix multiplication and the activation function calculation are re-executed to obtain the output value of the tanh activation function for gradient calculation.

[0025] Optionally, step S3 includes data parallelism, model parallelism and communication hiding. When the distributed training of the neural network model is performed in a data parallel manner, a complete set of network model parameters is stored on each graphics processor, and the data set is divided into several subsets and assigned to each graphics processor. During the training process, each graphics processor first uses its own local model parameters and data set subset to independently perform forward propagation and back propagation calculations and obtain local gradients. Subsequently, the gradients are aggregated through communication between the graphics processors. After obtaining the complete gradients, each graphics processor uses the gradients to update the local model parameters.

[0026] In model parallelism, the complete model parameters are divided among the GPUs. During the training process, each GPU is only responsible for the calculation and parameter update of the locally stored part of the network model.

[0027] The global reduction for aggregating gradients is decomposed into reduction and broadcast, which will overlap with the backpropagation calculation process and the forward propagation calculation process of the next iteration respectively.

[0028] After adopting the above technical solution, the beneficial effects of the present invention are:

[0029] The present invention improves the resource utilization of each hardware unit of the graphics processing unit (GPU), reduces the additional operator startup overhead and memory access overhead, and thus accelerates the reasoning speed of the PI-DeepONet model by customizing the SGEMM operator (single-precision matrix multiplication operator) and integrating it with the element-by-element calculation process. It uses data parallelism for distributed parallel training, optimizes the communication strategy between computing devices, overlaps the calculation process with the communication process, reduces the time required for model training, and improves the scalability of parallel training. BRIEF DESCRIPTION OF THE DRAWINGS

[0030] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the drawings required for use in the embodiments or the description of the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying creative labor.

[0031] Figure 1 This is the structure of PI-DeepONet;

[0032] Figure 2 It is the structure of Modified MLP;

[0033] Figure 3 It is the blocking strategy at the thread block level;

[0034] Figure 4 It is a shared memory access optimization;

[0035] Figure 5 It is the thread-level blocking strategy and memory access optimization;

[0036] Figure 6 It is communication optimization for distributed data parallelism;

[0037] Figure 7 is the acceleration effect of the matrix multiplication operator with fixed N and K sizes and changed M size, where (a) is the matrix multiplication without bias term addition in the experiment of matrix multiplication operator 1, (b) is the matrix multiplication with bias term addition in the experiment of matrix multiplication operator 1, (c) is the matrix multiplication without bias term addition in the experiment of matrix multiplication operator 2, (d) is the matrix multiplication with bias term addition in the experiment of matrix multiplication operator 2, (e) is the matrix multiplication without bias term addition in the experiment of matrix multiplication operator 3, and (f) is the matrix multiplication with bias term addition in the experiment of matrix multiplication operator 3;

[0038] Figure 8is the acceleration effect of the matrix multiplication operator with fixed M size and changed N and K sizes, where (a) is the matrix multiplication without bias term addition in the experiment of matrix multiplication operator 1, (b) is the matrix multiplication with bias term addition in the experiment of matrix multiplication operator 1, (c) is the matrix multiplication without bias term addition in the experiment of matrix multiplication operator 2, (d) is the matrix multiplication with bias term addition in the experiment of matrix multiplication operator 2, (e) is the matrix multiplication without bias term addition in the experiment of matrix multiplication operator 3, and (f) is the matrix multiplication with bias term addition in the experiment of matrix multiplication operator 3;

[0039] Fig. 9 is the time consumption ratio of each part in the model, where (a) is the end-to-end inference time and (b) is the GPU computing time;

[0040] Fig.10 This is a model reasoning acceleration experiment. In this experiment, the size of the matrix multiplication involved in the model is fixed to N=96 and K=96, and the data batch size is changed. (a) is the end-to-end reasoning time, and (b) is the GPU computing time.

[0041] Fig.11 This is a model reasoning acceleration experiment. The data batch size is fixed at 50,000 and the size of the matrix multiplication involved in the model is changed. (a) is the end-to-end reasoning time, and (b) is the GPU computing time.

[0042] Fig.12 This is a distributed data parallel training scalability experiment, where (a) decouples the global reduction operation into reduction distribution and global aggregation, and (b) decouples the global reduction operation into reduction and broadcast. DETAILED DESCRIPTION

[0043] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

[0044] An embodiment of the present application discloses a method for improving the speed of reasoning, solving and training of a deep operator network model.

[0045] Embodiment, according to Figures 1 to 12 A method for improving the inference and training speed of deep operator network models is shown.

[0046] 1. Background

[0047] 1.1 Deep Operator Network (DeepONet) and Physical Information-Based Deep Operator Network (PI-DeepONet)

[0048] The DeepONet class model is a deep neural network model constructed based on the operator form of the universal approximation theorem of neural networks proved by Chen et al., which states that a neural network with a single hidden layer can accurately approximate any linear / nonlinear continuous operator. Based on the above theorem, the DeepONet framework contains two deep neural networks, the backbone network and the branch network: the backbone network is used to extract the potential representation of the input function, and the branch network is used to extract the potential representation of the input coordinate points. The final result is obtained by merging the outputs of the two sub-networks through dot products. PI-DeepONet is a model constructed on the basis of DeepONet by incorporating the idea of ​​PINN. Specifically, the output of the DeepONet network structure is differentiable relative to its input. Therefore, the constraints and boundary value conditions of partial differential equations can be used to construct additional loss functions as a regularization mechanism for the network model, so that the training results of the model tend to satisfy the potential constraints of the PDE itself. The model structure of PI-DeepONet is as follows: Figure 1 shown.

[0049] 1.2 Distributed Parallel Training of Deep Learning Models

[0050] Efficient and scalable parallel strategies are the key to multi-device collaboration and parallel training of deep learning models. In recent years, many different parallel methods have been proposed based on different model characteristics and device conditions. They can be roughly divided into three types: data parallelism, model parallelism, and pipeline parallelism. Of course, the three types of strategies can also be mixed to train deep learning models according to actual needs.

[0051] The most common distributed parallel training strategy is data parallelism. In this mode, each computing device stores a complete set of model parameters. The training data set is divided into several subsets and distributed to different computing devices. Each device uses its own local model parameters and data to complete the forward propagation and back propagation calculations during the training process to obtain gradient values, and shares each other's gradient values ​​and updates model parameters at the appropriate time. When and how to share gradient values ​​between different computing devices will be the key to affecting the efficiency of parallel training. In the parameter server (PS) system, devices are divided into computing nodes and parameter servers. After the computing nodes complete the computing tasks, they aggregate the gradients to the parameter server to update the model parameters. The use of the ring all-reduce algorithm can make the system decentralized and homogeneous. During communication, each computing node sends a piece of data to the next device and receives a piece of data sent by the previous device according to certain rules. The communication cost is independent of the system scale, overcoming the performance bottleneck that the centralized structure may encounter when expanding the device scale. Many works are centered on optimizing the communication strategy of the ring all-reduce algorithm. When the computing devices in the parallel system are highly heterogeneous, the asynchronously updated data parallel strategy can be used to improve the training performance. In addition, ZeRO can be regarded as a special data parallel strategy, which divides the model's optimizer state, parameter gradients and even model parameters to different devices for maintenance, while the additional communication cost can be masked by calculation, reducing the memory pressure during data parallel training and greatly expanding the scope of use of data parallel strategies.

[0052] Unlike data parallelism, model parallelism will split the complete model into several parts and distribute them on different computing devices to address the challenge that the model parameters are too large to fit into the memory space of a single computing device. Pipeline parallelism uses the structural characteristics of deep neural networks that calculate layer by layer to distribute the complete network model layer by layer to different computing devices, organizing the computing devices together in a pipeline manner for parallel training.

[0053] 2. Optimize the kernel function of the PI-DeepONet neural network model

[0054] 2.1 Challenges

[0055] The PI-DeepONet neural network model combines the physical information constraint idea of ​​the PINN model with the deep operator network. At the same time, in order to solve the gradient pathology that may be caused by the PINN model structure, a network structure called modified MLP is used in PI-DeepONet. This is a simple extension of the standard fully connected network. It can effectively improve the accuracy and robustness of the PINN model by reducing the rigidity of the gradient flow. Specifically, the modified MLP introduces two converter networks, which first project the input variables into high-dimensional features, and then according to certain rules, multiply the high-dimensional features with the output features of each hidden layer through residual connections, such as Figure 2 As shown in the figure. The modified MLP can effectively improve the accuracy of PI-DeepONet, but it will also introduce additional computational overhead. Incorporating high-dimensional features into each hidden layer output feature requires multiplication interaction, and its calculation process is a combination of multiple element-by-element calculation processes. Frequently starting element-by-element operators during training will significantly reduce the utilization of computing power of computing devices. At the same time, the width of the fully connected layer used in the PI-DeepONet network model is small, but a large batch size is used during training, which will cause two effects. First, large-size features need to be transmitted between the fully connected layer and each element-by-element operator, further exacerbating the efficiency reduction caused by frequently starting element-by-element operators; at the same time, this will reduce the efficiency of the SGEMM operator, especially after the model introduces the PINN structure. In order to calculate the residual loss value, the back propagation process of each module will be called multiple times in an iterative training process. The number of calls to the SGEMM operator increases significantly, further affecting the overall performance of the neural network model in training.

[0056] In order to improve the training and reasoning speed of the PI-DeepONet neural network model, the present invention reduces the additional operator startup overhead and data movement overhead by operator fusion, and customizes the SGEMM operator for the implementation of the fusion operator.

[0057] 2.2 Customizing the SGEMM Operator

[0058] In order to achieve efficient operator fusion for the PI-DeepONet neural network model, the present invention customizes the SGEMM operator for it. First, the present invention designs a three-level block strategy of thread blocks, warps and threads to achieve a high computing memory access ratio, then the present invention uses a double buffering mechanism to cover up the delay caused by each level of memory access to a certain extent, and finally further reduces the memory access overhead by using vectorized memory access instructions and eliminating intra-unit conflicts in the shared memory access process.

[0059] 2.2.1 Block strategy

[0060] The fully connected layers in DeepONet-like models are usually small in width and use a larger batch size in training. After being converted into matrix multiplication, the calculation process is the multiplication of a tall matrix with a small matrix. First, at the thread block level, the present invention uses shared memory to improve the computational memory access ratio of each thread block. The block strategy at the thread block level, such as Figure 3 As shown, the present invention only divides the calculation along the dimension M and distributes it to each thread block, and no division is performed on the dimension N. In order to complete the calculation for which it is responsible, each thread block needs to obtain an A matrix block of size m*K and a B matrix block of size K*n (n=N). By placing the A matrix block and the B matrix block into the shared memory, the number of accesses to the global memory can be reduced, but due to the size of the shared memory, the A matrix block and the B matrix block required by each thread block will be divided into K / k blocks along the dimension K again and stored in batches in the shared memory for calculation. The present invention refers to the calculation of a batch as a large cycle, and the final calculation result of each thread block is the specification of the calculation results of K / k large cycles.

[0061] In each block, the 8 thread warps are arranged as follows Figure 3 As shown. The present invention gives priority to ensuring that the 32 threads in each thread bundle are arranged in a 4*8 shape to improve the computational memory access ratio of a single thread bundle. On this basis, the present invention allocates computations to each thread bundle according to the actual sizes of m and N and tries to achieve load balancing. M, N, and K all describe the sizes of the three matrices in matrix multiplication. The height and width of matrix A are M and K, the height and width of matrix B are K and N, and the height and width of matrix C are M and N; m, n, and k all describe the sizes of the three matrices after being divided into blocks.

[0062] At the thread level, the present invention converts the calculation of matrix multiplication into multiple vector outer product calculations and uses registers to improve the calculation memory access ratio of each thread. Under the calculation rules of matrix multiplication, the calculation process in each thread-level block is a nested loop of MNK, which will cause repeated shared memory access. The present invention adjusts the order of the nested loops to KMN, and uses registers corresponding to the thread-level block size (m*n) to store intermediate values ​​for accumulation, avoiding repeated shared memory access. At this time, the calculation actually becomes the vector outer product calculation form of matrix multiplication. In the calculation form of the vector outer product, the calculation of each large loop will be divided into k small loops. In each small loop, the thread needs to move the m*1 A matrix elements and the 1*n B matrix elements from the shared memory into its own register for the vector outer product calculation of this small loop.

[0063] 2.2.2 Double buffering mechanism

[0064] The present invention uses a double buffer mechanism to cover up the delay caused by memory access at various levels to a certain extent. In a graphics processing unit (GPU), a complete matrix multiplication (GEMM) calculation process can be divided into the following three steps: step one: data movement from global memory to shared memory; step two: data movement from shared memory to register; step three: using the data in the register to perform a floating-point multiplication and accumulation instruction (FFMA). These three steps are dependent on each other: step one will be executed at the beginning of each large loop, and several small loops will be executed after step one ends. Step two needs to be executed at the beginning of each small loop, and step three can only be executed after step two ends. The three steps can be pipelined by using double buffer space. Specifically, steps one and two are executed before the start of the first large loop, and the data required for the first large loop is pre-fetched into the shared memory, and the data required for the first small loop is pre-fetched into the register; thereafter, the data required for the next large loop is pre-fetched at the start of each large loop, and the memory access overhead of the global memory is overlapped with all subsequent small loop steps; at the start of each small loop, the data required for the next small loop is pre-fetched into the register, and the memory access overhead of the shared memory is overlapped with the floating-point multiplication and accumulation instructions.

[0065] 2.2.3 Memory Access Optimization

[0066] The present invention makes full use of the multi-level storage of the graphics processor by vectorizing memory access and eliminating conflicts within the unit. For access to shared memory, in each small loop, the A matrix elements with a size of m*1 and the B matrix elements with a size of 1*n are moved from the shared memory to the registers of the corresponding threads. The present invention uses vectorized read instructions to access the matrix elements stored in the shared memory, which can reduce the number of memory access instructions and further improve the computational memory access ratio. In the vector outer product calculation form of matrix multiplication, each thread can naturally use vectorized read instructions to access the B matrix elements stored continuously in the shared memory, but the memory access position is non-continuous when accessing the A elements in the shared memory. For this reason, the present invention needs to store the blocks of the A matrix in the shared memory in the transposed order when moving them from the global memory to the shared memory. At the same time, in order to avoid conflicts within the shared memory units in this process, the present invention adds shared memory padding to allow the memory access positions of every 8 threads in the same thread bundle to fall into different shared memory units, such as Figure 4 In the process of moving data from shared memory to register, in order to avoid conflicts within the unit, the present invention changes the thread-level block division, and each thread no longer calculates a continuous block of data (usually 8*8), but calculates four separate blocks of data (each block of data is usually 4*4), as shown in FIG. Figure 5 shown.

[0067] 2.3 Operator Fusion

[0068] The present invention first fuses the bias addition, tanh activation function and SGEMM operator to obtain the Linear-tanh operator. In order to ensure a high thread bundle occupancy rate, the present invention limits the upper limit of the number of registers that can be used by each thread to 128. At the same time, the PTX instruction is used to implement the tanh calculation to reduce the use of registers. Then, all the calculation steps in the residual connection of the modified MLP are fused with the Linear-tanh operator. In this step, the registers storing the A matrix elements and the B matrix elements are reused to store the U matrix elements and the V matrix elements. Each thread uses vectorized memory access instructions to obtain the required data according to the block of Linear-tanh output results. For the back-propagation calculation process, the present invention maintains the unfused state, but uses a customized SGEMM operator to replace the original operator.

[0069] The tanh activation function needs to use the output calculated by the activation function during forward propagation in the process of back propagation to find the gradient. If the fusion state of the Linear-tanh operator and the residual connection operator is to be maintained, the intermediate quantity of the fusion operator needs to be saved for the calculation of the gradient in the back propagation, or the output value of the tanh activation function needs to be recalculated in the back propagation. The present invention first attempts to add an additional memory access step in the forward propagation process of the fusion operator, and writes the output of the tanh activation function back to the global memory, but this will cause register overflow and break the overlapping state of calculation and memory access in the original operator, seriously affecting the calculation efficiency of the forward propagation of the fusion operator. Then the present invention attempts to use the output of the fusion operator during back propagation, and recalculate the output value of the tanh activation function through the inverse process of the residual connection calculation process, but this will introduce a large number of additional division calculations, causing the back propagation calculation speed to drop seriously. Finally, the strategy adopted by the present invention is to re-execute the calculation of SGEMM and the activation function during the back propagation process to obtain the output value of the tanh activation function for gradient calculation.

[0070] 3. Distributed parallel training

[0071] In order to efficiently utilize more computing devices for training and further shorten the training time of the PI-DeepONet model, the present invention customizes a suitable distributed training strategy for the PI-DeepONet model based on its structural characteristics.

[0072] 3.1 Data Parallelism and Model Parallelism

[0073] Data parallelism and model parallelism are the two most commonly used methods in distributed training. When using data parallelism to perform distributed training of neural network models, a complete set of network model parameters is stored on each GPU, and the data set is divided into several subsets and assigned to each GPU. During the training process, each GPU first uses its own local model parameters and dataset subset to independently perform forward propagation and back propagation calculations and obtain local gradients. Then, the gradients are aggregated through communication between the GPUs. After obtaining the complete gradients, each GPU uses the gradients to update the local model parameters. Unlike data parallelism, in model parallelism, the complete model parameters are divided among the GPUs according to certain rules. During the training process, each GPU is only responsible for the calculation and parameter update of the locally stored part of the network model.

[0074] The essence of distributed training is to use multiple computing devices to parallelize the calculations in the model training process. Usually, it will not cause additional computing overhead. The communication overhead required in the parallel process will be the key to affecting the efficiency of distributed training. Due to differences in model structure and training process, choosing different parallel strategies will result in different communication overheads. The PI-DeepONet model has a small number of parameters, but a large batch size is usually required to ensure the training effect during training. Considering the above two characteristics, if distributed training is performed using model parallelism, the size of the activation value tensor that needs to be transmitted between each graphics processor is large, which will cause a lot of communication overhead and reduce training efficiency; when using data parallelism for distributed training of the PI-DeepONet model, each graphics processor only needs to communicate the gradient tensor used for parameter update, which is the same size as the model weight, and the communication overhead is very small.

[0075] 3.2 Communication Hiding

[0076] When using data parallelism for distributed training, in order to maintain the consistency of the model on each GPU, communication is required before updating the model parameters so that each GPU can obtain the complete aggregated gradient, which will cause additional communication overhead during the training process. By combining tensor fusion with the wait-free backpropagation algorithm, the communication of this part of the gradient will start immediately after the parameter gradient calculation of a group of network layers is completed. This can overlap part of the communication overhead with the calculation overhead of backpropagation to achieve a certain degree of communication masking. However, simply overlapping the communication process with the calculation process of backpropagation cannot completely cover the communication. Dear pointed out that if the communication process can be decoupled into two sub-communication processes without additional overhead, they can be overlapped with the backpropagation calculation process and the forward propagation calculation process of the next iteration respectively to obtain a greater degree of communication masking.

[0077] The present invention decomposes the global reduction for aggregating gradients into reduction and broadcast, which will overlap with the back propagation calculation process and the forward propagation calculation process of the next iteration respectively. Figure 6 As shown in the figure, the parameters of the entire neural network model are divided into several parameter groups. During the calculation process of the back propagation stage, whenever the calculation of a layer of the network is completed, its gradient will be packed into the buffer of the parameter group where the network layer is located. When the gradients of all model parameters in a parameter group are collected, the gradients in the buffer are immediately reduced and communicated. The gradients on all graphics processors will be sent to the graphics processor No. 0 (GPU0) and the average value will be calculated, but the gradient broadcast will not be performed immediately. In this process, the back propagation calculation process of parameter group n-1 will overlap with the gradient reduction communication process of parameter group n. In the complete back propagation process, the gradients of the first group of parameters are calculated last. When the gradients of the first group of parameters are all packed into the buffer, the complete communication process and parameter update are performed on them, that is, the reduction and broadcast are performed, and the parameters in parameter group 1 are updated after the broadcast is completed. When all model parameters in parameter group 1 are updated, the forward propagation calculation of the next round of iteration can be started, and the gradients in the buffer of the subsequent parameter groups will continue to be broadcasted. The forward propagation process of parameter group n will overlap with the gradient broadcast communication of parameter group n-1.

[0078] 4. Experiment

[0079] This section introduces the experiments conducted in this paper and the hardware configuration used to conduct the experiments. The experiment in 4.1 is used to test the performance of the customized SGEMM operator, the experiment in 4.2 tests the reasoning speed improvement brought by the customized SGEMM operator and the operator fusion strategy to the overall model, and the experiment in 4.3 tests the scalability of the distributed data parallel strategy in Section 4. The experiments in 4.1 and 4.2 are both conducted on a single RTX4090 graphics processor (RTX4090 GPU), and the experiment in 4.3 is conducted on a cluster composed of V100 graphics processors (V100 GPUs).

[0080] 4.1 Testing the Performance of Custom SGEMM Operator

[0081] The present invention designs three SGEMM operators according to the strategy in 2.2, which have high efficiency in single-precision matrix multiplication calculations of N<=64, N<=96, and N<=128, respectively. Their scope of application covers the most commonly used matrix multiplication calculations in deep operator networks. Hereinafter, these three SGEMM operators will be referred to as operator 1, operator 2, and operator 3, respectively. In order to prove that the SGEMM operator customized for the deep operator network of the present invention is highly efficient, the present invention designs two groups of experiments: the first group is to fix the size of N and K, change the size of M, and then compare it with the single-precision matrix multiplication in the basic linear algebra library of the CUDA computing platform; the second group is to fix the size of M, change the size of N and K within the size range commonly used in deep operator networks, and then compare it with the single-precision matrix multiplication in the basic linear algebra library of the CUDA computing platform; at the same time, in the above two groups of experiments, the present invention also considers two situations with or without deviation.

[0082] In the first set of experiments, the present invention first fixes the size of N and K to 64, then changes the size of M in the range of 10000 to 100000, and compares operator 1 with the single-precision matrix multiplication of the basic linear algebra library of the CUDA computing platform with the highest efficiency in the corresponding size. The experimental results are as follows: Figure 7 (a) and Figure 7 Then the present invention changes the size of N and K to 96, and repeats the above experiment using operator 2. The experimental results are shown in Figure 7 (c) and Figure 7 As shown in (d) in the figure; finally, the present invention changes the size of N and K to 128, and repeats the above experiment using operator 3. The experimental results are shown in Figure 7 (e) and Figure 7 As shown in (f) in .

[0083] In the second set of experiments, the present invention fixes the size of M to 30000, changes the sizes of N and K in the range of 32 to 64, and compares operator 1 with the single-precision matrix multiplication of the basic linear algebra library of the CUDA computing platform with the highest efficiency in the corresponding size. The experimental results are as follows: Figure 8 (a) and Figure 8 As shown in (b) in the figure, the sizes of N and K are then changed in the range of 48 to 96, and the operator 2 is compared with the single-precision matrix multiplication with the highest efficiency in the corresponding size. The experimental results are shown in Figure 8 (c) and Figure 8 As shown in (d) in the figure; finally, the sizes of N and K are changed in the range of 64 to 128, and the operator 3 is compared with the single-precision matrix multiplication of the basic linear algebra library of the CUDA computing platform with the highest efficiency in the corresponding size. The experimental results are shown in Figure 8 (e) and Figure 8 As shown in (f) in .

[0084] It can be seen that within the size range involved in the two sets of experiments, the SGEMM operator customized by the present invention is more efficient than the implementation in the basic linear algebra library of the CUDA computing platform in most cases, and can achieve a maximum speedup of about 1.5 times. Therefore, for single-precision matrix multiplication operations in deep operator networks, the SGEMM operator customized by the present invention can provide higher computing efficiency, and these SGEMM operators also serve as a good foundation for subsequent operator fusion work.

[0085] 4.2 Test the reasoning speed improvement brought by the customized SGEMM operator and operator fusion strategy to the overall model

[0086] In order to verify that the strategy in 2.2 improves the model reasoning speed, the present invention uses modified MLP to build a PI-DeepONet model for solving the Burgers equation. It has 7 fully connected layers, each layer has 96 units, and uses the Tanh activation function. The present invention first counts the proportion of time consumed by various computing tasks in the reasoning process of the network model, such as Fig. 9 As shown in (a) in the figure. As can be seen from the figure, the largest proportion is single-precision matrix multiplication, but the element-by-element operator (excluding Tanh) also accounts for as high as 41.5%. These element-by-element calculations mainly come from the special residual connection structure of modified MLP. Each residual connection structure contains several matrix element-by-element multiplications and matrix additions. Their calculation amount is much smaller than that of single-precision matrix multiplication. The time they consume mainly comes from frequent operator startup and frequent access to global memory for reading or writing large-size tensors.

[0087] For the experiment of model reasoning acceleration, the present invention changes the batch size in the range of 1000 to 10000, and compares the reasoning speed of the original model and the optimized model at each size. In addition, two indicators, end-to-end reasoning time and GPU computing time, are considered at each size. The two sets of experimental results are shown in Figure 2. Fig.10 As shown in the figure, for the GPU computing time, Fig.10 As can be seen from (a) in Figure 2, the optimized model has achieved more than 3 times acceleration in all sizes, with the highest acceleration being about 5.6 times. Fig.10 As can be seen from (b) in the figure, the optimized model has a stable speedup of more than 2 times in most sizes.

[0088] The present invention also counts the proportion of time consumed by various computing tasks in the inference process of the optimized network model, such as Fig. 9As shown in (b) of the figure. It can be seen that the fusion operator composed of SGEMM, tanh activation function and residual connection takes up the largest proportion of time, followed by the fusion operator composed of SGEMM and tanh activation function. The independent element-by-element calculation only takes about 4.5% of the time. This shows that a large number of scattered element-by-element calculations in the model have been effectively fused together, and the frequent operator startup overhead and memory access overhead brought by these element-by-element operators have been eliminated.

[0089] In addition, in order to prove that the strategy in 2.2 can effectively bring about reasoning acceleration within the size range commonly used by the PI-DeepONet model, the present invention also conducted experiments on models of more sizes. In these experiments, the model still consists of 7 fully connected networks, tanh activation function and modifiedMLP residual connection structure. The present invention fixes the batch size at 50,000, and then changes the number of units in the fully connected network in the range of 32 to 128. For each model size, the present invention compares the reasoning speed of the original model and the optimized model, which includes two indicators: end-to-end reasoning time and graphics processor computing time. The experimental results are as follows. Fig.11 As shown in the figure, for GPU computing time, the optimized model achieves at least 3.5 times acceleration in all sizes, and up to about 5.7 times acceleration; for end-to-end reasoning time, when the number of units in the fully connected network is greater than 64, the optimized model can stably achieve about 2 times acceleration.

[0090] The experimental results show that the strategy in 2.2 significantly improves the inference speed of the PI-DeepONet model. This acceleration effect mainly comes from two aspects. Fig. 9 , we can see that the single-precision matrix multiplication calculation in the fully connected layer and the element-by-element calculation in the residual connection are the two aspects that take up the most time in the inference process of the model. For the former, the present invention improves its efficiency through a customized SGEMM operator; for the latter, the present invention greatly reduces the operator startup overhead and memory access overhead through operator fusion.

[0091] 4.3 Data Parallelism

[0092] In order to verify the effectiveness of the strategy in Section 3, the present invention continues to use the PI-DeepONet model constructed in 4.2 for distributed data parallel training. The present invention conducts two sets of experiments. The first set decouples the global protocol into protocol and broadcast, and the second set decouples the global protocol into protocol distribution and global aggregation. The results of the two sets of experiments are shown in Figure 2. Fig.12 As shown, the scalability of 83%-85% was always maintained in the experiment.

[0093] In summary, the present invention optimizes PI-DeepONet in parallel at multiple levels to improve its reasoning and training speed. First, for the special shape of matrix multiplication involved in PI-DeepONet during training and reasoning, a SGEMM operator is customized for it. Experiments show that within the model size commonly used by PI-DeepONet, the SGEMM operator customized by the present invention can achieve a speedup ratio of up to about 1.5 times compared to the best single-precision matrix multiplication implementation in the basic linear algebra library of the CUDA computing platform. Then, a large number of scattered element-by-element calculations in the model are fused with the SGEMM operator, which effectively reduces the additional memory access overhead and operator startup overhead, and the end-to-end reasoning speed of the model is accelerated by more than 2 times, and the graphics processor calculation time can be reduced by up to 5.7 times. Finally, data parallelism is used for distributed training, and 83%-85% scalability can be maintained through communication masking.

[0094] The above are only preferred embodiments of the present invention and are not intended to limit the present invention. Any modifications, equivalent substitutions, improvements, etc. made within the technical solution of the present invention should be included in the protection scope of the present invention.

Claims

1. A method for improving the speed of reasoning, solving and training of deep operator network models, characterized in that: The following steps are involved: S1. Customized SGEMM operator for PI-DeepONet neural network model: designed three-level block strategies of thread block, thread bundle and thread to achieve high computation-to-memory ratio, and optimized memory access mode to reduce memory access time overhead; S2, design operator fusion strategy; The large number of scattered element-by-element calculations in the PI-DeepONet neural network model are fused with matrix multiplication. The time consumed by these element-by-element calculations comes from frequent startup of operators and frequent access to global memory for reading or writing large-size tensors. These extra overheads are effectively eliminated through operator fusion. S3. Design a distributed parallel training strategy to reduce the time required for PI-DeepONet neural network model training by overlapping the time overhead of computation and communication; In step S2, the bias addition, tanh activation function and SGEMM operator are fused to obtain the Linear-tanh operator. The upper limit of the number of registers that each thread can use is limited to 128. At the same time, the PTX instruction is used to implement the tanh calculation to reduce the use of registers. All calculation steps in the residual connection of the modified MLP are fused with the Linear-tanh operator. In this step, the registers storing the elements of the A matrix and the B matrix are reused to store the elements of the U matrix and the V matrix. Each thread uses the vectorized memory access instruction to obtain the required data according to the tile of the Linear-tanh output result. For the back-propagation calculation process, the unfused state is maintained, and the customized SGEMM operator is used to replace the original operator. An additional memory access step is added during the forward propagation of the fusion operator to write the output of the tanh activation function back to the global memory. During the backward propagation, the output of the fusion operator is used to recalculate the output value of the tanh activation function through the inverse process of the residual connection calculation process. During the backward propagation process, the S matrix multiplication and the activation function calculation are re-executed to obtain the output value of the tanh activation function for gradient calculation; Step S3 includes data parallelism, model parallelism and communication hiding. When the distributed training of the neural network model is performed in a data parallel manner, a complete set of network model parameters is stored on each GPU, and the data set is divided into several subsets and assigned to each GPU. During the training process, each GPU first uses its own local model parameters and data set subset to independently perform forward propagation and back propagation calculations and obtain local gradients. Then, the GPUs aggregate the gradients through communication. After obtaining the complete gradients, each GPU uses the gradients to update the local model parameters. In model parallelism, the complete model parameters are divided among the GPUs. During the training process, each GPU is only responsible for the calculation and parameter update of the locally stored part of the network model. The global reduction for aggregating gradients is decomposed into reduction and broadcast, which will overlap with the backpropagation calculation process and the forward propagation calculation process of the next iteration respectively.

2. A method for improving the speed of reasoning, solving and training of deep operator network models according to claim 1, characterized in that: In step S1, a double buffer mechanism is used to cover up the delay caused by memory access at each level. In the graphics processor, the complete matrix multiplication calculation process can be divided into the following three steps: D1: data movement from global memory to shared memory; D2: data movement from shared memory to registers; D3: Use the data in the register to perform floating-point multiplication and accumulation instructions; Among them, there is a dependency relationship between D1, D2 and D3. D1 will be executed at the beginning of each large loop, and several small loops will be executed after D1 ends. D2 needs to be executed at the beginning of each small loop, and D3 can be executed only after D2 ends.

3. A method for improving the speed of reasoning, solving and training of deep operator network models according to claim 2, characterized in that: In step S1, D1, D2 and D3 can be pipelined by using double buffer space; specifically, D1 and D2 are executed before the first large loop starts, the data required for the first large loop is pre-fetched into the shared memory, and the data required for the first small loop is pre-fetched into the register; after that, the data required for the next large loop is pre-fetched at the beginning of each large loop, so that the memory access overhead of the global memory and all subsequent small loop steps are overlapped; At the beginning of each small loop, the data required for the next small loop is pre-fetched and placed in the register, overlapping the shared memory access overhead and the floating-point multiplication and accumulation instructions.

4. A method for improving the speed of reasoning, solving and training of deep operator network models according to claim 3, characterized in that: In step S1, the multi-level storage of the graphics processor is fully utilized by vectorizing memory access and eliminating conflicts within the unit. For shared memory access, in each small loop, the A matrix elements with a size of m*1 and the B matrix elements with a size of 1*n are moved from the shared memory to the registers of the corresponding threads. When the tiles of the A matrix are moved from the global memory to the shared memory, they need to be stored in the shared memory in the transposed order. Memory padding is added to make the memory access locations of every 8 threads in the same thread bundle fall into different shared memory units. In the process of moving data from the shared memory to the register, each thread no longer calculates a continuous block of data, but calculates four separated blocks of data.

Images