Method for processing an attention layer of a machine learning model

Abstract

A computer-implemented method of processing an attention layer of a neural network includes receiving an input sequence comprising a plurality of tokens and splitting this input into subsequences. A query tensor, key tensor and value tensor are generated for each subsequence. A first query tensor, key tensor and value tensor generated for a first subsequence are provided to a processing device, and based on these tensors, a partial attention tensor is calculated. The method iteratively provides a next query tensor, key tensor and value tensor generated for a next subsequence to the processing device, and the partial attention tensor is updated based on these next subsequence tensors. Upon receipt at the processing device of a final subsequence query tensor, key tensor and value tensor, a final attention tensor is output. This method provides a means of reducing memory consumption during attention mechanism calculations.

Description

METHOD FOR PROCESSING AN ATTENTION LAYER OF A MACHINE LEARNING MODELField

[0001] The present application relates to a computer-implemented method of processing an attention layer of a machine learning model. The application further relates to computer systems, computer program products and computer-readable storage media for accomplishing the same.Background

[0002] Foundation models based on deep neural networks (DNNs) have become increasingly popular, owing to their impressive performance on a variety of tasks. For example, transformer-based natural language processing models such as the Generative Pretrained Transformer (GPT) models provided by Open Al® are able to provide realistic natural language output, as well as being able to perform tasks including summarisation, classification, translation and so on. Furthermore, such models display emergent characteristics, such as the ability to perform arithmetic, encode data and so on. Similarly impressive models are available for image generation and other modalities.

[0003] Key to the performance of these models is the size of their training data set and the size of their parameter space. GPT-3 for example has 175 billion parameters, and was trained on approximately 500 billion tokens, with GPT-4 rumoured to have over 1 trillion parameters and a training data set of over 10 trillion tokens. The growth in parameter space is trending exponentially.

[0004] As a result of this enormous parameter space and training data set size, training large DNNs requires significant computational resource. Training large DNNs can be accomplished using aggregated computing resources, through parallel and distributed computation that makes use of specialist work accelerator hardware rather than relying on traditional general-purpose CPUs (central processing units). An example of work accelerator hardware is Graphical Processing Units (GPUs), such asthose provided by NVIDIA®. Other example work accelerator hardware includes Tensor Processing Units (TPUs) provided by Google®, Deep Learning Processor Units (DPUs) provided by AMD® and so on.

[0005] The memory capacity of work accelerators has not grown in a similar manner to the size of DNN parameter space. For example, from 2018 to 2021, model size grew by over 1000-fold, but high-end GPU memory increased only three-fold. This lack of memory capacity equates to a requirement for ever larger distributed computing systems, with associated cost. For example, to train a 10 billion parameter model efficiently, a system including 16 GPUs is required at an estimated cost of over 100K USD. This takes the training of such models out of reach of smaller businesses and academic institutions.Summary

[0006] The transformer architecture is commonly employed in Large Language Models (LLMs) such as GPT, owing to its adaptable scalability and powerful modelling and learning capabilities. Layers of such models include transformer blocks, which in turn include an attention mechanism, which captures the long-term contextual relationship between every pair of tokens in input sequences to the layer. Typically, in LLMs the attention mechanism will be a self-attention mechanism, with the relationships being between every pair of tokens in the same input sequence, though in some applications the relationships may be between different sequences (e.g. texts in different languages).

[0007] This mechanism provides an efficient and effective means of capturing contextual information, when compared to previous techniques such as recurrent neural networks and long-short term memory networks. Empirical evidence demonstrates that when the length (Lsequence) of the input sequence increases, the model is able to capture a longer context, resulting in improved performance. However, the execution of the attention mechanism in previous techniques relies on allocating memory in a quadratic manner in relation to the length of the inputsequence. In other words, the memory allocation is O(Lsequence2). This limits the possibility of training based on longer sequences, especially when making use of accelerator hardware, which has relatively limited memory capacity as discussed above. Typically, models are constrained to (-sequence of around 2000 (i.e. 2K).

[0008] Some prior techniques have attempted to address these issues by dividing the calculation of attention across a plurality of devices. Broadly, these fall into two types of techniques. The first technique segments the tensors used in the attention calculation (Q, K, V, discussed in detail below) and distributes the segments of tensors to different processing nodes for processing. The segments are then gathered from the processing nodes to reconstruct the complete tensors. The second technique involves using cross-node communication primitives during the calculation - in other words, the nodes must communicate during the calculation. These strategies do not address the underlying issue of the allocation of memory in a quadratic manner, and instead leverage more memory by using multiple devices.

[0009] To address the issues outlined above, the disclosure provides an approach for reducing memory consumption during attention mechanism calculations, which enables training on substantially longer sequences of 32K, 64K or more. In fact, the techniques can be applied to substantially any sequence length. The techniques are applicable to both local and distributed runtime contexts. In other words, the techniques can be applied to configurations in which a single accelerator is processing the attention mechanism, or where the processing is distributed over multiple accelerators.

[0010] According to a first aspect of the disclosure, there is provided a computer-implemented method of processing an attention layer of a neural network, the method comprising: receiving an input sequence comprising a plurality of tokens; splitting the input sequence into a plurality of subsequences; generating respective subsequence query tensors, subsequence key tensors and subsequence value tensors for each of the plurality of subsequences; providing, to a processing device, a first subsequence query tensor, first subsequence key tensor and first subsequence valuetensor associated with a first subsequence of the plurality of subsequences; calculating, by the processing device, a partial attention tensor based on the first subsequence query tensor, first subsequence key tensor and first subsequence value tensor; iteratively: providing a next subsequence query tensor, next subsequence key tensor and next subsequence value tensor associated with a next subsequence of the plurality of subsequences to the processing device; updating the partial attention tensor based on the next subsequence query tensor, next subsequence key tensor and next subsequence value tensor; and upon receipt at the processing device of a final subsequence query tensor, final subsequence key tensor and final subsequence value tensor associated with a final subsequence of the plurality of subsequences, outputting by the processing device a final attention tensor.

[0011] Updating the partial attention tensor for an ithiteration may comprise calculating an ithattention score from the ithsubsequence query tensor, ithsubsequence key tensor and ithsubsequence value tensor. Updating the partial attention tensor for the ithiteration may comprise storing the ithattention score in a location of the partial attention tensor corresponding to an ithrow and an ithcolumn.

[0012] Updating the partial attention tensor may further comprise, for every location in the ithrow before the ithcolumn, calculating an attention score based on a subsequence key tensor and subsequence value tensor stored in a previous iteration and the ithquery tensor, and updating the location to store the attention score. Updating the partial attention tensor may further comprise for every location in the the ithrow after the ithcolumn, storing the ithsubsequence query tensor.

[0013] Updating the partial attention tensor may further comprise: for every location in the ithcolumn before the ithrow, calculating an attention score based on a subsequence query tensor stored in a previous iteration and the ithsubsequence key tensor and ithsubsequence value tensor, and updating the location to store the attention score. Updating the partial attention tensor may further comprise for every location in the ithcolumn after the ithrow, storing the ithsubsequence key tensor and ithsubsequence value tensor.

[0014] Updating the partial attention tensor may comprise accumulating the ithattention score with previous attention scores stored in the ithrow of the partial attention tensor. Updating the partial attention tensor may comprise dividing the product of the partial query tensor and the partial query tensor, suitably before accumulating. Updating the partial attention tensor may comprise applying a softmax function to the product of the partial query tensor and the partial query tensor, suitably before accumulating.

[0015] The key tensors referenced herein may be transpose versions of the key tensors.

[0016] A shape of at least one, but suitably each of the subsequence query tensors, subsequence key tensorsand subsequence value tensors may be (BN, seq, hs). BN may represent a training batch size multiplied by a number of attention heads, seq may represent a length of each subsequence and hs may represent a number of attention heads. The method may comprise allocating a memory region of the processing device at an ithiteration having a size i * 4 * (BN, seq, hs).

[0017] The method may comprise releasing the storage of the previous iterations upon calculation of the final attention tensor.

[0018] The attention layer may be part of a transformer block of the neural network. The attention layer may be a self-attention layer. The neural network may have billions of parameters. The neural network may comprise a plurality of attention layers. The neural network may be an image processing neural network. The neural network may be a language model. The neural network may be a multimodal model. The values in the tensors may be stored in a low-precision number format, such as FP32, FP16, FP8 or a combination of number formats.

[0019] The length of each subsequence may be IK, 2K, 3K, 4K, 5K, 6K, 7K, 8K or 16K. The length of the sequence may be at least 16K, suitably at least 32K or 64K.

[0020] The method may be carried out during training of the neural network. The method may be carried out at inference time of the neural network.

[0021] The memory usage of the method may scale linearly with the length of the sequence.

[0022] The processing device may be a work accelerator. The work accelerator may be a graphics processing unit, tensor processing unit or deep learning processor unit. The processing device may comprise a memory and processing hardware. The processing device may be connected to a host computer, which may comprise a host processor and host memory. The processing device memory may be distinct from the host memory. The capacity of the processing device memory may be 100 gigabytes or less. The host computer and processing device may be collocated. The host computer may provide the tensors to the processing device. The final attention tensor may be output by the processing device to the host computer.

[0023] According to a second aspect of the disclosure, there is provided a computer-implemented method of processing an attention layer of a neural network, the method comprising: receiving an input sequence comprising a plurality of tokens; splitting the input sequence into a plurality of subsequences; generating respective subsequence query tensors, subsequence key tensors and subsequence value tensors for each of the plurality of subsequences; providing, to each of a plurality of processing devices, a respective subsequence query tensor, subsequence key tensor and subsequence value tensor; calculating, by each processing device, a respective partial attention tensor based on the provided subsequence query tensor, subsequence key tensor and subsequence value tensor; receiving, from each processing device, the respective partial attention tensor; and generating a final attention tensor from the received partial attention tensors.

[0024] Calculating the partial attention tensor for an ithdevice may comprise calculating an ithattention score from the ithsubsequence query tensor, ithsubsequence key tensor and ithsubsequence value tensor; and storing the ithattentionscore in a location of the partial attention tensor corresponding to an ithrow and an ithcolumn.

[0025] Further optional features of the second aspect are defined above in respect of the first aspect and can be combined in any combination.

[0026] The methods of the first and second aspects may be combined. Accordingly, there may be provided a method in which the subsequences are arranged into groups, and each group is provided to a respective processing device. At each processing device, the processing may be as defined in the first aspect. The final attention tensors calculated for each group may then be combined.

[0027] According to a third aspect of the disclosure, there is provided a computer system comprising a processor and a memory storing instructions, which when executed, implement any of the methods disclosed herein. The computer system may comprise a host computer and a processing device. The host computer may have memory storing first instructions for carrying out the steps defined in the methods above that are not stated to be carried out by the processing device. The processing device may have memory storing second instructions for carrying out the steps stated to be carried out by the processing device.

[0028] There is also provided one or more computer-readable storage media storing instructions, which when executed, cause a computer system to implement any of the methods disclosed herein. The disclosure also extends to a computer program product comprising instructions, which when executed, cause a computer system to implement any of the methods disclosed herein.Brief Description of the Drawings

[0029] To assist understanding of the present disclosure and to show how embodiments may be put into effect, reference is made by way of example to the accompanying drawings in which:

[0030] FIG. 1 is a schematic block diagram of an example system according to examples of the disclosure.

[0031] FIG. 2 is a schematic block diagram of an example machine learning model.

[0032] FIG. 3 is a schematic illustration of a transformer machine learning model.

[0033] FIG. 4 is a table illustrating operations and tensor shapes associated with a prior art method of calculating attention.

[0034] FIG. 5 is a schematic diagram illustrating the splitting of a sequence into subsequences with associated query, key and value tensors.

[0035] FIG. 6 is a schematic flowchart of a first example technique of the disclosure.

[0036] FIG. 7 is a table illustrating the accumulation of an attention tensor according to the method of FIG. 6.

[0037] FIG. 8 is a schematic block diagram of a system associated with a second example technique of the disclosure.

[0038] FIG. 9 is a schematic block diagram of an example computing system.Detailed Description

[0039] In overview, examples of the disclosure provide means of calculating attention tensors that consumes a memory in a manner that scales linearly with the length of the sequence. This is accomplished by splitting the input sequence into subsequences, each of which have corresponding query, key and value tensors. In one example, an attention tensor is generated cumulatively, by processing the query, key and value tensors of the subsequences in order, iteratively updating a partial attention tensor. In another example, the attention tensor is generated in a distributed manner, by distributing the query, key and value tensors of the subsequences to different devices.

[0040] FIG. 1 illustrates a computer system 100 according to an example of the disclosure. The computer system 100 comprises a processor 110, volatile memory 120, storage 130 and a work accelerator subsystem 200.

[0041] The processor 110 may take the form of a central processing unit (CPU), which is general purpose in nature. The CPU may be any suitable processor, including a multi-core processor. The volatile memory 120 may be random access memory, and is configured to transiently store any suitable data for the operation of the system 100, including instructions for execution by the CPU 110. The storage 130 comprises any suitable storage media for non-transient (i.e. permanent) data storage and may encompass suitable hard disk drives, solid state drives and so on. In some examples, the storage 130 may encompass secondary storage devices, such as the NVIDIA® NVMe SSD or the like.

[0042] The components 110, 120, 130 are components of the host computer system 100, as distinct from the components of the work accelerator subsystem 200. Accordingly, they may be referred to herein as the host processor 110, host memory 120 and host storage 130 to distinguish them from corresponding components of the work accelerator 200.

[0043] The work accelerator subsystem 200 is a subsystem which processing can be offloaded to by the host system 100. In the context of the present disclosure, elements of the processing of (e.g. the training of) a machine learning model are offloaded to the work accelerator subsystem 200. The work accelerator subsystem 200 includes specialised (i.e. special purpose) hardware that renders it suitable for carrying out the offloaded processing in an accelerated manner, when compared to the general purpose CPU 110.

[0044] An example of a suitable work accelerator subsystem 200 is a graphics processing unit (GPU). Whilst GPUs were initially designed for accelerating computer graphics and image processing, such GPUs have found great utility in the training of machine learning models. For example, the accelerator subsystem may be an NVIDIA®V100 GPU or an NVIDIA A10 GPU, or any other suitable GPU provided by NVIDIA or any other supplier.

[0045] Other example work accelerator subsystems 200 are possible and contemplated within the scope of the disclosure. Other example work accelerator subsystems 200 include Tensor Processing Units (TPUs) provided by Google®, Deep Learning Processor Units (DPUs) provided by AMD® and so on.

[0046] Whatever the particular accelerator subsystem 200 used, the accelerator subsystem 200 includes processing hardware 210 and a local memory 220, hereinafter referred to as an accelerator processor 210 and accelerator memory 220. In general, the memory capacity of the accelerator memory 220 will be substantially smaller than that of the host memory 120. For example, whereas the RAM 120 of the host computer may be as large as 512GB or 1024GB, the size of the available local memory may be 100GB or less, such as 80 GB, 32GB or 16GB. As discussed further hereinbelow, examples of the disclosure are able to make efficient use of the relatively small amount of available accelerator memory 220.

[0047] The host system 100 and its accelerator subsystem 200 may be collocated. That is to say, the host system 100 and its accelerator 200 may effectively form one computer system, located in a single chassis or housing. For example, the accelerator 200 subsystem may be disposed on a suitable expansion card (e.g. a PCIe card) in the chassis, or direct mounted to a motherboard on which the host processor 110 is located. The subsystem 200 and host system 100 may be connected by a suitable bus.

[0048] The system 100 is a simplified example version of system in accordance with the disclosure to illustrate the concepts discussed herein. It will be understood that in reality, a single host computer 100 may host a plurality of accelerator subsystems 200. Equally, each accelerator subsystem 200 may comprise multiple processors 210 - e.g. a GPU may comprise a plurality of streaming multiprocessors configured to operate in parallel.

[0049] Turning now to FIG. 2, a brief discussion follows of common concepts in the execution of machine learning models. FIG. 2 schematically illustrates the structure of an example machine learning model 300, of the sort processed by the example techniques discussed herein. The model 300 is a layered model, such as a deep neural network. The model 300 comprises a plurality of layers, generally indicated by reference numeral 301. The layers 301 are arranged in sequence from a first layer 302 or input layer to a final layer 305 or output layer, such that each layer 301 receives input from the preceding layer 301 and provides output to the subsequent layer 301. In the context of a neural network, each layer 301, may comprise a plurality of artificial neurons, which receives a plurality of inputs from the preceding layer. Each input has an assigned weight, and the output is calculated as a function of the weighted input, referred to as an activation function. The task of training the neural network involves determining the weights by minimising a suitable loss function based on training data. The weights are an example of model data (also referred to herein as model states or model parameters) of the model 300.

[0050] In the example of FIG. 2, the model 300 is a simplified version of a transformer deep neural network. The first layer 302 is an embedding layer, which generates a vector in an embedding space. The layers 303-1 to 303-n are transformer blocks, which each comprise an attention layer and feedforward layer (both not shown). The layer 304 is a normalisation layer, and layer 305 is a pooling layer.

[0051] It will be understood that this is merely an example of a model architecture to which the present techniques apply, which has a relatively straightforward layer structure in which layers feedforward. The techniques may be applied to substantially any layered machine learning model, including an attention mechanism, which may be part of a transformer layer or part of some other layer. This includes networks that have residual components or gating mechanisms such as mixture-of-experts (MoE) models.

[0052] As illustrated in FIG. 2, the training process comprises 3 phases, each of which is typically iterated millions of times when training a large DNN, across multiple training epochs, wherein each epoch uses the entirety of the training data set.

[0053] In each iteration, there is a forward propagation (FP) phase, in which a batch of training data is passed in a forward direction through the model 300 (i.e. from the first layer 302 to the last layer 305). A loss is computed, which reflects the error in the output of the model 300, using an objective function. This is then followed by a backward propagation (BP) phase, in which the loss value is propagated back through the model 300 (i.e. in the opposite direction to the FP), to compute gradients. The gradients are vectors representing the derivative of the loss function with respect to the weights. Finally, in the third phase, an optimiser uses the aggregated gradients to update the weights of individual model layers. Example optimisers include stochastic gradient descent (SGD) or the well-known Adam optimiser.

[0054] There now follows a general discussion of transformer layers, and in particular the attention calculations that are carried out therein. The transformer architecture is well-known in the art and will be familiarto the skilled reader, such that the following discussion provides useful context and background. Discussions herein of attention mechanisms, attention calculations, attention computation, computation of attention score and the like should be understood to be interchangeable.

[0055] FIG. 3 illustrates an example of the transformer model architecture in more detail, as originally set out in Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems 30 (2017), the contents of which are incorporated herein by reference in their entirety. As illustrated, the model includes N transformer blocks in the encoder of the architecture and N transformer blocks in the decoder of the architecture. Each transformer block includes a multiheaded attention layer as a first sublayer, followed by a feedforward layer as a second sublayer. In the decoder, the blocks additionally include a third sublayer in the form of another multiheaded attention layer. The first multiheaded attention layer is masked, in the sense the attention layer is modified such that it prevents tokens in thesequence from taking into account subsequent tokens (i.e. the attention is based on only previous tokens).

[0056] In addition, the model comprises embedding layers for the input and outputs that convert input tokens into an embedding space for processing, as well as a positional encoding mechanism. Furthermore, output from the decoder part of the model may be passed through a linear layer and a softmax layer to provide output probabilities for the next token in a sequence given an input sequence.

[0057] In general, the attention sublayers (hereinafter referred to simply as attention layers) in a typical LLM model are self-attention layers, meaning that the attention is calculated between each pair of tokens in the same sequence, rather than a cross-attention layer between two different sequences. Furthermore, the figure shows multi-head attention layers, in which the relevant computation is divided amongst multiple "heads" of attention, each computing attention for an allocated slice of the relevant input tensors.

[0058] Turning now to the attention mechanism in more detail, the mechanism relies on query vectors (q), key vectors (k) and value vectors (v), corresponding to each token in the input sequence. These vectors q, k, v are derived from the embeddings of respective tokens of the input sequence, by multiplying them by respective weight matrices Wq, Wk, Wv. These vectors for the whole sequence are packed together into respective query, key and value matrices, Q, K and V.

[0059] The core computation carried out is:

[0060] where dk is the number of dimensions of K.

[0061] Setting aside for a moment the termthe basis that it is a static value, and the softmax operation, the underlying computation carried out is QKTV.

[0062] Furthermore, given that training occurs across multiple heads, and will involve multiple sequences divided into training batches, the values for each batch and head are also packed together, such that in the following Q, K and V in practice are multi-dimensional tensors. The shapes of the tensors involved are as follows. The shape of each tensor Q, K, V is defined by the following parameters:• batch size (batch_size), which is the number of input sequences contained in each batch of training examples• the number of attention heads (num_heads)• the sequence length (sequencejength or seq). . . . . . . . . . . . hidden_dimension . . .• the head size (head—size or hs), which is - , wherein the numjieads hidden_dimension is the dimensionality of the query, key or value vector (as the case may be).

[0063] Consequently, each of Q, K and V has the following shape: (batch_size, num._h.eads, sequencejength, head_size).

[0064] The tensors can be reshaped to combine the dimensions representing the batches and the number of heads, which then gives a shape BN, seq, hs), in which BN is batch_size multiplied by num_heads.

[0065] In prior approaches, the attention calculation is carried out in the steps set out in the table shown in FIG. 4, with the shape of the tensors provided as input arguments and output also set out.

[0066] In the table of FIG. 4, the matmul function is the matrix multiplication operation. The softmax function transforms the vectors in the input tensor into probability distributions, with the parameter dim = -1 indicating the operation is applied to the last dimension of the input tensor.

[0067] Walking through the steps in the table, step 1 is the initial operation is the matrix multiplication of Q and KT(i.e. the transpose of K). The inputs are Q and KT, Q having the shape (BN, seq, hs) and KThaving the shape (BN, hs, seq) owing to the transposition. The output is a tensor with the shape (BN, seq, seq). This output must be stored as a first temporary tensor.

[0068] Step 2 is the application of the softmax function to the tensor W, which consequently receives the matrix W having shape (BN, seq, seq) as input. As the softmax function does not change the shape of the tensor, the output is a second temporary tensor, also with shape (BN, seq, seq). The second temporary tensor may overwrite the first temporary tensor.

[0069] Finally, in step 3 the second temporary tensor is multiplied by the tensor V. Accordingly, the inputs are respectively of shape (BN, seq, seq) and (BN, seq, hs), which results in a final output in the shape (BN, seq, hs).

[0070] Accordingly, the temporary memory allocation required for this approach is 3 * (BN, seq, hs) to account for the Q, K and V tensors and (BN, seq, seq) for the temporary tensor, summing to (3 * seq * hs+seq2) * BN. In big O notation, this approximates to O(BN * seq2), which thus scales quadratically with the length of the input sequence.

[0071] FIG. 5 illustrates attention calculation according to examples of the disclosure. In general terms, the present techniques take a plurality of input sequences 400, each of which comprises a plurality of tokens, tokeni to tokens where L is the length of the input sequence. These input sequences 400 are divided into n microsequences 410, of the same size, also referred to herein as subsequences orseq / where / is the index of the subsequence among the n subsequences.

[0072] Each subsequence has the same length subseqjength. A first subseqence, seqi, therefore includes tokeni to tokenSUbseq_iength, a second subsequence seq2 includes tokenSUbseq_iength+i to token2*subseq_iength and so on until a final subsequence seqnincluding tokeni-subseqjength to tokeni.

[0073] One example value of subseqjength is 4K. Particularly, in one example a 4K subseqjength was implemented in training a 7-billion parameter model using the Adam optimiser with mixed-precision. In the example, the mixed precision was FP32 (i.e. float32 format) for the parameters and FP16 (i.e. floatl6) for momentum values and gradients on an Nvidia A100. In other examples, the subseqjength may be IK, 2K, 3K, 5K, 6K, 7K, 8K or any other suitable value.

[0074] Each subsequence has corresponding subsequence Q, KTand V tensors, respectively indicated by reference numerals 421, 422, 423, and hereinafter referred to as Qi, KTi, Vi , based on the tokens comprised in the subsequence. Accordingly, seqi has Qi, KTi, and Vi, seq2 has Q2, KT2 and V2, through to seqnwith Qn, KTn, and Vn.

[0075] In practice, the subsequence tensors can be generated by generating Q, KTand V over the whole sequence 400, and then subsequently partitioning Q, KTand V into the subsequences in the dimension of the tensors representative of the sequence length. In examples, the host processor 110 partitions the tensors and provides them to the accelerator 200.

[0076] In terms of the calculation of QKTV based on these subsequences, asdiscussed above Q is partitioned into , KTis partitioned into / (Kpartitioned into I I.\vj

[0077] QKTV is therefore equal to the following expressions:

[0078] In other words, calculating QKTV can be achieved by summing QiKTiVi to QiKTnVn, Q2KTiVi to Q2KTnVn, and so on up to QnKTiVi to QnKTnVn. Put differently, QKTV calculation can be broken down into subsequences, by calculating QiKTiVi to QiKTnVn, Q2KTIVI to Q2KTnVn, and so on up to QnKTiVi to QnKTnVn, and then summing each row Qi..Qn. The techniques herein take advantage of this insight, which allows attention calculations to be carried out accumulatively across subsequences in a manner that avoids quadratic scaling of memory usage.

[0079] FIG. 6 illustrates a first example method of using Qi...Qn, KTi... KTnand Vi...Vnto calculate the attention score in a manner that is memory efficient. In overview, the attention is calculated accumulatively, subsequence-by-subsequence, based on the respective Qi, KTi, Vi tensors. The method illustrated in FIG. 6 can be executed on a single local processing device, such as work accelerator 200, such that an entire sequence is processed by the single local device.

[0080] As illustrated in step S601, the value of i is initially 1. Then in step S602, a loop begins through the subsequences. In step S602, Qi, KTi, Vi are provided to the work accelerator 200.

[0081] In step S603, Qi, KTi, Vi are used to calculate a partial attention tensor. For the first iteration, the partial attention tensor is calculated based on Qi, KTi, Vi alone. However, in each subsequent iteration, the newly received Qi, KTi, Vi are used to update the partial attention tensor stored in the previous iteration to account for thenew subsequence. In other words, step S603 accumulates the attention tensor over the iterations.

[0082] In step S604, the method determines whether the end of the subsequences has been reached - i.e. whether i = n. If no, i is incremented in step S605 and the loop moves on to the next iteration. If i = n is true, then the loop terminates and the partial attention tensor is used to output the final attention tensor in step S606.

[0083] FIG. 7 illustrates the process of accumulating the attention tensor in more detail.

[0084] The first row 701 of the table of FIG. 7 illustrates the calculation carried out on receipt of Qi, KTi, Vi, which are associated with the first subsequence, seqi. The partial attention tensor 750, illustrated in column 711, is calculated from these tensors.

[0085] As each of Qi, KTi, Viare available, the top left corner (i.e. the first row of the first column) of the partial attention tensor 750 can be calculated by multiplying Qi, KTi, and Vi. As the entire top row of the tensor 750 will eventually be calculated as a product of Qi, Qi is populated across the remainder of the entire top row aside from the top left corner. Similarly, as the entire first column is a product of KTi and Vi, the product of these tensors is calculated and populated in the entirety of the first column aside from the top left corner.

[0086] The second row 702 of the table of FIG. 7 illustrates the calculation carried out on receipt of Q2,KT2,V2, which are associated with the second subsequence, seq2. The partial attention tensor 750 is updated based on the newly received values. The second row of the second column is calculated by multiplying Q2, KT2 and V2. In the second row of the first column KTiVi is already stored, and Q2 is now available, so Q2KTIVI can be calculated. The remainder of the second row is populated with Q2. In the second column, first row, Qi is already stored from the previous iteration and KT2 and V2 are now available, so QIKT2V2 can be calculated. The remainder of the second column is populated with KT2V2. As QiKTiVi and QIKT2V2 are now both available, thesevalues in the first row can be accumulated (i.e. summed together). Similarly, as 2KTIVI and CbKT2V2 are now both available, these values in the second row can be accumulated. Although in the following (and herein generally) reference is made to further columns and rows of the tensor, it will be understood where the contents of those rows are shown as or have said to have been accumulated, only the accumulated value is stored, and not the components of that value that have been summed to arrive at the accumulated value.

[0087] The third row 703 of the table of FIG. 7 illustrates the calculation carried out on receipt of 3,KT3,V3, which are associated with the third subsequence, seqs. The partial attention tensor 750 is updated based on the newly received values. The third row of the third column is calculated by multiplying 3, KT3 and V3. The two preceding columns of the third row already have KTiVi and KT2V2 respectively stored therein, so can be multiplied by newly received 3 to complete them. The remainder of the third row is populated with 3.

[0088] Similarly, in the third column the two preceding rows already have Qi and 2 respectively stored therein from the previous iterations. As KT3 and V3 are now available, IKT3V3 and CbKT3V3 can be calculated. The remainder of the third column is populated with KT3V3.

[0089] Once again, QIKT3V3 can be accumulated with the previously accumulated QiKTiVi + QIKT2V2. Similarly, Q2KT3V3 is accumulated with the previously accumulated Q2KTIVI + Q2KT2V2.

[0090] As indicated in row 704 of the table of FIG. 5, this generalises to the appending of the ithrow with Qi and ithcolumn with KTiVi, and completing any available locations at which all of Q, KTand V are present. Furthermore, where Q, KTand V are all present, this can be summed with the previous accumulation of Q, KTand V in that row.

[0091] Row 705 of the table of FIG. 5 illustrates the penultimate step associated with the penultimate subsequence seqn-i, in which the n-lthrow with is appended with cm-i and the n-lthcolumn with KTn-iVn-i, and any available locations at which all of Q, K and V are present are completed and summed with earlier accumulated values.

[0092] Finally row 706 shows the final step associated with receipt of Qn, KTn, Vn. In this step the final row and column are appended, completing the attention tensor. At this stage, each row of the final column is completed, and each value is summed with the accumulated value for the respective row (i.e. QiKTiVi through to QiKTn-iVn-i for a given row i) is summed to generate the final output tensor.

[0093] The table of FIG. 5 also illustrates the temporary memory usage of this approach. In the first row 701, the temporary storage is 1 * (1 + 3) * (BN, seq, hs). The (1 + 3) * (BN, seq, hs) part of this expression reflects that there are 3 tensors stored for Qi, Ki, Vi, plus one temporarily stored tensor, each of shape (BN, seq, hs). The temporarily stored tensor is the partial attention tensor 750 discussed above and shown in FIG. 7. The 1* reflects that for this first step there is 1 set of these tensor (based on Qi, Ki, Vi) stored.

[0094] As the QKTV calculation is carried out together within the temporarily stored tensor, this may be the only temporary tensor required, in contrast to the techniques discussed above with reference to FIG. 4, where QKTis calculated in a first step, necessitating the storage of a temporary tensor of shape (BN, seq, seq), followed by subsequently multiplying the temporarily-stored QKTby V.

[0095] In the second row 702, the temporary storage is 2* (1 + 3) * (BN, seq, hs), reflecting that now Qi, Ki, Vi and 2, K2, V2 are stored, along with their respective temporary tensors. The temporary memory usage continues to grow in this linear fashion, such that in the penultimate step of row 705, the memory usage is (n-1) * (1 + 3) * (BN, seq, hs). The memory usage peaks upon receipt of the final subsequence tensors Qn, KTn, Vn. At this point, the storage is the aforementioned (n-1) * (1 + 3) * (BN, seq, hs) + 3 * (BN, seq, hs) for Qn, KTn, Vn. In big O notation, this is O(BN * n * seq * hs), which can be approximated to O(BN * n * seq).

[0096] In the final step set out in row 706, the storage of the previous n-1 steps is released upon calculation of the final output vector, such that the storage is merely (BN, seq, hs).

[0097] In terms of the execution of the examples discussed above, accelerator 200 allocates a working space in accelerator memory 220 for storing the temporary tensor 750. As the values accumulate as each subsequence is received, the previously stored values are effectively overwritten with each new accumulation of previous values.

[0098] As noted above, for simplicity, the examples omit the application of the softmax function and the divisionon the basis that they are not relevant to the shape of the tensors and thus not material to the scaling of memory application. However, in reality, to provide the attention tensor as output, both the division by dkand the application of the softmax function is required. This can be carried out as part of the accumulation process. For example, as each Qj KTjVj becomes available, the Oj KTj is divided by - ef and the softmax function is applied before accumulating it with theprevious value, such that each stored Q.jKTjVj is in fact softmax Vt.\ yd / c /

[0099] FIG. 8 illustrates a second example technique of using Qi...Qn, KTi... KTnand Vi...Vnto calculate the attention score in a manner that is memory efficient. In overview, in this example, the attention tensor is calculated in a distributed manner, by issuing the tensors associated with each subsequence to different devices. In other words Qi, KTi, Vi are issued to a devices

[0100] In more detail, in the example a coordinating computer system 500 comprises a storage 501, which stores the model 300 and the tensors Qi, K\ Vi in a suitable tensor storage 502.

[0101] The system 500 distributes Qi, KTi, Vi to the ithdevice, such that device 100-1 receives Qi, KTi, Vi , device 100-2 receives Q2, KT2, V2 and so on, with device 100-n receiving Qn, KTn, Vn.

[0102] The system 500 also distributes Qi...Qnto each device 100. In other words Qi...Qnare synchronised across devices 100. The respective devices 100 then calculate their respective part of the partial attention tensor 750, and provide their partial attention tensor back to the coordinating computer system 500. This may for example be carried out as a reduction operation, for example using the well-known MPI communication primitives.

[0103] The approaches set out in FIG. 6-7 and FIG. 8 can be combined. In other words, groups of consecutive subsequences may be distributed to different devices 100, with each device processing its respective group of subsequences in turn and the results being combined.

[0104] In more detail, in examples where a group is calculated on a local device (i.e. some, but not all, Qi, KTi, Vi are on the same device), first Qi, KTi and Vi with a groupi are calculated, with the result stored to a temporary tensor. Then, 2, KT2 and V2 are calculated with groupi, and then the temporary tensor is updated. This continues until all Qi KTi and Vi in groupi are calculated. A similar process is carried out on another local device with group2. The results may then be combined to derive QKTV for the entire sequence.

[0105] In the discussion above, reference is made to tokens. Tokens should be broadly understood to be the minimum units processable by the model. Whilst in natural language processing examples the tokens may be words, punctuation marks and the like, the disclosure is not limited to tokens of this nature. In other natural language processing examples, the tokens may be individual characters.

[0106] In yet further examples, the attention calculation may form part of a model that processes data of other modalities, in which case the tokens may be representative of data other than natural language. For example, the model may be an image processing model, and the tokens may be image patches or the like. The image processing model may be for one of the following tasks: image recognition, object detection, segmentation, image super-resolution, video understanding, imagegeneration, text-image synthesis and visual question answering. In other examples, the model may be a model that processes sensor data or audio data. Furthermore, the disclosure extends to multi-modal models, that receive input in one modality (e.g. text) and provide output in another modality (e.g. images).

[0107] The techniques discussed herein are applicable during training of the model. Particularly, during the FP part of training, attention scores may be generated based on input training samples. Equally, the techniques discussed herein are applicable to execution of an attention layer at inference time.

[0108] Advantageously, the above-described examples provide an improved means of managing memory during execution of a machine learning model. For example, the techniques may make optimal use of available accelerator memory, enabling the calculation of attention over very long sequences. These techniques may allow the use of fewer instances of expensive accelerator hardware to train large models with billions of parameters, thus ensuring such techniques can be carried out by smaller businesses or academic institutions.

[0109] For completeness, FIG. 9 schematically shows a non-limiting example of a computing system 1200 that can enact one or more of the methods and processes described above. Computing system 1200 is shown in simplified form. Computing system 1200 may take the form of computer system 100 or 500 described above, or any other computer device discussed herein. Computing system 1200 may take the form of one or more personal computers, server computers, and / or other computing devices.

[0110] Computing system 1200 includes a logic processor 1202, volatile memory 1204, and a non-volatile storage device 1206. Computing system 1200 may optionally include a display subsystem 1208, input subsystem 1210, communication subsystem 1212, and / or other components not shown in FIG. 9.

[0111] Logic processor 1202 includes one or more physical devices configured to execute instructions. For example, the logic processor may be configured to executeinstructions that are part of one or more applications, programs, routines, libraries, objects, components, data structures, or other logical constructs. Such instructions may be implemented to perform a task, implement a data type, transform the state of one or more components, achieve a technical effect, or otherwise arrive at a desired result.

[0112] The logic processor may include one or more physical processors (hardware) configured to execute software instructions. Additionally or alternatively, the logic processor may include one or more hardware logic circuits or firmware devices configured to execute hardware-implemented logic or firmware instructions. Processors of the logic processor 1202 may be single-core or multi-core, and the instructions executed thereon may be configured for sequential, parallel, and / or distributed processing. Individual components of the logic processor optionally may be distributed among two or more separate devices, which may be remotely located and / or configured for coordinated processing. Aspects of the logic processor may be virtualised and executed by remotely accessible, networked computing devices configured in a cloud-computing configuration. In such a case, these virtualised aspects are run on different physical logic processors of various different machines, it will be understood.

[0113] Non-volatile storage device 1206 includes one or more physical devices configured to hold instructions executable by the logic processors to implement the methods and processes described herein. When such methods and processes are implemented, the state of non-volatile storage device 1206 may be transformed — e.g., to hold different data.

[0114] Non-volatile storage device 1206 may include physical devices that are removable and / or built-in. Non-volatile storage device 1206 may include semiconductor memory (e g., ROM, EPROM, EEPROM, FLASH memory, etc.), and / or magnetic memory (e.g., hard-disk drive), or other mass storage device technology. Non-volatile storage device 1206 may include nonvolatile, dynamic, static, read / write, read-only, sequential-access, location-addressable, file-addressable, and / or content-addressable devices. It will be appreciated that non-volatile storage device 1206 is configured to hold instructions even when power is cut to the non-volatile storage device 1206.

[0115] Volatile memory 1204 may include physical devices that include random access memory. Volatile memory 1204 is typically utilised by logic processor 1202 to temporarily store information during processing of software instructions. It will be appreciated that volatile memory 1204 typically does not continue to store instructions when power is cut to the volatile memory 1204.

[0116] Aspects of logic processor 1202, volatile memory 1204, and non-volatile storage device 1206 may be integrated together into one or more hardware-logic components. Such hardware-logic components may include field-programmable gate arrays (FPGAs), program- and application-specific integrated circuits (PASIC / ASICs), program- and application-specific standard products (PSSP / ASSPs), system-on-a-chip (SOC), and complex programmable logic devices (CPLDs), for example.

[0117] The terms "module," "program," and "engine" may be used to describe an aspect of computing system 1200 typically implemented in software by a processor to perform a particular function using portions of volatile memory, which function involves transformative processing that specially configures the processor to perform the function. Thus, a module, program, or engine may be instantiated via logic processor 1202 executing instructions held by non-volatile storage device 1206, using portions of volatile memory 1204. It will be understood that different modules, programs, and / or engines may be instantiated from the same application, service, code block, object, library, routine, API, function, etc. Likewise, the same module, program, and / or engine may be instantiated by different applications, services, code blocks, objects, routines, APIs, functions, etc. The terms "module," "program," and "engine" may encompass individual or groups of executable files, data files, libraries, drivers, scripts, database records, etc.

[0118] When included, display subsystem 1208 may be used to present a visual representation of data held by non-volatile storage device 1206. The visual representation may take the form of a graphical user interface (GUI). Because the herein described methods and processes change the data held by the non-volatile storage device, and thus transform the state of the non-volatile storage device, the state of display subsystem 1208 may likewise be transformed to visually represent changes in the underlying data. Display subsystem 1208 may include one or more display devices utilising virtually any type of technology. Such display devices may be combined with logic processor 1202, volatile memory 1204, and / or non-volatile storage device 1206 in a shared enclosure, or such display devices may be peripheral display devices.

[0119] When included, input subsystem 1210 may comprise or interface with one or more user-input devices such as a keyboard, mouse, touch screen, or game controller.

[0120] When included, communication subsystem 1212 may be configured to communicatively couple various computing devices described herein with each other, and with other devices. Communication subsystem 1212 may include wired and / or wireless communication devices compatible with one or more different communication protocols. As non-limiting examples, the communication subsystem may be configured for communication via a wireless telephone network, or a wired or wireless local- or wide-area network. In some embodiments, the communication subsystem may allow computing system 1200 to send and / or receive messages to and / or from other devices via a network such as the Internet.

[0121] Although at least some aspects of the embodiments described herein with reference to the drawings comprise computer processes performed in processing systems or processors, aspects of the invention also extend to computer programs, particularly computer programs on or in a carrier, adapted for putting aspects of the invention into practice. The program may be in the form of non-transitory source code, object code, a code intermediate source and object code such as in partially compiled form, or in any other non-transitory form suitable for use in theimplementation of processes according to aspects of the invention. The carrier may be any entity or device capable of carrying the program. For example, the carrier may comprise a storage medium, such as a solid-state drive (SSD) or other semiconductorbased RAM; a ROM, for example, a CD ROM or a semiconductor ROM; a magnetic recording medium, for example, a floppy disk or hard disk; optical memory devices in general; etc.

[0122] The examples described herein are to be understood as illustrative examples of embodiments of the invention. Further embodiments and examples are envisaged. Any feature described in relation to any one example or embodiment may be used alone or in combination with other features. In addition, any feature described in relation to any one example or embodiment may also be used in combination with one or more features of any other of the examples or embodiments, or any combination of any other of the examples or embodiments. Furthermore, equivalents and modifications not described herein may also be employed within the scope of examples of the invention, which is defined in the claims. The labels "first", "second", "third" etc. herein are merely used to distinguish between similar elements and need not imply a particular ordering.

Claims

28CLAIMS1. A computer-implemented method of processing an attention layer of a neural network, the method comprising:receiving an input sequence comprising a plurality of tokens;splitting the input sequence into a plurality of subsequences;generating respective subsequence query tensors, subsequence key tensors and subsequence value tensors for each of the plurality of subsequences;providing, to a processing device, a first subsequence query tensor, first subsequence key tensor and first subsequence value tensor associated with a first subsequence of the plurality of subsequences;calculating, by the processing device, a partial attention tensor based on the first subsequence query tensor, first subsequence key tensor and first subsequence value tensor;iteratively:providing a next subsequence query tensor, next subsequence key tensor and next subsequence value tensor associated with a next subsequence of the plurality of subsequences to the processing device;updating the partial attention tensor based on the next subsequence query tensor, next subsequence key tensor and next subsequence value tensor; andupon receipt at the processing device of a final subsequence query tensor, final subsequence key tensor and final subsequence value tensor associated with a final subsequence of the plurality of subsequences, outputting by the processing device a final attention tensor.

2. The method of claim 1, wherein updating the partial attention tensor for an ithiteration comprises:calculating an ithattention score from the ithsubsequence query tensor, ithsubsequence key tensor and ithsubsequence value tensor; andstoring the ithattention score in a location of the partial attention tensor corresponding to an ithrow and an ithcolumn.

3. The method of claim 2, wherein updating the partial attention tensor further comprises:for every location in the ithrow before the ithcolumn, calculating an attention score based on a subsequence key tensor and subsequence value tensor stored in a previous iteration and the ithquery tensor, and updating the location to store the attention score;for every location in the the ithrow after the ithcolumn, storing the ithsubsequence query tensor.

4. The method of claim 3, wherein updating the partial attention tensor further comprises:for every location in the ithcolumn before the ithrow, calculating an attention score based on a subsequence query tensor stored in a previous iteration and the ithsubsequence key tensor and ithsubsequence value tensor, and updating the location to store the attention score;for every location in the ithcolumn after the ithrow, storing the ithsubsequence key tensor and ithsubsequence value tensor.

5. The method of claim 2 or any claim dependent thereon, wherein updating the partial attention tensor comprises accumulating the ithattention score with previous attention scores stored in the ithrow of the partial attention tensor.

6. The method of any preceding claim, wherein:a shape of each of the subsequence query tensors, subsequence key tensors and subsequence value tensors is (BN, seq, hs), wherein BN represents a training batch size multiplied by a number of attention heads, seq represents a length of each subsequence and hs represents a number of attention heads; andthe method comprises allocating a memory region of the processing device at an ithiteration having a size i * 4 * (BN, seq, hs).

7. A computer-implemented method of processing an attention layer of a neural network, the method comprising:receiving an input sequence comprising a plurality of tokens;splitting the input sequence into a plurality of subsequences;generating respective subsequence query tensors, subsequence key tensors and subsequence value tensors for each of the plurality of subsequences;providing, to each of a plurality of processing devices, a respective subsequence query tensor, subsequence key tensor and subsequence value tensor;calculating, by each processing device, a respective partial attention tensor based on the provided subsequence query tensor, subsequence key tensor and subsequence value tensor;receiving, from each processing device, the respective partial attention tensor; andgenerating a final attention tensor from the received partial attention tensors.

8. The method of claim 7, wherein calculating the partial attention tensor for an ithdevice comprises:calculating an ithattention score from the ithsubsequence query tensor, ithsubsequence key tensor and ithsubsequence value tensor; andstoring the ithattention score in a location of the partial attention tensor corresponding to an ithrow and an ithcolumn.

9. The method of any preceding claim, wherein the attention layer is part of a transformer block of the neural network.

10. The method of any preceding claim, carried out during training of the neural network.

11. The method of any of claims 1 to 9, carried out at inference time of the neural network.

12. The method of any preceding claim, wherein the processing device is a work accelerator.

13. The method of any preceding claim, wherein the attention layer is a self attention layer.

14. The method of any preceding claim, wherein the neural network is an image processing neural network.

15. A computer system comprising a processor and a memory storing instructions, which when executed, implement the method of any preceding claim.

16. A computer-readable storage medium storing instructions, which when executed, cause a computer system to implement the method of any of claims 1 to

Images