An apparatus for training a softmax function in a large language model

Abstract

This invention provides a device for training the Softmax function in a large language model. The upper part of the device is a forward propagation path, and the lower part is a backward propagation path. The forward propagation path includes e x The exponential function unit, adder, and divider, in e x Registers are inserted between the exponential function unit, adder, and divider; the backpropagation path includes two multipliers, one adder A1, and a multiplexer MUX for reconstructing the data path; the two multipliers are left multiplier B1 and right multiplier B2; the multiplexer MUX is used to change the data flow within the device; the forward and backward propagation paths share two random access memories, RAM1 and RAM2. This invention can be applied to the computation of the Softmax function at various stages of training, thereby better utilizing computational and storage resources to achieve higher performance and energy efficiency.

Description

Technical Field

[0001] This invention relates to an apparatus for training the Softmax function in a large language model. Background Technology

[0002] The Transformer is a classic model for NLP (Natural Language Processing) proposed by Google in 2017. The Transformer model uses a self-attention mechanism, allowing for parallel training and the acquisition of global information about the samples. Popular models such as BERT and GPT are also based on the Transformer architecture.

[0003] In recent years, Transformer-based deep neural networks (DNNs) have achieved remarkable results in fields such as natural language processing (NLP), computer vision (CV), and speech processing. Transformer-based models are typically pre-trained on large-scale datasets and then fine-tuned for downstream tasks. With the expanding application scenarios of Transformer-type models, training (fine-tuning) these models on edge platforms has become crucial, considering data privacy and real-time processing requirements. However, due to the massive number of parameters and high computational complexity of Transformer-type models, deploying the fine-tuning training process on resource-constrained edge platforms faces numerous challenges. Transformer-type models consist of Transformer layers, which utilize an attention mechanism called self-attention. As the model size and the length of the sample sequences processed by the model continue to increase, the computational cost of Softmax in the attention mechanism during inference and training also increases, becoming one of the bottlenecks restricting model deployment efficiency. Current hardware designs for Softmax are mostly geared towards the inference stage of the model, and their applications are primarily concentrated on traditional convolutional neural networks. Existing technical solutions mainly use mathematical transformations to transform the complex exponential function (e) in the Softmax function. x The division operation is converted into a solution with lower complexity and more suitable for hardware implementation. Summary of the Invention

[0004] Purpose of the invention: The technical problem to be solved by the present invention is to address the shortcomings of the prior art by providing a device for training the Softmax function in a large language model, wherein the large language model is a Transformer-type model, characterized in that the upper part of the device is a forward propagation path and the lower part is a backward propagation path.

[0005] The forward propagation path includes e x The exponential function unit, adder, and divider, in e x Registers were inserted between the exponential function unit, the adder, and the divider;

[0006] The backpropagation path includes two multipliers, one adder A1, and a multiplexer MUX for reconstructing the data path; the two multipliers are the left multiplier B1 and the right multiplier B2.

[0007] The multiplexer MUX is used to change the data flow within the device;

[0008] The forward propagation path and the backward propagation path share two random access memories, RAM1 and RAM2.

[0009] During the forward propagation phase, the Softmax function is calculated row-wise on the M matrix, where a row vector m = (m1, m2, ..., mn) in the matrix M is a vector. s ), m s This represents the s-th element in vector m, where the subscript indicates the element's position within the vector. The length of vector m is s. The result vector obtained by Softmax is n = (n1, n2, ..., nn...). s The i-th element n in n i The calculation process is expressed as follows:

[0010]

[0011] Where i takes values ​​from 1 to s;

[0012] The calculation process of formula (1) includes three stages: the calculation of each element. calculate, Summation and division.

[0013] The forward propagation path uses a serialization process to perform the forward propagation calculation: the input data are the elements m1, m2, ..., m in the vector m. s Input data is stored in RAM1; after the calculation begins, m1, m2, ..., m s Retrieve sequentially from RAM1, first passing through e x The exponential function unit performs exponential operations, with each element m1, m2, ..., m s ex Calculation results The accumulation is performed in the adder, while e x The calculation result is stored in RAM2; after all elements in vector m have undergone exponentiation, the adder accumulates the results. It was stored in RAM2; then, The data is transmitted to the divisor input port of the divider. The numbers are retrieved from RAM2 and sequentially transferred to the dividend input port of the divider. The outputs of the divider are n1, n2, ..., n in sequence. s .

[0014] During serialization, the division calculation needs to wait until all elements in vector m have completed e. x The division calculation can only begin after the exponentiation operation and the accumulated result are obtained, dividing the calculation into two stages: e x The exponentiation and accumulation phase and the division phase are performed in a pipelined manner.

[0015] Set the input matrix during the forward propagation of the Softmax function. Let M be a real matrix with s rows and s columns, and let the row vectors in matrix M be represented as m1, m2, ..., m2, ..., m3, respectively. s , where m s Let m represent the vector in the s-th row of matrix M. s =(m s1 m s2 , ..., m ss ), m ss Represents vector m s The s-th element in the vector m has an index indicating its position in the vector. s The length is s; after m1 completes the exponentiation and accumulation operations and performs the division operation, m2 starts the new exponentiation and accumulation operations without m1 having to wait for the calculation to complete; the calculations in the two stages are performed in parallel.

[0016] During the backpropagation phase of the backpropagation path, the gradient propagated from the subsequent network structures of Softmax is dn. These subsequent network structures refer to the computations following the Softmax function in Transformer-type models, such as matrix multiplication. dn represents the gradient of the output data n of the Softmax function during the forward propagation phase relative to the final output of the model. During the backpropagation phase, the gradient dm of vector m relative to the final output of the model needs to be calculated. dm is obtained using the following formula:

[0017] dm = dn·(diag(n)-n) T ·n)

[0018] Where diag represents a diagonal matrix function, T represents matrix transpose; and the element dm in dm... i The calculation process can be broken down as follows:

[0019]

[0020] Where 1≤i≤s.

[0021] When the backpropagation path of the device is used for backpropagation calculation, RAM1 stores the calculation results n1, n2, ..., n of the forward propagation. s RAM2 stores the gradient values ​​dn1, dn2, ..., dn passed from subsequent Softmax network structures. s dn s Let represent the s-th element in vector dn, where the subscript indicates the position of the element in the vector, and the length of vector dn is s.

[0022] In the process of calculating the gradient dm of vector m, the backpropagation path reconstructs the internal data flow through a multiplexer, thereby completing the calculation of each element in dm. For example, in calculating any element dm in dm... i The calculation process is as follows (dm) i Let represent the i-th element in vector dm, where the index i indicates the element's position within the vector. The length of vector dm is s.

[0023] When calculating dm i During the first clock cycle of the calculation, the connection method of the backpropagation path has the following characteristics:

[0024] At this moment, the data read from RAM1 is n. i The data read from RAM2 is dn i ;

[0025] The input to the left-hand multiplier B1 is n. i and -n i The output is -n i ·n i ;

[0026] The input to the multiplier B2 on the right is n. i and dn i The output of B2 is n i ·dn i ;

[0027] The left input of adder A1 is n i ·dn i The input above A1 is 0, indicating the start of accumulation. The output of A1 in the current cycle is n. i ·dn i ;

[0028] In calculating the t-th clock cycle of dm1, where 2≤t≤s+1, the connection method of the back propagation path has the following characteristics:

[0029] At this moment, the data read from RAM1 is n. t n t This represents the t-th element in vector n, where the index indicates the element's position within the vector. All elements of vector n are stored in RAM1; the data read from RAM2 is dn. t-1 ;dn t-1 Let represent the (t-1)th element in vector dn, and let the subscript indicate the position of the element in the vector. All elements of vector dn are stored in RAM2.

[0030] The left input of the left multiplier B1 is kept at n by the control of the multiplexer MUX. i The input above B1 is -n t The output of B1 in the current cycle is -n i ·n t .

[0031] The input above the multiplier B2 on the right is dn. t-1 At this point, the input to the left of B2 is the output of B1 in the previous cycle - n. i ·n t-1 The output of B2 is -n i ·n t-1 ·dn t-1 ;

[0032] The left input of adder A1 is the output of B2 - n i ·n t-1 ·dn t-1 The upper input of A1 is controlled by the multiplexer MUX to be the output of A1 in the previous cycle. At this time, A1 has the function of an accumulator.

[0033] After s+1 cycles, the accumulated output of A1 is At this point, dm was calculated. i .

[0034] Furthermore, the computational throughput of the hardware modules can be further improved through parallel design. Computation can be performed in parallel by two or more devices, taking advantage of the parallelism that exists during the training of Transformer-type models.

[0035] Beneficial Effects: This invention proposes a flexible and efficient hardware architecture for Softmax computation during the training process of Transformer-type models. Utilizing a pipelined design approach, it can be applied to the computation of the Softmax function at various stages of training, thereby better utilizing computational and storage resources to achieve higher performance and energy efficiency. The solution proposed in this invention is the first of its kind and is an effective one. Attached Figure Description

[0036] The present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments, and the advantages of the present invention in the above and / or other aspects will become clearer.

[0037] Figure 1 This is a diagram of the Softmax training module structure.

[0038] Figure 2 This is a schematic diagram of a parallelization design with a parallelism of 4 in the traditional scheme.

[0039] Figure 3 This is a schematic diagram of the forward propagation calculation for serialized processing.

[0040] Figure 4 This is a schematic diagram of the forward propagation calculation in pipeline design.

[0041] Figure 5 This is a flowchart of the Softmax calculation process for each vector in the M matrix.

[0042] Figure 6 This is a schematic diagram of the backpropagation path of the Softmax training module.

[0043] Figure 7 This is a schematic diagram of the data flow during the first clock cycle of the dm1 calculation process.

[0044] Figure 8 This is a schematic diagram of the data flow during the t-th clock cycle of the dm1 calculation process. Detailed Implementation

[0045] This invention provides an apparatus for training the Softmax function in a large language model. Figure 1 This diagram illustrates the structure of the Softmax module, which can be used for computation at each training stage. It includes basic arithmetic units and a corresponding storage module. Two parts of Random Access Memory (RAM) are used to cache data required for computation or intermediate results. RAM1 and RAM2 store different data for different training stages and are shared by all stages. The upper part of the module represents the forward propagation path, and the lower part represents the backward propagation path.

[0046] The forward propagation path contains e x The computation modules include an exponential function unit (exp), adders, and dividers. Registers are inserted between these computation modules to implement pipelined processing, thereby improving hardware frequency and throughput.

[0047] The backpropagation path includes two multipliers, one adder, and a multiplexer (MUX) for reconstructing the data path. During different computation cycles in the backpropagation phase, the multiplexer can modify the data flow within the module to achieve accurate computation. The supported data formats of the computation module can be adjusted as needed. For example, if the data is floating-point, the exp unit can be implemented using a corresponding floating-point arithmetic IP core; if the data is fixed-point, a look-up table or a corresponding mathematical approximation unit can be used.

[0048] Streamlined forward propagation path:

[0049] During the forward propagation phase, the Softmax function is calculated row-wise on the M matrix. A row vector in the matrix M is represented as m = (m1, m2, ..., m...). s The result vector obtained by Softmax is n = (n1, n2, ..., nn) s ), the elements n in n i The calculation process is expressed as follows:

[0050]

[0051] Its calculation can be divided into three stages: the calculation of each element. calculate, Summation and division.

[0052] In traditional neural network accelerators, parallelization is frequently employed in hardware design related to Softmax inference to improve throughput. For example, a traditional design with a parallelism of 4... Figure 2 As shown: Simultaneous computation is achieved by using multiple parallel modules. The summation is achieved using an addition tree, followed by calculation of each result using multiple dividers. The advantage of parallel processing is that throughput can be continuously improved by increasing the degree of parallelism. However, in Transformer-type models, the matrix dimension in the self-attention mechanism is uncertain because the length *s* of the sample sequence processed by the model is constantly changing. Therefore, in a fixed hardware design, the degree of parallelism is a difficult quantity to determine. A high degree of parallelism means that when the sample sequence length *s* is small, the computation cannot fully utilize hardware resources; a low degree of parallelism means that when the sample sequence length *s* is large, the hardware cannot correctly complete the required computation.

[0053] Considering the variable characteristics of Softmax computation in Transformer-type models, this invention uses a serialization and pipeline design method to adapt to the computational characteristics of the model. At the same time, this invention still has the space for parallelization design to further improve the computational throughput of hardware modules. Parallel computation by multiple modules can utilize the parallelism provided by s×s matrices, multi-head attention, and multi-batch training.

[0054] Serialization process:

[0055] like Figure 3 As shown, the forward propagation path uses serialization to complete the forward propagation calculation. The input data are the elements m1, m2, ..., mn in the vector m. s The values ​​are stored in RAM1. After the calculation begins, m1, m2, ..., m... s The elements are retrieved sequentially from RAM1, first undergoing exponentiation in the exp unit, and the exponentiation of each element is calculated. x Calculation results The accumulation is performed in the adder, while e x The calculation result is stored in RAM2. After all elements in vector m have undergone exponentiation, the adder accumulates the results. It was stored in RAM2. Then, The data is transmitted to the divisor input port of the divider. The values ​​are retrieved from RAM2 and sequentially transferred to the dividend input port of the divider. Therefore, the outputs of the divider are n1, n2, ..., n in sequence. s .

[0056] Pipeline design: In the serialized processing described above, the division calculation requires waiting for all elements in vector m to complete e. x The calculation can only begin after the exponential operation and the accumulated result are obtained, and the waiting time between each stage reduces the computational efficiency. Therefore, in this invention, the calculation is divided into two stages: (1) e x (1) Exponentiation and accumulation; (2) Division. And a pipeline is inserted between the two stages to improve computational efficiency and throughput.

[0057] like Figure 4 As shown, the row vectors in matrix M are denoted as m1, m2, ..., m... s , where m1=(m 11 m 12 , ..., m 1s ), m2=(m 21 m 22 , ..., m 2sAfter m1 completes the exponentiation and accumulation operations and performs the division operation, m2 begins the new exponentiation and accumulation operations without requiring m1 to wait for the calculation to complete. The two stages of calculation can be performed in parallel, without any waiting time between hardware modules, effectively improving hardware utilization and computational throughput. The calculation flowchart for each vector in matrix M is shown below. Figure 5 As shown.

[0058] Reconfigurable backpropagation path:

[0059] During the backpropagation phase, the gradient propagated from the subsequent network structure of Softmax is dn, and dm can be calculated using the following formula:

[0060] dm = dn·(diag(n)-n) T ·n)

[0061] With elements in dm i Taking (1≤i≤s) as an example, the calculation process can be expanded as follows:

[0062]

[0063] For example:

[0064]

[0065] The structure of the backpropagation training path of the Softmax module is as follows: Figure 6 As shown, when used for backpropagation calculations, RAM1 stores the forward propagation calculation results n1, n2, ..., n. s RAM2 stores the gradient values ​​dn1, dn2, ..., dn passed from subsequent Softmax network structures. s In the process of calculating the gradient dm of vector m, the backpropagation path reconstructs the internal data flow through a multiplexer, thereby accurately completing the calculation of each element in dm.

[0066] Taking the calculation process of dm1 as an example, in the first clock cycle of the calculation, the connection method and data flow of the back propagation path are as follows: Figure 7 As shown:

[0067] (1) In the figure, the black solid line represents the valid path controlled by the multiplexer, and the dashed line represents the path that is not selected.

[0068] (2) At this time, the data read in RAM1 is n1, and the data read in RAM2 is dn1.

[0069] (3) The inputs of the left multiplier B1 are n1 and -n1, and the output is -n1·n1.

[0070] (4) The inputs of the multiplier B2 on the right are n1 and dn1, and the output of B2 is n1·dn1.

[0071] (5) The left input of adder A1 is n1·dn1, and the top input of A1 is 0, indicating the start of accumulation. The output of A1 in the current cycle is n1·dn1.

[0072] In calculating the t-th clock cycle of dm1 (2≤t≤s+1), the connection method and data flow of the back propagation path are as follows: Figure 8 As shown, it has the following characteristics:

[0073] (1) At this time, the data read from RAM1 is n t The data read from RAM2 is dn t-1 .

[0074] (2) The left input of multiplier B1 is kept at n1 by the control of the multiplexer, and the top input of B1 is -n. t The output of B1 in the current cycle is -n1·n t .

[0075] (3) The upper input of multiplier B2 is dn t-1 At this point, the input to the left of B2 is the output of B1 in the previous cycle - n1·n t-1 Therefore, the output of B2 is -n1·n t-1 ·dn t-1 .

[0076] (4) The left input of adder A1 is the output of B2, i.e., -n1·n t-1 ·dn t-1 The upper input of A1 is controlled by a multiplexer to become the output of the previous cycle of A1, at which point A1 functions as an accumulator.

[0077] After s+1 cycles, the accumulated output of A1 is

[0078] The calculation process for the remaining elements in dm is the same as that for dm1.

[0079] Example

[0080] Taking the training process of a Transformer-type model in a natural language processing task as an example, this paper illustrates the computational efficiency improvement of the forward propagation path in a serialized pipelined design compared to a traditional parallel design. The sample length of the natural language processing dataset is set to be between 1 and 128, meaning the Softmax function needs to process vectors with a length of 1 ≤ s ≤ 128. In traditional parallel designs, to correctly perform the Softmax function calculation, a hardware design with a parallelism of 128 is required (e.g., ...). Figure 2 Assuming the mean length s of all samples in the dataset is μ, then the hardware utilization of the parallelization design scheme during the entire training process is... In the forward propagation path of the serialized pipelined design proposed in this invention, the utilization rate of hardware resources during the entire training process is... Comparing hardware utilization reveals that η1 ≥ η2 only when μ ≥ 127. In actual training, the mean of the sample length s (1 ≤ s ≤ 128) is much smaller than μ, therefore η2 is greater than η1. Thus, the forward propagation path of the serialized pipelined design proposed in this invention exhibits better hardware resource utilization and computational efficiency than traditional parallelized design schemes.

[0081] This invention provides an apparatus for training the Softmax function in a large language model. Many methods and approaches exist for implementing this technical solution; the above description is merely a preferred embodiment of the invention. It should be noted that those skilled in the art can make various improvements and modifications without departing from the principles of this invention, and these improvements and modifications should also be considered within the scope of protection of this invention. All components not explicitly stated in this embodiment can be implemented using existing technologies.

Claims

1. An apparatus for training the Softmax function in a large language model, wherein the large language model is a Transformer-type model, characterized in that, The upper part of the device is the forward propagation path, and the lower part is the reverse propagation path; The forward propagation path includes e x The exponential function unit, adder, and divider, in e x Registers were inserted between the exponential function unit, the adder, and the divider; The backpropagation path includes two multipliers, one adder A1, and a multiplexer MUX for reconstructing the data path; the two multipliers are the left multiplier B1 and the right multiplier B2. The multiplexer MUX is used to change the data flow within the device; The forward propagation path and the backward propagation path share two random access memories, RAM1 and RAM2. The forward propagation path uses a serialization process to perform the forward propagation calculation: the input data is the elements m1, m2, ..., m in the vector m. s Input data is stored in RAM1; after the calculation begins, m1, m2, ..., m s Retrieve sequentially from RAM1, first passing through e x The exponential function unit performs exponential operations, with each element m1, m2, ..., m s e x Calculation results The accumulation is performed in the adder, while e x The calculation results are stored in RAM2; After all elements in vector m have undergone exponentiation, the adder accumulates the result. It was stored in RAM2; then, The data is transmitted to the divisor input port of the divider. The values ​​are retrieved from RAM2 and sequentially transferred to the dividend input port of the divider. The outputs of the divider are n1, n2, ..., n in sequence. s ; During serialization, the division calculation needs to wait until all elements in vector m have completed e. x The division calculation can only begin after the exponentiation operation and the accumulated result are obtained, dividing the calculation into two stages: e x The exponentiation and accumulation phase and the division phase are performed in a pipelined manner. Set the input matrix during the forward propagation of the Softmax function. Let M be a real matrix with s rows and s columns, and let the row vectors in matrix M be represented as m1, m2, ..., m. s , where m s Let m represent the vector in the s-th row of matrix M. s =(m s1 ,m s2 ,…,m ss ), m ss Represents vector m s The s-th element in the vector m has an index indicating its position in the vector. s The length is s; after m1 completes the exponentiation and accumulation operations and performs the division operation, m2 starts the new exponentiation and accumulation operations without m1 having to wait for the calculation to complete; the calculations in the two stages are performed in parallel.

2. The apparatus according to claim 1, characterized in that, During the forward propagation phase, the Softmax function is calculated row-wise on the M matrix, where a row vector m = (m1, m2, ..., mn) in the matrix M is a vector. s ), m s This represents the s-th element in vector m, where the subscript indicates the element's position within the vector. The length of vector m is s. The resulting vector obtained by Softmax is n = (n1, n2, ..., nn...). s The i-th element n in n i The calculation process is expressed as follows: Where i takes values ​​from 1 to s.

3. The apparatus according to claim 2, characterized in that, The calculation process of formula (1) includes three stages: the calculation of each element. calculate, Summation and division.

4. The apparatus according to claim 3, characterized in that, During the backpropagation phase of the backpropagation path, the gradient propagated from the subsequent network structure of Softmax is dn. The subsequent network structure refers to the calculation after the Softmax function in the Transformer-type model. dn represents the gradient of the output data n of the Softmax function in the forward propagation phase relative to the final output result of the model. During the backpropagation phase, it is necessary to calculate the gradient dm of vector m relative to the final output of the model.

5. The apparatus according to claim 4, characterized in that, dm is calculated using the following formula: dm=dn·(diag(n)-n T ·n) Where diag represents the diagonal matrix function, T represents the matrix transpose; and dn is an element in dm. i The calculation process can be broken down as follows: Where 1≤i≤s.

6. The apparatus according to claim 5, characterized in that, When the backpropagation path of the device is used for backpropagation calculation, RAM1 stores the calculation results n1, n2, ..., n from the forward propagation. s RAM2 stores the gradient values ​​dn1, dn2, ..., dn passed from the subsequent Softmax network structure. s dn s Let represent the s-th element in vector dn, where the subscript indicates the position of the element in the vector, and the length of vector dn is s.

7. The apparatus according to claim 6, characterized in that, In the process of calculating the gradient dm of vector m, the backpropagation path reconstructs the internal data flow through a multiplexer, thereby completing the calculation of each element in dm. i The calculation process is as follows: When calculating dm i During the first clock cycle of the calculation, the connection method of the backpropagation path has the following characteristics: At this moment, the data read from RAM1 is n. i The data read from RAM2 is dn i ; The input to the left-hand multiplier B1 is n. i and -n i The output is -n i ·n i ; The input to the multiplier b2 on the right is n. i and dn i The output of B2 is n i ·dn i ; The left input of adder A1 is n i ·dn i The input above A1 is 0, indicating the start of accumulation. The output of A1 in the current cycle is n. i ·dn i ; In calculating the t-th clock cycle of dm1, where 2≤t≤s+1, the connection method of the back propagation path has the following characteristics: At this moment, the data read from RAM1 is n. t n t This represents the t-th element in vector n, where the index indicates the element's position within the vector. All elements of vector n are stored in RAM1; the data read from RAM2 is dn. t-1 ;dn t-1 Let represent the (t-1)th element in vector dn, and let the subscript indicate the position of the element in the vector. All elements of vector dn are stored in RAM2. The left input of the left multiplier B1 is kept at n by the control of the multiplexer MUX. i The input above B1 is -n t The output of B1 in the current cycle is -n i ·n t ; The input above the multiplier B2 on the right is dn. t-1 At this point, the input to the left of B2 is the output of B1 in the previous cycle - n. i ·n t-1 The output of B2 is -n i ·n t-1 ·dn t-1 ; The left input of adder A1 is the output of B2 - n i ·n t-1 ·dn t-1 The upper input of A1 is controlled by the multiplexer MUX to be the output of A1 in the previous cycle. At this time, A1 has the function of an accumulator. After s+1 cycles, the accumulated output of A1 is At this point, dm was calculated. i .

8. The apparatus according to claim 7, characterized in that, The calculation is performed in parallel using two or more devices.

Images