Attention mechanism computing and model inference methods, apparatuses, devices, and media

Abstract

The application provides an attention mechanism calculation and model inference method, device, equipment and medium, wherein the method comprises the following steps: dividing a sequence into multiple segments, respectively assigning processing units to the segments, and controlling the processing units to calculate Q vectors, K vectors and V vectors of the segments; controlling the processing units to respectively perform QKV calculation to obtain segment calculation results of the processing units, wherein the QKV calculation involves the Q vectors, the K vectors and the V vectors calculated by the processing units themselves and the Q vectors calculated by other processing units; and controlling the processing units to respectively perform accumulation calculation to obtain attention calculation results of the sequence at the segments corresponding to the processing units, wherein the accumulation calculation involves the segment calculation results calculated by the processing units themselves and the segment calculation results calculated by other processing units. The method, device, equipment and medium provided by the application relieve the memory pressure, avoid the waste of storage resources, and improve the flexibility of subsequent parallel calculation.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence technology, and in particular to a method, apparatus, device, and medium for attention mechanism computation and model reasoning. Background Technology

[0002] As a core component of the Transformer model, the attention mechanism can effectively capture contextual information in sequence data and dynamically assign weights to different positions, thereby improving the model's ability to process sequences.

[0003] However, as the model size and sequence length increase, the model faces challenges in computational complexity and memory limitations when handling long sequence inference tasks. When the model parameters and sequence length reach a certain level, multiple graphics cards or even multiple computing nodes are often required to perform model inference together.

[0004] How to perform attention mechanism calculations in a multi-GPU environment remains a pressing issue in the field of artificial intelligence. Summary of the Invention

[0005] This invention provides a method, apparatus, device, and medium for attention mechanism calculation and model inference, to solve the problem of attention mechanism calculation in environments with multiple graphics cards.

[0006] This invention provides a method for calculating attention mechanisms, comprising:

[0007] The sequence is divided into multiple segments, and processing units are assigned to each segment. Each processing unit is controlled to calculate the query Q vector, key K vector, and numerical V vector of the assigned segment.

[0008] Each of the processing units is controlled to perform QKV calculations to obtain the fragment calculation results of each processing unit. The QKV calculations include the Q vector, K vector, and V vector calculated by the processing unit itself, as well as the Q vector calculated and transmitted by other processing units.

[0009] Each of the processing units is controlled to perform cumulative calculations, and the attention calculation results of the sequence are obtained at the segment corresponding to each processing unit. The cumulative calculations include the segment calculation results calculated by the processing unit itself and the segment calculation results calculated and transmitted by other processing units.

[0010] The attention mechanism calculation method provided by the present invention further includes:

[0011] While controlling each of the processing units to perform QKV calculations, at least one of Q-vector transmission and fragment calculation result transmission is performed between the processing units.

[0012] According to an attention mechanism calculation method provided by the present invention, the step of simultaneously controlling each of the processing units to perform QKV calculations and performing at least one of Q-vector transmission and segment calculation result transmission between the processing units includes:

[0013] While controlling each of the processing units to perform QK calculations in the current round of QKV calculations, the system also controls each of the processing units to transmit the fragment calculation results obtained in the previous round of the current round.

[0014] According to an attention mechanism calculation method provided by the present invention, the step of simultaneously controlling each of the processing units to perform QKV calculations and performing at least one of Q-vector transmission and segment calculation result transmission between the processing units includes:

[0015] While controlling each of the aforementioned processing units to perform SV calculation in QKV calculation, the system also controls each of the aforementioned processing units to perform Q vector transmission.

[0016] The SV calculation is performed on the result S of the QK calculation and the V vector; in one round of the QKV calculation, the Q vector used in the QK calculation is the same as the Q vector transmitted by the Q vector.

[0017] The attention mechanism calculation method provided by the present invention further includes:

[0018] Each of the aforementioned processing units is controlled to perform subsequent processing on the attention calculation results based on its own stored segment processing weights;

[0019] The segment processing weight is obtained by dividing the segment into subsequent processing weights required for subsequent processing of the sequence based on the segments allocated by the processing unit.

[0020] According to the attention mechanism calculation method provided by the present invention, the step of dividing the sequence into multiple segments and allocating processing units to each segment includes:

[0021] The original sequence is divided into multiple sequences, and a set of processing units is assigned to each sequence.

[0022] The sequence is divided into multiple segments, and each segment is assigned a processing unit from the set of processing units.

[0023] The present invention also provides a model reasoning method, comprising:

[0024] Obtain the target sequence;

[0025] The target sequence is input into the model for inference, and the inference result output by the model is obtained;

[0026] The model includes an attention mechanism, and during the inference process of the model, the calculation of the attention mechanism is performed based on the attention mechanism calculation method described above.

[0027] The present invention also provides an attention mechanism computing device, comprising:

[0028] The segmentation module is used to divide the sequence into multiple segments, allocate processing units to each segment, and control each processing unit to calculate the query Q vector, key K vector, and numerical V vector of the allocated segment.

[0029] The first calculation module is used to control each of the processing units to perform QKV calculations respectively, and to obtain the fragment calculation results of each processing unit. The QKV calculations include the Q vector, K vector and V vector calculated by the processing unit itself, and the Q vector calculated and transmitted by other processing units.

[0030] The second calculation module is used to control each of the processing units to perform cumulative calculations, and to obtain the attention calculation results of the sequence at the segment corresponding to each processing unit. The cumulative calculations include the segment calculation results calculated by the processing unit itself and the segment calculation results calculated and transmitted by other processing units.

[0031] The present invention also provides a model reasoning device, comprising:

[0032] The acquisition module is used to acquire the target sequence;

[0033] The reasoning module is used to input the target sequence into the model for reasoning and obtain the reasoning result output by the model.

[0034] The model includes an attention mechanism, and during the inference process of the model, the calculation of the attention mechanism is performed based on the attention mechanism calculation method described above.

[0035] The present invention also provides an electronic device, including a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor executes the program to implement the attention mechanism calculation method or model inference method as described above.

[0036] The present invention also provides a non-transitory computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the attention mechanism computation method or model inference method as described above.

[0037] The present invention also provides a computer program product, including a computer program that, when executed by a processor, implements the attention mechanism calculation method or model reasoning method as described above.

[0038] The attention mechanism calculation and model inference method, apparatus, device, and medium provided by this invention perform QKV calculation involving Q vectors calculated by the processing unit itself and transmitted from other processing units at each processing unit, and perform cumulative calculation of segment calculation results calculated by the processing unit itself and transmitted from other processing units at each processing unit. Thus, in an environment where attention mechanism calculation is performed by multiple processing units, each processing unit can obtain the attention calculation result of the complete sequence. Therefore, each processing unit only needs to store the weights for each segment in the subsequent operations of attention calculation, which reduces the storage pressure on the processing unit, avoids the waste of storage resources, and improves the flexibility of subsequent parallel calculations by multiple processing units. Attached Figure Description

[0039] To more clearly illustrate the technical solutions in this invention or related technologies, the accompanying drawings used in the description of the embodiments or related technologies will be briefly introduced below. Obviously, the accompanying drawings described below are some embodiments of this invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0040] Figure 1 This is one of the flowcharts illustrating the attention mechanism computation method provided by related technologies;

[0041] Figure 2 This is the second flowchart illustrating the attention mechanism calculation method provided by related technologies;

[0042] Figure 3 This is one of the flowcharts illustrating the attention mechanism method provided by the present invention;

[0043] Figure 4 This is the second flowchart of the attention mechanism method provided by the present invention;

[0044] Figure 5 This is a flowchart illustrating the model reasoning method provided by the present invention;

[0045] Figure 6 This is a schematic diagram of the attention calculation device provided by the present invention;

[0046] Figure 7 This is a schematic diagram of the structure of the model reasoning device provided by the present invention;

[0047] Figure 8 This is a schematic diagram of the structure of the electronic device provided by the present invention. Detailed Implementation

[0048] To make the objectives, technical solutions, and advantages of this invention clearer, the technical solutions of this invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some, not all, of the embodiments of this invention. All other embodiments obtained by those skilled in the art based on the embodiments of this invention without creative effort are within the scope of protection of this invention.

[0049] With the continuous development of deep learning technology, the Transformer model has achieved remarkable results in fields such as natural language processing and speech recognition. One of the core components of the Transformer model is the attention mechanism, which can effectively capture contextual information in a sequence and dynamically assign weights to different positions in the sequence, thereby improving the model's ability to process sequences.

[0050] The attention mechanism computation mainly consists of three parts: Query (Q), Key (K), and Value (V) vectors. During the attention mechanism computation, the input sequence first undergoes a linear transformation, converting it into the aforementioned Q, K, and V vectors. Then, the dot product of the Q vector and all K vectors is calculated, and the result is normalized using softmax. Finally, the Q vector and the V vector are weighted and summed to obtain the output, which is the result of the attention mechanism computation. The execution of the attention mechanism allows the Transformer model to flexibly focus on important information when processing sequences, thereby improving the model's performance.

[0051] However, as the model size and sequence length increase, the Transformer model faces challenges in computational complexity and memory limitations when handling long sequence inference. When the model parameters and sequence length reach a certain level, a single graphics card (GPU) or a single computing node often cannot handle such a large amount of computation and storage requirements, thus necessitating the use of multiple graphics cards or even multiple computing nodes to process the data.

[0052] In related technologies, an implementation method for attention mechanisms in multi-GPU environments (i.e., multi-card environments) is provided. Considering that the QKV calculations in the attention mechanism do not have a strict order, as long as all sequence calculations are completed, the sequence can be segmented. During calculation, each card first processes its own portion of the sequence's QKV, and then, while calculating the QKV, the calculated KV sequences are ring-transmitted between the cards. Through multiple transmissions, each card can obtain the KV sequences calculated by all cards, thus obtaining the complete attention calculation result for its respective portion of the sequence.

[0053] Figure 1 , Figure 2 This is a flowchart illustrating the attention mechanism computation method provided by related technologies, such as... Figure 1 As shown, devices 1 and 2 represent different graphics cards. During the attention mechanism calculation process, each device puts its calculated query block into the attention mechanism calculation. In addition, key value blocks (K-V blocks) are transmitted between devices, so that the KV sequences calculated by each device can be ring-transmitted between devices. Figure 2 In this scenario, a sequence is divided into four segments. Each card calculates Q, K, and V for its respective segment. Specifically, for segment 1, we obtain query Q1, key K1, and value V1; for segment 2, we obtain query Q2, key K2, and value V2; for segment 3, we obtain query Q3, key K3, and value V3; and for segment 4, we obtain query Q4, key K4, and value V4. Taking the attention calculation for segment 1 as an example, we can calculate the query Q1 with the key K1 and value V1 calculated by card 1 itself. Then, we can calculate the query Q1 with the key K2 and value Q2, key K3 and value Q3, and key K4 and value Q4 calculated by other cards. This gives us the attention calculation result for segment 1.

[0054] After the above method is executed, each card only stores the attention calculation results of a portion of the sequence, thereby reducing storage requirements. Furthermore, since computation and data transmission can be performed in parallel, the data transmission time can be effectively masked, which also helps to improve the efficiency of attention calculation.

[0055] However, in subsequent calculations after the attention mechanism is completed, each card needs to store the complete FFN (Feedforward Neural Network) weights, that is, in Figure 1 , Figure 2 Each blockwise feedforward requires the application of the full FFN weights. This can lead to memory waste in large models, hindering efficient parallelization of subsequent computations.

[0056] Even after the above calculations are completed, integrating the attention calculation results of the partial sequences on each card by grouping the calculation cards still requires an additional transmission of calculation results, which will also affect the inference efficiency of the model.

[0057] To address the above problems, embodiments of the present invention provide an attention mechanism calculation method. Figure 3 This is one of the flowcharts illustrating the attention mechanism method provided by the present invention, such as... Figure 3As shown, the method includes:

[0058] Step 310: Divide the sequence into multiple segments, assign processing units to each segment, and control each processing unit to calculate the query Q vector, key K vector, and numerical V vector of the assigned segment.

[0059] The sequence referred to here is the sequence for which attention mechanism calculations are required. Specifically, the sequence can be a text sequence, an image sequence, a speech sequence, etc., but this embodiment of the invention does not impose specific limitations on it.

[0060] After obtaining the sequence, it can be split into multiple segments. Furthermore, the splitting of the sequence can be based on the number of available processing units. These processing units can be GPUs or other types of units that can be used for attention mechanism computation; this embodiment of the invention does not specifically limit this. For example, with four available processing units, the sequence can be divided into four parts, resulting in four segments. It is understood that each segment obtained is a portion of the sequence.

[0061] After the sequence is split, the resulting segments can be assigned to processing units. Here, the allocation rule is to assign one segment to each processing unit. Thus, by running multiple processing units, the attention mechanism can be computed on the sequence composed of multiple segments.

[0062] After assigning segments to be processed to each processing unit, each processing unit can be controlled to calculate the query vector (Q), key vector (K), and value vector (V) for its assigned segment. This yields the Q, K, and V vectors for each segment at each processing unit. For example, in the case of four processing units, processing unit 1 performs vector calculations on segment 1, obtaining the Q1, K1, and V1 vectors for segment 1; processing unit 2 performs vector calculations on segment 2, obtaining the Q2, K2, and V2 vectors for segment 2; processing unit 3 performs vector calculations on segment 3, obtaining the Q3, K3, and V3 vectors for segment 3; and processing unit 4 performs vector calculations on segment 4, obtaining the Q4, K4, and V4 vectors for segment 4.

[0063] Step 320: Control each of the processing units to perform QKV calculations to obtain the fragment calculation results of each processing unit. The QKV calculations include the Q vector, K vector, and V vector calculated by the processing unit itself, as well as the Q vector calculated and transmitted by other processing units.

[0064] Specifically, after each processing unit has calculated the Q vector, K vector, and V vector of the segment allocated to it, it can be controlled to perform QKV calculations.

[0065] The QKV calculation here assumes that each processing unit not only possesses its own calculated Q vector, K vector, and V vector, but also receives Q vectors calculated by other processing units. In other words, processing units can transmit Q vectors to each other; that is, each processing unit can send its own calculated Q vector to other processing units and can receive Q vectors calculated by other processing units.

[0066] Therefore, at any processing unit, the QKV calculation involves the Q vectors calculated by all processing units, as well as the K and V vectors calculated by that processing unit itself. Based on this, the fragment calculation result obtained by performing QKV calculation at any processing unit is actually the result of performing QKV calculation on the complete Q vector of the sequence and the K and V vectors of the fragment assigned by that processing unit. It is an intermediate result obtained during the attention mechanism calculation process for the sequence.

[0067] Step 330: Control each processing unit to perform cumulative calculations, and obtain the attention calculation results of the sequence at the segment corresponding to each processing unit. The cumulative calculations include the segment calculation results calculated by the processing unit itself and the segment calculation results calculated and transmitted by other processing units.

[0068] Specifically, each processing unit can also transmit the calculated fragment results to each other. That is, each processing unit can send the calculated fragment results to other processing units and receive the calculated fragment results from other processing units.

[0069] In this way, each processing unit can contain the fragment calculation results obtained from all processing units.

[0070] Therefore, each processing unit can be controlled to perform a reduction calculation on the segment processing results calculated by all processing units, so that each processing unit can obtain the cumulative result of the calculation results of all segments, that is, the attention calculation result of the sequence.

[0071] Compared to related technologies that transmit K and V vectors at each processing unit, the method provided in this embodiment of the invention transmits Q vectors and fragment calculation results at each processing unit. That is, the amount of transmission between processing units is comparable to that of related technologies.

[0072] However, in related technologies, after the attention mechanism calculation is completed, each processing unit obtains an incomplete attention calculation result. That is, each processing unit can only obtain the attention calculation result for its assigned segment. Therefore, complete FFN weights need to be configured at each processing unit to ensure the execution of subsequent calculations. In the method provided in this embodiment of the invention, after the attention mechanism calculation is completed, each processing unit assigned to each segment obtains the complete attention calculation result for the entire sequence. Therefore, each processing unit only needs to store the FFN weights of each segment to execute subsequent calculations.

[0073] Therefore, compared to related technologies that require storing complete FFN weights at each processing unit, the attention mechanism calculation method provided by this invention can obtain the complete attention calculation result of the entire sequence at each processing unit under the premise of equivalent transmission volume, that is, without increasing the additional transmission volume. Thus, each processing unit does not need to store the FFN weights of each segment, which greatly reduces the memory pressure of the processing unit and improves the flexibility of subsequent parallel data calculation.

[0074] The method provided in this invention performs QKV calculations involving Q vectors calculated by the processing unit itself and those transmitted from other processing units at each processing unit, and performs summation calculations of segment calculation results calculated by the processing unit itself and those transmitted from other processing units at each processing unit. As a result, in an environment where attention mechanism calculations are performed by multiple processing units, each processing unit can obtain the attention calculation results of the complete sequence. Thus, each processing unit only needs to store the weights for each segment in subsequent attention calculation operations, reducing the storage pressure on the processing units, avoiding the waste of storage resources, and improving the flexibility of subsequent parallel calculations by multiple processing units.

[0075] Based on the above embodiments, the method further includes:

[0076] While controlling each processing unit to perform QKV calculations, at least one of the following is performed between the processing units: Q vector transmission and fragment calculation result transmission.

[0077] Specifically, in order to enable each processing unit to apply the Q vector calculated by other processing units to the QKV calculation when each processing unit performs QKV calculation separately, Q vector transmission is required between processing units.

[0078] In this process, Q-vector transmission between processing units can be performed synchronously with QKV calculations performed by each processing unit. That is, while each processing unit is performing QKV calculations, the Q-vectors calculated by each processing unit can be transmitted. The parallel execution of QKV calculations and Q-vector transmission can mask the time spent on Q-vector transmission, thereby improving the computational efficiency of the attention mechanism.

[0079] It is understandable that Q-vector transmission can be performed simultaneously with QKV computation, which can be achieved through circular transmission of Q-vectors among multiple processing units. That is, while each processing unit performs QKV computation using its own calculated Q-vector, it can simultaneously send its calculated Q-vector to the next processing unit and receive the Q-vector calculated by the previous processing unit. Thus, after each processing unit completes its QKV computation using its own calculated Q-vector, it can then use the received Q-vector from the previous processing unit to perform its own QKV computation, while simultaneously sending the Q-vector calculated by the previous processing unit to the next processing unit and receiving the Q-vector sent by the previous processing unit. In this way, each QKV computation uses the Q-vector received in the previous computation, and each time a QKV computation is performed, the Q-vector used in the current computation is sent to other processing nodes, thereby achieving synchronization between QKV computation and Q-vector transmission.

[0080] In addition, in order for each processing unit to apply the fragment calculation results obtained by other processing units to the accumulation calculation when each processing unit performs the accumulation calculation separately, it is necessary to transmit the fragment calculation results between the processing units.

[0081] In this process, the transmission of fragment computation results between processing units can be performed synchronously with the QKV computation performed by each processing unit. That is, while each processing unit is performing QKV computation, the fragment computation results obtained by each processing unit can be transmitted. The parallel execution of QKV computation and fragment computation results can mask the time of fragment computation result transmission, thereby improving the computational efficiency of the attention mechanism.

[0082] It is understandable that transmitting fragment calculation results simultaneously with QKV calculation can be achieved through circular transmission of fragment calculation results among multiple processing units. That is, while each processing unit performs QKV calculation using a set of Q vectors, it can send the fragment calculation results obtained from the previous QKV calculation to the next processing unit, and simultaneously receive the vectors obtained from the previous QKV calculation from the previous processing unit. Therefore, after each set of Q vectors completes its QKV calculation and the next set of Q vectors begins its QKV calculation, circular transmission of previously obtained fragment calculation results can be achieved between processing units, thus realizing synchronization between QKV calculation and fragment calculation result transmission.

[0083] Based on any of the above embodiments, step 320, in which each processing unit is controlled to perform QKV calculation, may include the following two steps:

[0084] Each processing unit is controlled to perform QK calculations to obtain the calculation result S;

[0085] Each processing unit is controlled to perform SV calculations to obtain the fragment processing results.

[0086] Here, the calculation result S is the result of QK calculation, which is the intermediate calculation result for obtaining the fragment processing result. S can be specifically expressed in the following form:

[0087] S = softmax(QK) T )

[0088] After obtaining the calculation result S, the SV calculation involves performing a dot product between S and V. In this embodiment of the invention, the dot product result of SV is denoted as the fragment processing result.

[0089] Based on any of the above embodiments, the step of simultaneously controlling each processing unit to perform QKV calculations and performing at least one of Q-vector transmission and fragment calculation result transmission between the processing units includes:

[0090] While controlling each processing unit to perform QK calculations for the current round of QKV calculations, the system also controls each processing unit to transmit the fragment calculation results obtained from the previous round of the current round.

[0091] Specifically, after dividing the QKV computation into two steps, the transmission of fragment computation results can be performed synchronously with the QK computation in the QKV computation.

[0092] It is understandable that during the process of controlling each processing unit to perform QKV calculation, each processing unit can perform QKV calculation in multiple rounds, and in each round, a set of Q vectors participates in the QKV calculation.

[0093] For a QKV calculation in one round, that is, in the QKV calculation of the current round, when each processing unit performs QK calculation, the fragment calculation results calculated by each processing unit in the previous round before the current round have been obtained. Therefore, the fragment calculation results obtained by each processing unit in the previous round before the current round can be transmitted synchronously.

[0094] It should be noted that, for any processing unit, the segment calculation result transmitted at this time includes the segment calculation result it obtained in the previous round, as well as the segment processing result it received from other processing units in the previous round.

[0095] Therefore, in multi-round QKV calculations, the fragment calculation results obtained by each processing unit in each round can be transmitted to other processing units. Thus, when each processing unit completes the QKV calculation, each processing unit has the fragment calculation results of all processing units.

[0096] Based on any of the above embodiments, the step of simultaneously controlling each processing unit to perform QKV calculations and performing at least one of Q-vector transmission and fragment calculation result transmission between the processing units includes:

[0097] While controlling each processing unit to perform SV calculation in QKV calculation, the processing unit is also controlled to perform Q vector transmission.

[0098] The SV calculation is performed on the result S of the QK calculation and the V vector; in one round of the QKV calculation, the Q vector used in the QK calculation is the same as the Q vector transmitted by the Q vector.

[0099] Specifically, after dividing the QKV computation into two steps, the Q vector transfer can be performed synchronously with the SV computation in the QKV computation.

[0100] It is understandable that during the process of controlling each processing unit to perform QKV calculation, each processing unit can perform QKV calculation in multiple rounds, and in each round, a set of Q vectors participates in the QKV calculation.

[0101] For a QKV calculation in one round, that is, in the current round of QKV calculation, after each processing unit performs QK calculation, the calculation result S of the current round has been obtained, and the Q vector used for the QKV calculation in the current round has been used up. Therefore, when performing SV calculation, the Q vector used for QK calculation in the current round can be transmitted so that other processing units can use the Q vector to perform QKV calculation.

[0102] Thus, in multi-round QKV calculations, Q vectors are transmitted per processing unit, so that each processing unit can obtain all Q vectors before the end of the multi-round QKV calculations, thereby realizing QKV calculations for all Q vectors and the K and V vectors of the segments allocated to the processing units.

[0103] Based on any of the above embodiments, the method further includes:

[0104] Each processing unit is controlled to perform subsequent processing on the attention calculation results based on its own stored segment processing weights;

[0105] The segment processing weight is obtained by dividing the segment into subsequent processing weights required for subsequent processing of the sequence based on the segments allocated by the processing unit.

[0106] Specifically, after step 330 is completed, each processing unit obtains the complete attention calculation result for the entire sequence. Therefore, when each processing unit performs attention calculation based on the attention calculation result, only the segment processing weights corresponding to the segments allocated by each processing unit are needed, thereby enabling subsequent processing of the segments allocated by each processing unit.

[0107] Here, for each processing unit, the segment processing weight is obtained by segmenting the sequence according to the subsequent processing weights required for subsequent processing. The segment processing weights are used to implement subsequent processing for a single segment in the sequence.

[0108] It is understandable that, compared to related technologies where each processing unit can only obtain the attention calculation results of its assigned segment, so it is necessary to configure candidate processing weights for the complete sequence at each processing unit, in the method provided by the embodiments of the present invention, since each processing unit can obtain the attention calculation results of the complete sequence, it is only necessary to configure the segment processing weights of the segment at each processing unit to realize subsequent calculations, thereby reducing the memory pressure on the processing unit.

[0109] Based on any of the above embodiments, step 310, which involves dividing the sequence into multiple segments and assigning processing units to each segment, includes:

[0110] The original sequence is divided into multiple sequences, and a set of processing units is assigned to each sequence.

[0111] The sequence is divided into multiple segments, and each segment is assigned a processing unit from the set of processing units.

[0112] Specifically, the original sequence is the sequence that requires attention mechanism computation. For long original sequences, in order to reduce the complexity and memory constraints of performing attention mechanism computation on long sequences, the original sequence can be split into multiple shorter sequences.

[0113] For each sequence obtained from the splitting, a set of processing units can be allocated for the attention mechanism computation specific to that sequence. That is, each sequence corresponds to a set of processing units, and different sequences correspond to different sets of processing units, thereby enabling parallel attention computation across sequences. It is understood that a set of processing units can include multiple processing units.

[0114] Based on this, each sequence can be split into multiple segments, and each segment can be assigned a processing unit from the set of processing units corresponding to that sequence, thus achieving a one-to-one correspondence between segments and processing units. Furthermore, segments in different sequences correspond to processing units in different sets of processing units.

[0115] For example, for an original sequence of length 512k, the original sequence can be split into four sequences of length 128k, and each sequence can be assigned a set of processing units. For example, assuming there are a total of 128 processing units, these 128 processing units can be divided into four sets of processing units, each set containing 32 processing units. Each set of processing units can store the complete model weights, and the complete model weights can be further divided within each set of processing units.

[0116] The method provided in this invention, when the original sequence is long, can divide the original sequence into sequences and assign a set of processing units to each sequence to achieve parallel computation of multiple processing unit sets. Based on the method provided in this invention, in an environment of parallel computation of multiple processing unit sets, since each processing unit can obtain the attention calculation result of the complete sequence for transmission between sets, compared to related technologies that require each processing unit to first perform intra-set transmission to obtain the attention calculation result of the complete sequence before it can be used for inter-set transmission, this reduces one transmission of attention calculation result between processing units.

[0117] Based on any of the above embodiments Figure 4 This is the second flowchart of the attention mechanism method provided by the present invention, as shown below. Figure 4 As shown, the processing unit in this embodiment of the invention is a GPU card. For scenarios involving four GPU cards and requiring attention calculation on a 128k sequence, this embodiment provides an attention mechanism calculation method, including the following steps:

[0118] The 128k sequence is divided into four segments and distributed to four cards. The allocation of these four segments across the four cards is denoted as P1_1, P2_2, P3_3, and P4_4, respectively. P1_1 represents segment 1 stored on the first card, P2_2 represents segment 2 stored on the second card, and so on. The corresponding calculations are as follows:

[0119] In step S1, the 128k sequence input is divided into four groups according to the number of cards, and each group processes a 32k segment.

[0120] Step S2: Calculate Q, K, and V for each card, resulting in Q, K, and V vectors of length 32k for each card, i.e., Q1_1, K1_1, V1_1, Q2_2, K2_2, V2_2… Here, Q1_1, K1_1, and V1_1 represent the Q, K, and V vectors calculated for segment 1 on the first card, Q2_2, K2_2, and V2_2 represent the Q, K, and V vectors calculated for segment 2 on the second card, and so on.

[0121] Step S3: Each card performs a 32k QK calculation (Q1_1*K1_1, Q2_2*K2_2, Q3_3*K3_3, Q4_4*K4_4) to obtain the result S sequence. While each card is performing SV calculation, the Q sequence is ring-transmitted between the cards. That is, the Q sequence of card 1 is passed to card 2, the Q sequence of card 2 is passed to card 3, and so on. For example, the first card receives the Q sequence calculated by the fourth card and records it as Q4_1 at the first card. The second card receives the Q sequence calculated by the first card and records it as Q1_2 at the second card. The third card receives the Q sequence calculated by the second card and records it as Q2_3 at the third card. The fourth card receives the Q sequence calculated by the third card and records it as Q3_4 at the fourth card.

[0122] After SV calculation, each card obtains a partial result for the current round, which is denoted here as the fragment calculation result.

[0123] Step S4: When each card performs the next round of QK calculation (Q4_1*K1_1, Q1_2*K2_2, Q2_3*K3_3, Q3_4*K4_4), it rings the partial results obtained from each card in the previous round, and performs an accumulation reduce operation based on the partial results existing at each card after the transmission.

[0124] Step S5, and so on, until all Q sequences and all partial results have been ring-transmitted between the cards. At this point, after the attention mechanism is calculated, each card has the attention calculation result of the complete sequence length.

[0125] Step S6: Based on the partial weights distributed across each card, perform subsequent processing.

[0126] Based on any of the above embodiments Figure 5 This is a flowchart illustrating the model reasoning method provided by the present invention, as shown below. Figure 5 As shown, the method includes:

[0127] Step 510, obtain the target sequence;

[0128] Step 520: Input the target sequence into the model for inference and obtain the inference result output by the model;

[0129] The model includes an attention mechanism, and during the inference process of the model, the calculation of the attention mechanism is performed based on the attention mechanism calculation method provided in the above embodiments.

[0130] The model inference method provided in this embodiment of the invention is applied to model inference of models that include attention mechanisms. For example, the model here can be a Transformer model used for tasks such as text classification and sentiment analysis, or a large language model used for question answering systems and dialogue generation, or a model used for computer vision tasks such as image classification, object detection, and object tracking. This embodiment of the invention does not specifically limit the model in this regard.

[0131] Here, the target sequence is the sequence to be used for model inference. For example, in a natural language processing scenario, the target sequence can be text; in a computer vision scenario, the target sequence can be an image sequence. This embodiment of the invention does not specifically limit this. The target sequence is the model input.

[0132] After obtaining the target sequence, it can be input into the model. The model then performs inference on the target sequence to obtain and output the inference result corresponding to the target sequence. For example, for text classification tasks, the inference result can be the text type to which the target sequence belongs; for sentiment analysis tasks, the inference result can be the sentiment type contained in the target sequence; for dialogue generation tasks, the inference result can be the response to the target sequence; for image classification tasks, the inference result can be the image type to which the target sequence belongs. This embodiment of the invention does not impose specific limitations on these aspects.

[0133] During the model's inference process, the attention mechanism included in the model can be executed based on the attention mechanism calculation method provided in the above embodiments. Thus, in an environment where multiple processing units perform attention mechanism calculations, each processing unit can obtain the attention calculation results of the complete sequence. As a result, each processing unit only needs to store the weights for each segment in the subsequent operations of attention calculation, which reduces the storage pressure on the processing units, avoids the waste of storage resources, improves the flexibility of subsequent parallel computing by multiple processing units, and thus reduces the model inference implementation cost and improves the flexibility of model inference.

[0134] The attention calculation device provided by the present invention will be described below. The attention calculation device described below and the attention calculation method described above can be referred to in correspondence.

[0135] Figure 6 This is a schematic diagram of the attention calculation device provided by the present invention, as shown below. Figure 6 As shown, the device includes:

[0136] The segmentation module 610 is used to divide the sequence into multiple segments, allocate processing units to each segment, and control each processing unit to calculate the query Q vector, key K vector, and numerical V vector of the allocated segment.

[0137] The first calculation module 620 is used to control each of the processing units to perform QKV calculations respectively, and to obtain the fragment calculation results of each processing unit. The QKV calculations include the Q vector, K vector and V vector calculated by the processing unit itself, and the Q vector calculated and transmitted by other processing units.

[0138] The second calculation module 630 is used to control each of the processing units to perform cumulative calculations, and to obtain the attention calculation results of the sequence at the segment corresponding to each processing unit. The cumulative calculations include the segment calculation results calculated by the processing unit itself and the segment calculation results calculated and transmitted by other processing units.

[0139] The apparatus provided in this invention performs QKV calculations involving Q vectors calculated by the processing unit itself and transmitted from other processing units at each processing unit, and performs summation calculations of segment calculation results calculated by the processing unit itself and transmitted from other processing units at each processing unit. As a result, in an environment where attention mechanism calculations are performed by multiple processing units, each processing unit can obtain the attention calculation results of the complete sequence. Thus, each processing unit only needs to store the weights for each segment in subsequent attention calculation operations, reducing the storage pressure on the processing units, avoiding the waste of storage resources, and improving the flexibility of subsequent parallel calculations by multiple processing units.

[0140] Based on any of the above embodiments, the device further includes a transmission module, used for:

[0141] While controlling each of the processing units to perform QKV calculations, at least one of Q-vector transmission and fragment calculation result transmission is performed between the processing units.

[0142] Based on any of the above embodiments, the transmission module is specifically used for:

[0143] While controlling each of the processing units to perform QK calculations in the current round of QKV calculations, the system also controls each of the processing units to transmit the fragment calculation results obtained in the previous round of the current round.

[0144] Based on any of the above embodiments, the transmission module is specifically used for:

[0145] While controlling each of the aforementioned processing units to perform SV calculation in QKV calculation, the system also controls each of the aforementioned processing units to perform Q vector transmission.

[0146] The SV calculation is performed on the result S of the QK calculation and the V vector; in one round of the QKV calculation, the Q vector used in the QK calculation is the same as the Q vector transmitted by the Q vector.

[0147] Based on any of the above embodiments, the device further includes a post-processing module for:

[0148] Each of the aforementioned processing units is controlled to perform subsequent processing on the attention calculation results based on its own stored segment processing weights;

[0149] The segment processing weight is obtained by dividing the segment into subsequent processing weights required for subsequent processing of the sequence based on the segments allocated by the processing unit.

[0150] Based on any of the above embodiments, the segmentation module is specifically used for:

[0151] The original sequence is divided into multiple sequences, and a set of processing units is assigned to each sequence.

[0152] The sequence is divided into multiple segments, and each segment is assigned a processing unit from the set of processing units.

[0153] The model reasoning apparatus provided by the present invention is described below. The model reasoning apparatus described below and the model reasoning method described above can be referred to in correspondence.

[0154] Figure 7 This is a schematic diagram of the structure of the model inference device provided by the present invention, as shown below. Figure 7 As shown, the device includes:

[0155] Module 710 is used to acquire the target sequence;

[0156] The reasoning module 720 is used to input the target sequence into the model for reasoning and obtain the reasoning result output by the model.

[0157] The model includes an attention mechanism, and during the inference process of the model, the calculation of the attention mechanism is performed based on the attention mechanism calculation method provided in the above embodiments.

[0158] The apparatus provided in this embodiment of the invention allows the attention mechanism included in the model to be executed based on the attention mechanism calculation method provided in the above embodiment during the model inference process. Thus, in an environment where multiple processing units perform attention mechanism calculations, each processing unit can obtain the attention calculation results of the complete sequence. As a result, each processing unit only needs to store the weights for each segment in the subsequent operations of attention calculation, which reduces the storage pressure on the processing units, avoids the waste of storage resources, improves the flexibility of subsequent parallel computing by multiple processing units, and thus reduces the model inference implementation cost and improves the flexibility of model inference.

[0159] Figure 8 An example is a schematic diagram of the physical structure of an electronic device, such as... Figure 8 As shown, the electronic device may include: a processor 810, a communication interface 820, a memory 830, and a communication bus 840, wherein the processor 810, the communication interface 820, and the memory 830 communicate with each other via the communication bus 840. The processor 810 can call logical instructions in the memory 830 to execute an attention mechanism calculation method, which includes:

[0160] The sequence is divided into multiple segments, and processing units are assigned to each segment. Each processing unit is controlled to calculate the query Q vector, key K vector, and numerical V vector of the assigned segment.

[0161] Each of the processing units is controlled to perform QKV calculations to obtain the fragment calculation results of each processing unit. The QKV calculations include the Q vector, K vector, and V vector calculated by the processing unit itself, as well as the Q vector calculated and transmitted by other processing units.

[0162] Each of the processing units is controlled to perform cumulative calculations, and the attention calculation results of the sequence are obtained at the segment corresponding to each processing unit. The cumulative calculations include the segment calculation results calculated by the processing unit itself and the segment calculation results calculated and transmitted by other processing units.

[0163] Alternatively, processor 810 may invoke logical instructions in memory 830 to execute a model inference method, the method comprising:

[0164] Obtain the target sequence;

[0165] The target sequence is input into the model for inference, and the inference result output by the model is obtained;

[0166] The model includes an attention mechanism, and during the inference process of the model, the calculation of the attention mechanism is performed based on the attention mechanism calculation method described above.

[0167] Furthermore, the logical instructions in the aforementioned memory 830 can be implemented as software functional units and, when sold or used as independent products, can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention, or the part that contributes to related technologies, or a portion of the technical solution, can be embodied in the form of a software product. This computer software product is stored in a storage medium and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute all or part of the steps of the methods described in the various embodiments of the present invention. The aforementioned storage medium includes various media capable of storing program code, such as USB flash drives, portable hard drives, read-only memory (ROM), random access memory (RAM), magnetic disks, or optical disks.

[0168] On the other hand, the present invention also provides a computer program product, the computer program product comprising a computer program that can be stored on a non-transitory computer-readable storage medium, wherein when the computer program is executed by a processor, the computer is capable of executing the attention mechanism computation method provided by the above methods, the method comprising:

[0169] The sequence is divided into multiple segments, and processing units are assigned to each segment. Each processing unit is controlled to calculate the query Q vector, key K vector, and numerical V vector of the assigned segment.

[0170] Each of the processing units is controlled to perform QKV calculations to obtain the fragment calculation results of each processing unit. The QKV calculations include the Q vector, K vector, and V vector calculated by the processing unit itself, as well as the Q vector calculated and transmitted by other processing units.

[0171] Each of the processing units is controlled to perform cumulative calculations, and the attention calculation results of the sequence are obtained at the segment corresponding to each processing unit. The cumulative calculations include the segment calculation results calculated by the processing unit itself and the segment calculation results calculated and transmitted by other processing units.

[0172] Alternatively, when a computer program is executed by a processor, the computer can execute the model reasoning methods provided by the above methods, which include:

[0173] Obtain the target sequence;

[0174] The target sequence is input into the model for inference, and the inference result output by the model is obtained;

[0175] The model includes an attention mechanism, and during the inference process of the model, the calculation of the attention mechanism is performed based on the attention mechanism calculation method described above.

[0176] In another aspect, the present invention also provides a non-transitory computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the attention mechanism computation method provided by the above methods, the method comprising:

[0177] The sequence is divided into multiple segments, and processing units are assigned to each segment. Each processing unit is controlled to calculate the query Q vector, key K vector, and numerical V vector of the assigned segment.

[0178] Each of the processing units is controlled to perform QKV calculations to obtain the fragment calculation results of each processing unit. The QKV calculations include the Q vector, K vector, and V vector calculated by the processing unit itself, as well as the Q vector calculated and transmitted by other processing units.

[0179] Each of the processing units is controlled to perform cumulative calculations, and the attention calculation results of the sequence are obtained at the segment corresponding to each processing unit. The cumulative calculations include the segment calculation results calculated by the processing unit itself and the segment calculation results calculated and transmitted by other processing units.

[0180] Alternatively, when the computer program is executed by a processor, it is implemented to perform the model inference methods provided by the methods described above, the method comprising:

[0181] Obtain the target sequence;

[0182] The target sequence is input into the model for inference, and the inference result output by the model is obtained;

[0183] The model includes an attention mechanism, and during the inference process of the model, the calculation of the attention mechanism is performed based on the attention mechanism calculation method described above.

[0184] The device embodiments described above are merely illustrative. The units described as separate components may or may not be physically separate. The components shown as units may or may not be physical units; that is, they may be located in one place or distributed across multiple network units. Some or all of the modules can be selected to achieve the purpose of this embodiment according to actual needs. Those skilled in the art can understand and implement this without any creative effort.

[0185] Through the above description of the embodiments, those skilled in the art can clearly understand that each embodiment can be implemented by means of software plus necessary general-purpose hardware platforms, and of course, it can also be implemented by hardware. Based on this understanding, the above technical solutions, in essence or the parts that contribute to the related technology, can be embodied in the form of software products. This computer software product can be stored in a computer-readable storage medium, such as ROM / RAM, magnetic disk, optical disk, etc., and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute the methods described in the various embodiments or some parts of the embodiments.

[0186] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, and not to limit them; although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some of the technical features; and these modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the spirit and scope of the technical solutions of the embodiments of the present invention.

Claims

1. A method for calculating an attention mechanism, characterized in that, include: The sequence is divided into multiple segments, and processing units are assigned to each segment. Each processing unit is controlled to calculate the query Q vector, key K vector, and numerical V vector of the assigned segment. Each of the processing units is controlled to perform QKV calculations to obtain the fragment calculation results of each processing unit. The QKV calculations include the Q vector, K vector, and V vector calculated by the processing unit itself, as well as the Q vector calculated and transmitted by other processing units. Each of the processing units is controlled to perform cumulative calculations, and the attention calculation results of the sequence are obtained at the segment corresponding to each processing unit. The cumulative calculations include the segment calculation results calculated by the processing unit itself and the segment calculation results calculated and transmitted by other processing units.

2. The attention mechanism calculation method according to claim 1, characterized in that, Also includes: While controlling each of the processing units to perform QKV calculations, at least one of Q-vector transmission and fragment calculation result transmission is performed between the processing units.

3. The attention mechanism calculation method according to claim 2, characterized in that, The step of controlling each of the processing units to perform QKV calculations while simultaneously transmitting Q vectors and transmitting fragment calculation results between the processing units includes at least one of the following: While controlling each of the processing units to perform QK calculations for the current round of QKV calculations, the system also controls each of the processing units to transmit the fragment calculation results obtained in the previous round for the current round.

4. The attention mechanism calculation method according to claim 2, characterized in that, The step of controlling each of the processing units to perform QKV calculations while simultaneously transmitting Q vectors and transmitting fragment calculation results between the processing units includes at least one of the following: While controlling each of the aforementioned processing units to perform SV calculation in QKV calculation, the system also controls each of the aforementioned processing units to perform Q vector transmission. The SV calculation is performed on the result S of the QK calculation and the V vector; in one round of the QKV calculation, the Q vector used in the QK calculation is the same as the Q vector transmitted by the Q vector.

5. The attention mechanism calculation method according to claim 1, characterized in that, Also includes: Each of the aforementioned processing units is controlled to perform subsequent processing on the attention calculation results based on its own stored segment processing weights; The segment processing weight is obtained by dividing the segment into subsequent processing weights required for subsequent processing of the sequence based on the segments allocated by the processing unit.

6. The attention mechanism calculation method according to claim 1, characterized in that, The process of dividing the sequence into multiple segments and assigning processing units to each segment includes: The original sequence is divided into multiple sequences, and a set of processing units is assigned to each sequence. The sequence is divided into multiple segments, and each segment is assigned a processing unit from the set of processing units.

7. A model reasoning method, characterized in that, include: Obtain the target sequence; The target sequence is input into the model for inference, and the inference result output by the model is obtained; The model includes an attention mechanism, and during the inference process of the model, the calculation of the attention mechanism is performed based on the attention mechanism calculation method as described in any one of claims 1 to 6.

8. An attention mechanism computing device, characterized in that, include: The segmentation module is used to divide the sequence into multiple segments, allocate processing units to each segment, and control each processing unit to calculate the query Q vector, key K vector, and numerical V vector of the allocated segment. The first calculation module is used to control each of the processing units to perform QKV calculations respectively, and to obtain the fragment calculation results of each processing unit. The QKV calculations include the Q vector, K vector and V vector calculated by the processing unit itself, and the Q vector calculated and transmitted by other processing units. The second calculation module is used to control each of the processing units to perform cumulative calculations, and to obtain the attention calculation results of the sequence at the segment corresponding to each processing unit. The cumulative calculations include the segment calculation results calculated by the processing unit itself and the segment calculation results calculated and transmitted by other processing units.

9. A model reasoning device, characterized in that, include: The acquisition module is used to acquire the target sequence; The reasoning module is used to input the target sequence into the model for reasoning and obtain the reasoning result output by the model. The model includes an attention mechanism, and during the inference process of the model, the calculation of the attention mechanism is performed based on the attention mechanism calculation method as described in any one of claims 1 to 6.

10. An electronic device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, characterized in that, When the processor executes the program, it implements the attention mechanism calculation method as described in any one of claims 1 to 6, or the model reasoning method as described in claim 7.

11. A non-transitory computer-readable storage medium having a computer program stored thereon, characterized in that, When the computer program is executed by the processor, it implements the attention mechanism calculation method as described in any one of claims 1 to 6, or the model reasoning method as described in claim 7.

12. A computer program product, comprising a computer program, characterized in that, When the computer program is executed by the processor, it implements the attention mechanism calculation method as described in any one of claims 1 to 6, or the model reasoning method as described in claim 7.

Images