Coding LLaMA 2 from scratch in PyTorch - KV Cache, Grouped Query Attention, Rotary PE, RMSNorm

Umar Jamil · Advanced ·🧠 Large Language Models ·2y ago

About this lesson

Full coding of LLaMA 2 from scratch, with full explanation, including Rotary Positional Embedding, RMS Normalization, Multi-Query Attention, KV Cache, Grouped Query Attention (GQA), the SwiGLU Activation function and more! I explain the most used inference methods: Greedy, Beam Search, Temperature Scaling, Random Sampling, Top K, Top P I also explain the math behind the Rotary Positional Embedding, with step by step proofs. Repository with PDF slides: https://github.com/hkproj/pytorch-llama Download the weights from: https://github.com/facebookresearch/llama Prerequisites: 1) Transformer explained: https://www.youtube.com/watch?v=bCz4OMemCcA 2) LLaMA explained: https://www.youtube.com/watch?v=Mn_9W1nCFLo Chapters 00:00:00 - Introduction 00:01:20 - LLaMA Architecture 00:03:14 - Embeddings 00:05:22 - Coding the Transformer 00:19:55 - Rotary Positional Embedding 01:03:50 - RMS Normalization 01:11:13 - Encoder Layer 01:16:50 - Self Attention with KV Cache 01:29:12 - Grouped Query Attention 01:34:14 - Coding the Self Attention 02:01:40 - Feed Forward Layer with SwiGLU 02:08:50 - Model weights loading 02:21:26 - Inference strategies 02:25:15 - Greedy Strategy 02:27:28 - Beam Search 02:31:13 - Temperature 02:32:52 - Random Sampling 02:34:27 - Top K 02:37:03 - Top P 02:38:59 - Coding the Inference

Full Transcript

hello guys welcome to my new coding video in this video we will be coding llama 2 from scratch and just like the previous video in which I was I was coding the Transformer model from from zero while coding I will also explain all the aspects of llama so all the each building block of the Llama architecture and I will also explain the math behind the rotary positional encoding I will also explain grouped query attention the KV cache so we will not only have a theoretical view of these Concepts but also a practical one if you are not familiar with the Transformer model I highly recommend you watch my previous video on the Transformer model and then if you want you can also watch the previous video on how to code a Transformer model from xero because it will really help you a lot in understanding what we are doing in this current video if you already watched my previous video on the architecture of llama in which I explain the concepts it will be also really helpful if you didn't it's I will try to explain all the concept but not as much in the as as in the in detail as in the last video so please if you have time if you want you please watch my previous video on the Llama architecture and then watch this video because this will really give you a deeper understanding of what is happening and um let's review uh Lama the Llama architecture here we have a comparison between the architecture of the standard Transformer as introduced in the attention is all you need paper and the architecture of llama the first thing we notice is that the Transformer was at encoder decoder model and in the previous video we actually trained it on a translation task on how to translate for example from English to Italian while Lama is a large language model so la the the goal of a large language model is actually to work with the what is called the next token prediction task so given a prompt the model tries to come up with the next token that completes this prompt in the most coherent way so in the in a way that it makes sense the answer and we keep ask asking the model for the successive tokens based on the previous tokens so this is why it's called a causal model so each output depends on the previous tokens which is also called The Prompt contrary to what I have done to in my previous video on the coding the Transformer model in this video we will not start by coding the single building blocks of llama and then come up with a bigger picture but we will start from the bigger picture so we will first make the skeleton of the architecture and then we will build each block I find that this is a better way for explaining llama also because it's the model is simpler even if the single building blocks are much more complex and this is why it's better to first look at how they interact with each other and then zoom in into the their inner workings let's start our journey with the embeddings so this block here let me use the laser so this block here so we are given an input and we want to convert it into embeddings let's also review what are embeddings these are my slides from my previous video on the Transformer and as you can see we start with an input sentence so we are given for example the sentence your cat is a lovely cat we tokenize it so we split into single tokens we match we map each token into its position in the vocabulary the vocabulary is the list of all the words that our model can recognize these tokens actually most of the time are not single words and what I mean is that the model doesn't just split the word by white space into single words usually the most commonly used tokenizer is the bpe tokenizer which means byte pair encoding tokenizer in which the single tokens can also be a sequence of letters that are not necessarily mapped to a single word it may be a part of a word or maybe it may be a white space it may be multiple words with maybe a single digit Etc and the embedding is a mapping between the number so the input IDs that represent the position of the token inside of the vocabulary to a vector this Vector in the original Transformer was of size 512 while in llama in the base model so the 7 billion model it's a 4096 uh the dimension is 4 times 2096 which means this Vector will contain 4096 numbers each of these vectors represents somehow the meaning of the word because each of these vectors are actually parameter vectors that are trained along the model and somehow they they capture the meaning of each word so for example if we take the word cat and dog they will have an embedding that is more similar compared to cat and the tree if we compare the distance of this two vectors so they we could say the euclidean distance these two vectors and this is why they are called embedding this is why we say they capture the meaning of the word so let's start coding the model we open we open I will open Visual Studio code and first of all we make sure that we have the necessary libraries I will also share the repository of this project so the only thing we need is Storage Center space which is the tokenizer that we use in Lama and tkodm the second thing is from the official repository flama you should download the download script it's this file download.sh that allows you to download the weights of the Llama model in my case I have downloaded The Dilemma 27 billion along with the tokenizer and this is the smallest model actually and I will not even be able to use the GPU because my GPU is not powerful enough and so I will run the model on the CPU I think most of you guys will do the same because it's actually a little big for a normal computer unless you have a powerful GPU so let's start coding it let's create a new file model Dot pi and let's start our journey import the necessary things so we import torch import and also george.nn okay these are the basic thing that we always import I remember we also need math and then also data classes and this is all the Imports we need most of the code is actually based on the original llama code so you don't be surprised if you see a lot of similarities but I simplified a lot of parts to remove things that we don't need especially for example the parallelization and um I also try to add a lot of comments to show you all the shapes changed in each tensor and that's it so let's start so the first thing I want to create is the class that represents the parameters of the model so here we already see that we have two type of heads one is the number of heads for the queries so number of heads for the queries and here we have the number of heads for the K and the V so the the keys and the values because with them we will see later with the grouped query attention we don't have to necessarily have the same number of heads for the query key and values like in the original Transformer but we can have multiple number of heads and we will see why and how they work this will be set when we load the tokenizer uh these two parameters indicate the dimension the hidden dimension of the FF and lay so they feed forward the layer the basic idea is that they try to when they introduce the grouped query attention they try to keep the number of parameters because with the grouped query attention we reduce the number of heads of the K and B but they incremented the number of parameters of the feed forward layer so as the number of the total parameters of the model Remains the Same this allows to to compare the full um the base Transformer so with the old heads for the query the key and values with the one they use in Lama which has a reduced number of heads for the knb but this is just a decision in an architectural decision then we have some EPS this is a number that is very small and we will see why we need it oh my God and these are all the parameters we need also here we have two parameters that we will later use for the KV cache and I will explain you later what is it and how it works let's start as I saw as I said before let's start with implementing the skeleton of the entire model and then we Implement each Single part and while implementing each Single part I will also review the the background and how it works and the maths behind it this is the main class that will represent the entire model so all the model we can see here so all this model here except for the soft Max we make sure that we have set the vocabulary size save the values this is the number of layers of the model so this is represents this block here is repeated many times one after another just like in the Transformer the base Transformer if you remember correctly this block here and this block here were repeated one after another many times and here it's repeated 32 times and the last the output of the last layer is then sent to this RMS Norm then to the linear Etc I'm using the same names as used in the transformer in the code from from the Llama repository because when we will load the weights of the model the names must match otherwise the torch doesn't know where to load the weights in and this is the reason I'm trying to keep the same name I only change some names to make them more clear but most of the other names are same this is the model list so this is the list of the layers we will create later the encoder block which is each of these blocks here this is the encoder block for now we just create the skeleton so we have a list of these blocks then we have a normalization and the normalization is the RMS normalization we will implement it later we need to tell him the size of the features and the EPS is a very small number that is needed for the normalization calculation so that the we never divide it by zero basically and then we have the output layer okay then we need to pre-compute the frequencies of the rotary um positional encodings so let's do it I created this method and then we go to implement it and I will show you how it works let me check I think we have a parenthesis okay this is the base Transformer model so first of all we have n layers we have first of all the token at the input embeddings so we convert the the input into embeddings then we pass it through a list of layers the last layer output is sent to the normalization then to the output so this logic will be more clear in the forward method okay here you will see one thing that is different from the previous Transformers is that the sequence length that we want is always one and this is because we are using the KV cache and we will see why so in the previous uh let's review here here when we give the input we want to give the prompt and the model will give us a soft Max of the next token but with the KV cache we don't need to give all the previous tokens because maybe we already computed them in the previous iteration so we only need to give the latest token and then the model will put the next token while the intermediate tokens so the the cache of the previous tokens will be kept by the model in its cache because here we have a KV cache but we will see this mechanism later so for now just remember that the input here that we will get is one token at a time so we will get a batch with sequence length okay so embed size as you can stand this tokens dot shape so and we make sure that the sequence length is actually one and second thing is this model is only good for inferencing not for training because for training of course we need to not have the KV cache and we only we need to be able to process multiple tokens but our goal is actually to use the pre-trained Llama weights so we convert the tokens into embeddings as you can see we had the dim Dimension so the dimension of the embeddings which is 4096 for the base model but depending on the model size it can be different okay I promise I will explain this line here and this line here much more in detail just in two minutes let me finish it and I will explain everything together so this is basically we are we pre-compute something about the positional encoding that we then give to the successive layers but uh let's finish writing it and then we I explain expand this method and this one and everything will be clear so suppose we with this one we retrieve something that is needed for computing the positional encoding which then we feed to the next layers so we constructively apply all the encoder blocks and finally we apply the normalization just like here so we apply these blocks one after another many times then we apply the normalization and then we calculate the output using the linear layer and finally we return the output so this is the skeleton of the model so we take the input we convert it into embeddings this part I will explain later we give this input embeddings with something about the positional encodings to these blocks one after another we take the output of the last one and give it to the RMS Norm we take the output of the RMS Norm give it to the linear layer and then during the inference we will apply the softmax now let's concentrate on the positional encodings let's first review how they worked in the original Transformer as you remember in the original Transformer so the transform and the attention is all you need we first take the sentence we convert into embedding vectors so if the vectors of size 512 we then add another Vector which has the same size so 512 that represents the position of that token inside the sentence so every sentence every token in the first position of a sentence will receive this Vector every token in the second position of a sentence will be have this Vector added to it and every token in the third position of a sentence will have this Vector added to it these vectors are pre-computed because they only depend on the position not on the word they are applied to and this is why they are called absolute position encoding because they only strictly depend on the position of the word inside of the sentence while in the contrary in the rotary positional embeddings they are a little different let's go have a look first of all the rotary positionally encoding them or embeddings they are computed right before the calculation of the attention and they are only applied to the Q and the K matrices not to the V let's see why the first thing we need to understand is the difference between absolute positional encodings and relative ones the absolute positional encodings are the one we just saw so they they deal with one token at a time and each token gets its own embedding while the relative positional embeddings they come into play during the calculation of the attention and the calculation of the attention is basically done using the dot product because it's a query multiplied by the transpose of the keys divided by the square root of the model so there is a DOT product in between and that we want with the relative positional encodings we we change this dot product so that we introduce a new Vector that indicates the distance between the two tokens involved in the dot product so for example in the original Transformer we have this formula here so query multiplied by the transpose of the keys divided by the square root of the model while in the relative positional encodings which are not the one used in llama so this is just an introduction in the relative positional encodings we have another Vector here that represents the distance between these two tokens and we compute the attention mechanism like this the rotary position embeddings the one that are used in Lama are something in between the absolute and the relative absolute because they each token will get its own embedding but relative because the attention mechanism will be evaluated using the relative distance between two token let's see the relative position the rotary position embeddings were introduced in the in the paper uh from this company J and uh the authors of this paper they wanted to find an inner product that works like this so first of all what is the inner product we are we are all familiar with the dot product the inner product can be thought of as a generalization of the dot product so it's a fun it's a an operation that has some properties that reflect what is the dot product so the authors of this paper wanted to find an inner product between the two vectors the query and key such that they only depend on them so this is the symbol for inner products so this inner product only depends on the embedding of the two tokens involved so XM and xn and the relative distance of these two tokens so the distance between them for example if the first token is in position two and the second position token is in position five so m equal to n equals five the distance between two will be 3 or -3 according to the order and they wanted to find a DOT product that has this property that only depends on the embedding of the first token of the on the embedding of the second token and the relative distance between them then this so that if this function G is built in this way then we achieve that objective that is we take the first token so the query for example we multiply it by the W Matrix this is actually download so in the vanilla Transformer but okay Suppose there is no W Matrix here we convert it into a complex number in this form we take the key Vector we transform into a complex number into this form and we Define the inner product in this way this inner product will basically depend only on the distance between these two tokens so they wanted to find an encoding mechanism such that the attention mechanism which is based on a DOT product which is an inner product behaves like this so it only depends on the embeddings of the vector and the distance between them and if we for example this formulation here if we apply it on a vector of Dimension two so we think of embedding of with only two Dimensions it becomes in this form here this is due to the errors formula so each complex number thanks to the ls formula can be written as the sine plus a the cosine plus assign and this Matrix here reminds us of the rotation Matrix let me give you an example suppose our original Vector is here and if we multiply this vector v0 by this Matrix here the resulting Vector will be rotated by the angle Theta so this is why they are called rotary positional embeddings because the Matrix here represents a rotation of the vector so when we want to visualize how the rotary position embedding work we have to think that they will map it into a vector space and they will rotate each word to an angle that is a multiple of a base angles of theta and for proportional to the Theta angle proportional according to its position so that's two two tokens that occupy similar positions will have similar inclination and the two tokens have different positions will have different inclinations and this is the idea behind the rotary positional embeddings but how do we actually compute them into um in the in the code in the pythorch well to compute them we need to build a matrix like this and actually as you can see this Matrix is actually full of zeros and so when we calculate the embedding in this way we will do if we do a matrix multiplication we will be doing a lot of operations that are useless because most of these items are zero so the authors of the paper proposed another form that is more computationally efficient and this form basically says that we take the embedding of the vector to which we want to apply the positional encodings so for example this one this is a vector so the First Dimension the second dimension the third dimension and the last Dimension so if this is for example the vanilla Transformer this would be x d which should be 512 we multiplied element Wise by this Matrix plus another Vector that is actually based on this Vector but with positions and the signs changed so this is actually in the first position we have the second dimension with its sign change in the second position we have actually the first dimension in the third position we have the fourth dimension but with sine chain XL so it depends only on the embedding of the word but in a com changes with the science and position change and then we multiply this element wise with another Matrix that you can see here this vector then this will be the encoding of the token we are talking about and now what we can pre-compute is this Matrix here because it doesn't depend on the token to which we apply it to and this Matrix here because it doesn't depend on the token we apply it to and it did depend on M so it's the position of the word and Theta what is Theta Theta is a series of numbers defined like this and so let's first build the code to pre-compute this and this here let's do it I will first write the code and later I will show you how it works this data parameter 10000 comes from the paper it's written here ten thousand we first need to make sure that the dimension of the word to which we are applying the embedding is actually even because we can in the paper it's written that the this rotary positional encodies cannot be applied to an embedding which has a odds odd Dimension so cannot be 513 must be 512 or 514 or any other even number and this is first written in the paper even okay now we build the Theta parameters which is a sequence and the shape of this data will be head Dimension divided by two because we will apply this embeddings to each head so not right after the embedding but after we have split them into multi-head so each token the token of each head we check the size of the dimension of each head and we divide it by two because y divide bit by 2 because in the paper they also divide it by two so D divided by 2 here okay so what's the formula here the formula is Theta of I is equal to 10 000 to the power of minus 2 multiplied by I minus 1 divided by dimension for I equal to 1 2 Etc up to Dimension divided by 2. so now we are Computing this part here so this part here which is a series so I here it starts from 1 we will start from zero so we don't have to do I minus 1. and Theta is equal to 1 over the ta so ten thousand to the power of Theta numerator divided by that dimension why do we do 1 over Theta well because this is to the power of minus two so something to the power of a negative number is 1 over that something to the power of the positive exponent and then this will result in a matrix with shape head Dimension divided by two so shape is head Dimension divided by two now we construct the positions what are the positions because we want to build these two matrices they depend on Theta so the series of theta that goes from Theta 1 to Theta Dimension divided by 2 and that we already have now we need to build the M's because the amps the possible positions of a token can be many we basically give as input to this function the maximum sequence length that we can afford multiplied by two because we can we have also the prompt which may be long so we say okay let's pre-compute all the possible Theta and M for all the possible positions that our model will see and the all the possible positions is given by this parameter sequence length so now construct the positions and the shape is sequence length which is m now we need to multiply M by all the sequence of thetas but each M with all the Theta so for example if we have m equal to 1 we need m one Theta 1 m 1 Theta 2 m 1 Theta divided D divided by two then we need M2 Theta 1 m 2 Theta 2 m 2 Theta 3. so for that we will use a outer product the outer product I will show you later basically means multiply all the elements of the first Vector with all the elements of the second vector all the possible combinations so so for example here we have a frequency is equal to torch other product um and Theta okay so what we are doing is we are doing the auto the other product within M which is the positions multiplied by Theta this will basically take the first element of the first vector and multiply with all the elements of the second Vector then take the second element of the first vector and multiply it with all the elements of the second Vector etc etc so if we start with the shape let me say shape off m is the sequence length let's say outer product with head Dimension divided by two this will result in a tensor of sequence length by add Dimension divided by 2. so for each position we will have all the Theta then for the second position we will have all the data for the third position we will have all the Theta and so on now we want to write these numbers into a complex form and I will show you why I multiplied by m multiplied by Theta where R is equal to 1 as follows so we compute let me also write the shape and then I explain you how it works this is a here too okay let's write some formulas I could also you know not explain all the proofs so I know this in the next few minutes will be a little boring because I will be explaining all the math behind it but of course I don't think just you just like me you like to watch just some code and say okay this is how it's done no I like to actually give a motivation before behind every operation we do and that's I think one of the reasons you are watching this video and not just reading the code from The Meta repository so um let's do some math the first thing we need to review is how complex numbers work okay a complex number is a number in the form a plus I multiplied by B where a is called the real part and the B is called the imaginary part and I is a number such that I to the power of 2 is equal to -1 so the complex numbers were introduced to represent all the numbers that involve somehow the square root of a negative number as you know from school the square root of in every number cannot be calculated but so that's why we introduce this constant I which is the negative number which is the square root of -1 and so we can represent the square root of negative numbers and they can also be helpful in Vector calculations and we will see how because we have the Euler's formula the Euler's formula says that e to the power of I multiplied by X is equal to cosine of x plus I multiplied by sine of x so it allow us to represent a complex number in the exponential form into a sum of two trigonometric functions the cosine and the design and this will be very helpful later because our goal our goal is to calculate this matrices here the cosine of M Theta and the Sines of M Theta and the first thing we did is we calculated all the Theta 1 then we calculated all the positions then we calculated all the possible possible combinations of positions and thetas so what we did is we calculated a vector that represents the Theta so Theta 1 Theta 2 up to Theta D divided by 2 then we calculated all the possible M's M can be 1 can be 2 can be whatever so sequence length let's say then we calculated the product of each of them which all of for each for all the possible data so for example we created a new Matrix that for has M1 Theta 1 m 1 Theta 2 and 1 theta 3 up to m up to um 1 Theta D divided by 2. and then M2 Theta 1 m 2 Theta 2 and 2 theta 3 etc etc until m 2 Theta D divided by 2. these numbers are still not complex number they're just real numbers because Theta is a real number m is a real number but they are not complex numbers then we convert them into complex numbers so what we do with the last operation here this one here we convert each of these numbers into Polar into its Polar form a number in polar form is a number that can be written as R multiplied by e to the power of I Theta which can be written as R cosine of theta plus I sine of theta y because it can be represented in the graphical let's say graphical plane x y as you know complex numbers can be represented into the 2D plane X Y where the real part is on the X and the imaginary part is on the Y so we are actually representing a vector of size let's say r with an inclination of theta because as you know the projection of this Vector on the real part is R cos theta plus I um the the the projection on the y axis is sine of theta and here I forgot or yeah I forgot an R here or scene of theta so this is a another way of representing complex numbers and what we are doing is we are calculating this Matrix and then converting all these numbers into their complex form so we are converting it into another Matrix that has r equal to one and this number here for example this item here will become cosine of um one Theta 1 plus I sine of m 1 Theta 1. this number here will become another number so this is only one number this is become another complex number that is the cosine of m 1 Theta 2 plus I sine of m 1 Theta 2 etc etc for all the because we are not increasing the numbers the over the total numbers this shape of the tensor also doesn't change it just becomes a more complex number so instead of having M Theta 1 it becomes cosine of M Theta 1 plus I M Theta 1. why do we need this form here because we need signs and cosines and later we will see how we will use them now the point is imagine we are given a num a vector because we want to apply this positional encodings to a vector so how to apply them because the vector we will will be giving us as a list of Dimensions from X1 to the last Dimension just like in the original Transformer we have a vector of size 512 in this case it will be much smaller because it's the dimension of each head and as you remember each head doesn't watch the full dimension of the vector embedding Vector but a part of it so but for us okay imagine it's only one head so if it's only one head we will watch the full Dimension so for now don't consider the the multihead just suppose that we only are working with one head so imagine we have a token with its full Dimensions so in the case of llama 4096 dimensions and in the case of the vanilla Transformer 512 how to apply it let's do some math actually let's do some more math so we are given we are given suppose a smaller embedding Vector because we want to do the calculation and not go crazy otherwise 4096 is a little difficult to prove I want to make a list a list of operations on this Vector until we arrive to this form here so let's start suppose our embedding Vector is only made of four dimensions X1 X 2 X 3 and X 4. okay the first thing we do is I will do some Transformations and that will later translate them into code so for now just follow the Transformations I'm doing this is the transformation number one I want to group successive tokens successive Dimensions so into another dimension so X1 and X2 become another dimension in this tensor and X3 and X4 become oops very badly written x 4 become another dimension in this tensor the total number of items is still four but I added another dimension and this has a size 4 by 1 right this one has two by two by one so I split it into multiple tensors and okay now this next thing I do I consider this first number of this part to be the real part of the complex number and this one to be the imaginary part of the complex number and the same for the second Vector here so I do another transformation that we will call to in which X1 plus I X two and then X3 plus I X 2. this Vector has less items because now two numbers become one complex number now I multiply this element wise with the vector that we pre-computed before as you remember before we pre-computed this one cosine of M1 Theta 1 plus I of M1 Theta 1 cosine of M1 Theta 2 plus I because they suppose this position this token here suppose his position is M1 because we need also the m so this suppose this token here his position is M1 so we take all this row here M1 and this will become our new Matrix here we only have four dimensions so four dimensions means we have a Theta 1 and Theta 2 because D divided by 2 until D divided by 2 so element wise with cosine of m 1 Theta 1 plus I sine of m 1 Theta 1 and then we have cosine of M1 Theta 2. plus I of sine of M1 Theta 2. now we have a element wise product between the first item of this Matrix and the first item of this Matrix actually they are two vectors and then we have the the product of two complex numbers this complex number here and this complex number here so let's do let's see how to compute the product of two complex numbers because I don't want to write very long Expressions I will call this one F1 so F1 is the cosine of M1 Theta 1 and F2 is the sine of M1 Theta 1. and for the same reason I will call this one F3 and F 4. now let's compute the product of the first item of this vector and the first item of this vector so X 1 plus i x two multiplied by f 1 plus i f 2. this is equal to x 1 f 1. plus I X One F2 Plus i x 2 F1 then we have this product i x 2 multiplied by i x 2 but it will become I squared X2 F2 I squared we know it's equal to -1 so it will become minus X2 F2 this one can then be written as real part so all the terms that don't have I x 1 f 1. minus x 2 f 2. Plus I that multiplies X1 F2 x 1 f 2 plus X2 F1 okay this is how to compute the product of two complex numbers let's do it so the first number here in the resulting Matrix from this element wise multiplication will be X1 F1 x 1 F1 minus X2 F2 X2 F2 Plus I of X1 F2 Plus x 2 F1 the second element we don't need to do do this multiplication because they have the similar structure as the first one so we just change the X1 with X3 X2 with X4 this is X4 and F1 with F3 and F2 with F4 so the resulting Matrix will be x 3 F3 minus X4 F4 Plus I x three F4 Plus x 4 F3 this one can then be split back so this complex number we can split the real part and the imaginary part and this we will call it transformation number three so we can split it into two in a tensor of two Dimensions one is the real part and one is the complex part where is X1 f 1. minus X 2. f 2. then X 1 f 2. Plus x 2 F1 this is the first tensor the second tensor will be X3 F3 minus X 4 F 4 and the second will be X3 F4 plus x 4 F3 I'm really sorry for the bad handwriting but it's my touchpad is raw not so good then we do another transformation that in which we flatten all these values so this will be the first item this the second the third and the fourth so we remove this Dimension the inner dimension so we flatten this Matrix and it will become X1 F1 minus X to F2 the second items will will be X1 F2 Plus X 2 x 1 F2 plus 2X 2 F1 then we have x 3 f 3. minus x 4 F4 then we have X3 F4 Plus X 4 f 3. let's compare this resulting Matrix with what is in the paper so let's compare it with this one oops let me Zoom a little bit okay the resulting Matrix is exactly the same as what is in the paper so X1 multiplied by F1 which is the cosine as you remember M1 Theta 1 so X1 multiplied by F1 plus minus X2 so minus X2 multiplied by F2 which is the sign as you can see here sine of M1 Theta 1 here it's not M1 because it's for the generic M but we said m equal to M1 and the second dimension is also correct so it's X1 F2 so X1 um with X1 here because here is the sum so the order doesn't matter so X1 F2 so X1 multiplied by The Sign Plus X2 F1 X2 F1 and um let me check if we can use Okay the third dimension is X3 F3 so X3 multiplied by the cosine minus X4 sine minus X4 sine then we have X3 F4 so X3 F4 F4 is the sign of theta 2. and plus X4 F3 X4 F3 F3 is the cosine of theta 2 in also in this case because we have there are we have the sum here inside the the order doesn't matter so as you can see we started with a vector of Dimension 4 but it could be of Dimension n and we multiply we did some transformation we then multiplied with the Matrix that we pre-computed here then we did some other transformation and the end result is exactly as doing this operation so now let's translate this into code because this is actually what we need to apply the embedding Vector to this Vector here so to this token how to apply the embeddings the rotary position Bendix through this series of Transformations so I could have also written the code and not tell you anything but I like to give proof to what I do so that you know that what I'm doing is actually um actually described in the paper and we are actually doing it according to the paper there is also a visualization in the paper that is really helpful so what we did here for example is we transform the embedding Vector into split it into a new tensor which has half Dimension but Group by grouping two consecutive Dimensions so the two consecutive Dimension X1 and X2 and X3 and X4 plus then we multiply it we transform it to a complex number we multiplied with M Theta that we pre-computed and this visualization of why we do this is present in the paper X in particular at this figure here so here they say if you have a word with n Dimensions you need of course D of D dimensions then of course you will have a Theta of D half Theta because we have a Theta 1 Theta 2 up to Theta D half we group successive Dimensions into a new complex number that if we project it on the complex plane it will result into this Vector so the X1 X2 Vector you can see here and then we multiply it with the complex number M Theta 1. this will result in the number being rotated by the angle indicated by m Theta 1 and this is the encoded number is the encoded token and this is exactly what we are doing with our Matrix Transformations that I show you right now now let's translate this into code apply rotary embeddings X is the token to which we want to apply the Rotary embeddings fixed complex is the output of this function but only for the position of this token so only for all the positions of this token because this will have all the all the data for all the possible positions but we only need the positions for this particular token and then we need the device the first thing we do is the transformation number one I think I call it yeah this one here and number one and number two so the first thing we do is we transform the two consecutive Dimensions into a new tensor and then we visualize it as a complex number these operations are supported by pi torch so we do them so we create X complex is equal to okay this operation here is basically saying take two consecutive dimensions and group them and then we transform this intermediate answer into a complex tensor by using the view as complex operation from torch let me write some comments so we are starting from B sequence length h head Dimension because I saw that before this x is actually not the original Vector but it's already the one divided into with its head Dimension because we we will have a multi multi-head attention but if there is no multi head attention then this head Dimension is actually the full dimension of the token so 4096. then we um we have this tensor here but this tensor has two Dimensions less than this one it doesn't have the batch Dimension and it doesn't have the head Dimension so we need to add it so take the fixed complex and we added the two Dimensions that it's missing and here we are doing okay let me write all the transformations here we are going from here to divide by 2y divide by 2 because every two consecutive pairs are becoming one complex number and here we go from sequence length to head Dimension divided by two we are mapping it to 1 because this is the batch Dimension sequence length then the head Dimension one and then head Dimension divided by two now we multiply them together so we do this operation here so element wise multiplication which will result in a rotation as we saw in the figure before so that's why I call it X rotated is equal to X complex multiplied by the complex number of the frequencies in this case we are doing sequence length Edge at Dimension divided by two then we we multiplied we obtain this result and then we first transform it into um we transform the complex number into a tensor in which the first item is the real part of the complex number and then the complex part the imaginary part and then we flatten it so let's do it this operation view as real will transform the tensor like this so it's this transformation here in which we transform the complex number into a tensor of two dimensions because that's why you can see this additional Dimension here and then we flatten it you can just say to flatten it with the shape of the original um with the original tensor we gave it so become and this is how we calculate the embedding so given a tensor of representing a token or a list of tokens because we have the batch dimensions we can apply the embeddings like this during doing all these Transformations that we have done here and that are represented in this code and they are all equivalent to doing this operation as written on the paper now we need to go forward with our Transformer by implementing the rest the next thing that we can Implement is this RMS Norm because it's present at the output of the of the Transformer but also at the input so let's go review again the architecture we can see that we have the normalization the RMS normalization here but we also have it here and here so let's implement it let's also visualize how the RMS Norm works if you want to have a deep understanding of how normalization Works in my previous video about Lama I actually described why we need normalization how it was historically done and how it works also at the auto grad level so I will not repeat the same lecture here I will just briefly introduce how it works but if you want to have a better understanding please watch my previous video so as you remember in the original Transformer we use the layer normalization a layer normalization worked like this we have an input where we have some items suppose item one item two up to item 10 each item has three features so A1 a2a3 what we did with the layer normalization we computed two statistics one for each item so mu and sigma so the mean and the C Sigma and we standardize each item normalize each element of this input Matrix using this formula here which transforms it into a distribution with zero mean and variance of one and this formula comes from probability statistics so as you know if you have any random variable with it's a move and sigma if you do the variable minus its mean divided by the standard deviation so the square root of the variance it will result into a gaussian of mean 0 and standard and the variance of one we then multiply this with the gamma parameter and we also add a beta parameter here but this was done in the layer normalization in the lamba we use RMS normalization and let's see the difference in RMS normalization the paper of the RMS normalization claims that we don't need to obtain the effect of layer normalization we don't need to compute two statistics that is the mean and the variance and actually they claim that the normal the effect given by Larry normalization can be obtained without re-centering the values so without re-centering them around the mean of zero but just by scaling however the variance in the layer normalization was computed using the mean because if you remember the formula of the variance is x minus the mean of the distribution to the power of 2 divided by n so um to compute the variance we needed the mean but we wanted to avoid Computing the mean because we don't need it this is not what the RMS paper claims the RMS paper claims that we don't need the mean and we don't need to recenter so we need to compute a statistic that doesn't depend on the mean that's why they introduced these statistics here which is the root mean square that doesn't depend on the mean and in practice gives the same normalization effect as the layer normalization and we also have a gamma parameter also here that is learnable and that's multiplied so as you can see the only difference between the layer normalization RMS normalization is that we don't re-center the values and looks like that re-centering was not necessary as written in the paper because they say in this paper we hypothesize that the reselling invariance is the reason for the success of layer Norm rather than the re-centering invariance so they just rescale the values according to the RMS statistic and this is what we will do in our code so let's build this block so the aps value you can see here is used is used as a denominator let me go back here is used here as the added to the denominator So to avoid a division by zero and then we have the gamma the gamma parameter and this is it then we Define the function Norm where X is um batch sequence length dimension okay so we'll turn X applied by torch dot r s Square t r s Square t stands for the one over the square root and that's it [Music] we multiply by gamma so we have as you can see weight is actually is the number a list of ones with the dimension dim so dim multiplied by the sequence length team results in b equals length team where B is the batch dimension and here what we are doing is this is our sqrt is equal to 1 over sqrt of X just as a reminder and in the dimensions here are multiplied by P Square is length one which results in this length Dimensions so what we are doing is exactly multi just this formula here so just multiply it by 1 over RMS and then multiply it with gamma here now that we have also built the RMS Norm let's go check our next building block which is this encoder block so what is the encoder block let's go back to the Transformer here we have the encoder block is all this block here that contains a normalization it contains a self attention here it contains skip connections you can see here another normalization another skip connection and a feed forward layer here I think the easiest one to start with is the fall feed forward but we can also we can also okay let's start first build the encoder block and then we will build the attention and finally the feed forward so we first build the skeleton of this then the attention and then this let's go we save some parameters what is the head Dimension is the dimension of the vector divided by the number of heads so 4096 divided by uh here is the divide by 32 because as we can see here we have the dimension of the vector of the embedding Vector is 4096 but we have 32 heads so each head will see 4096 divided by 32 items from each token then we have a self-attention block I Define it but don't build it right now just Define the skeleton then we have the feed forward then we have the normalization before the self attention tension this is our RMS Norm and this is the motivation behind this argument Norm abs then we have an after the feed forward is it after it's after the attention not after the feed forward before defeat forward block and then we have Norm APS okay okay now let's implement the forward method start pause indicates the position of the token I kept the same variable number as in the original code it's actually the position of the token because as you remember we will be dealing with only one token at a time so start pause indicates the position of the token we are dealing with these are the pre-computed frequencies so we need to the skip connection and yeah okay the hidden is equal to X Plus the attention so we calculate the attention of what of the normalized version of this input so we first apply the normalization and then we calculate this potential and to the attention we also give the frequencies because as you remember the the rotary positional encodings are up kind of they come into play when we calculate the attention and these operations involve tensors of size B sequence length Dimension which is X plus the skip connection plus the output of the attention P sequence length dimension which results in B sequence length dimension then we have another we have the application of the feed forward with its script connection so out is equal to h plus and before we send it to the feed forward as I said before we applied the normalization and this is the output now we need to build the self attention and the feed forward let's start with the harder part first so the self-attention because I think it's more interesting before we build the self-attention let's review how self-attention worked in the original Transformer and how it will work here so okay this is the original paper from the the original paper of the Transformer so attention is all you need let's review the self-attention mechanism in the original Transformer and then we will see it how it works in Lama in the attention is all you need we have an input which is sequenced by the model so a sequence of tokens each token modeled by a vector of size D model we transform them into query key and values which are the same input we multiply by a w Matrix which is a parameter Matrix which results in a new Matrix which has this dimension of sequence by D model and we then split them into the number of heads that we have such that each time each Vector that represents the token is split into suppose we have four heads so each Vector each day each head will will see a part of the embedding of each token so if the token was 512 in size for example the embedding Vector the first head will watch 120 I8 dimensions of this Vector the second head will watch the next 108 Dimensions the next head will watch the next 108th Dimension etc etc etc we then calculate the attention between all these smaller matrices so q k and V this results in head one had two head three and had four we then concatenate them together we multiply with the w Matrix and and this is the output of the multi-head attention in this case it's called self-attention because there is the same input that acts as a query as key and values in case the query comes from one place and the key and the values come from another place in that case it's called cross attention and that kind of attention is used in multi-model multimodal architectures for example when you want to combine for example pictures with captions or music with text or you want to translate from one language to another so you have kind of multi-modality and you want to connect the two together but in our case we are modeling a language so self-attention is what we need actually attention is all we need so let's watch the how it works in Lama okay ilama we need to talk about a lot of things before we build the self-attention we need to review how The self-attention Works in Lama how is the key what is the KV cache what is the grouped query attention and actually how the inference works so we need to review all this stuff before we proceed with the code otherwise it will be very hard to follow the code so let's first talk talk about the inferencing given suppose we have a model so that has been trained on this particular line so we are the line is love that can quickly seize the gentle heart and this is a line from danteligiri you can see this from the apostle from the from The Inferno fifth Canto it's not the first line actually but this is Paula Francesca by the way and we will we have a model that has been trained on this line love that can quickly seize the gentle heart now a model that has been trained on this particular line using the next token prediction should have an input that is built in this way so the start of sentence and then the the the tokens that represent the sentence then the target should be the same sentence with the end of sentence because the Transformer is a sequence to sequence model it Maps one input sequence into an output sequence of the same size this means that the first token will be mapped to the first token of the output the second token of the input will be mapped to the second out token of the output the third token of the input will be mapped to the third token of the output but it's not a one-to-one correspondence because of the Mask of the causal mask that we apply during the self-attention to predict this particular token can for example the model doesn't only watch the same token in the input so that but also watch all the previous tokens so the model to predict can needs to access not only that but also SOS love that and the self-attention mechanism with its causal mask will access all the previous tokens but not the next ones this means that when we do the inferencing we should do it like this we start with the start of sentence and the model will output the first word to Output the next token we need to give the previous output token as input also so we always append the last token of the output to the input to predict the successive tokens so for example to Output that we need to give SOS love to Output the next token we take this that and we put it in the input so that we can get the next word to Output the next token we need to append the previous output to the input to get the new output quickly now when we do this job the model is actually we are giving this input which is a sequence of four tokens and the model will produce a sequence of four tokens as output but this is not really convenient when we do the inferencing because the model is doing a lot of dot products that are not necessary that have already been built um for example what I want to say is that in order to get this last token quickly we are we need to access all the previous context here but we don't need to Output love that can because we don't care we already have these tokens we only care about the last one however we can't just tell them the Transformer model to not output the various tokens we need to change the calculations in such a way that we only receive at the output of the Transformer only one token so that all the other tokens are not even calculated and this will make the interesting fast and this is the job of the KV cache let me show you with some diagrams as you can see at every step of the token we are only interested in the last token output by the model because we already have the previous ones however the model needs to access all the previous tokens to decide which token to Output because the model needs to access all the prompt to act to Output to the next token and we do this using the KV cache to reduce the amount of computation so let's do with some examples suppose we do the same job that we did before so the inferencing of that model we give the first token so the SOS this will be multiplier this is the self-attention so it will be multiplied by the transposed of the keys this will produce this Matrix here which is one by one so you can check the dimensions 1 by 4096 multiplied by 4006 by 1 will output a matrix that is one by one this will be multiplied by the values and this will result in the output token so this is the inferencing step one in which the only token we give is the start of sentence then we take this token output the and the token at the output we this is actually not the token because this has to be mapped to the linear layer etc etc but suppose this is already the token and we append it to the input so it becomes the second input of the um of the input so this is SOS and this is the last output we multiply it by the transpose of the keys we we get this Matrix here we multiplied by the values and we get the two output tokens as output because it's a sequence to sequence model then we append the output of the previous as the input and we multiplied by the transposed by the keys we got this Matrix here we then multiplied by the v's and we get three tokens as output we then append the output of the last one at the Q we multiplied by the transpose of the keys we get this Matrix here we multiplied by the V and we get this sequence as output but we see some problems and then the one that I told you before we are doing a lot of computations that we don't need first of all this um this dot products that we are Computing here because this is the self-attention so Q multiplied by the transpose of the keys will result in a lot of dot products that results in this Matrix these dot products that you see here highlighted in violets that have been already computed at the previous steps because we are at the step number four but these have already been computed at the previous plus not only they have been computed already we don't need them because we only are interested in what the latest token that we added as the input so the to the prompt um what is this tokens dot product with all the other tokens because these tokens dot product with all the other tokens will result in the output of the last token the one we are interested in so if there is a way to not do all these computations again and also to not output all the previous tokens that we actually don't need because we always access the lastest latest token yes we just use the KV cache in the KV cache what we do is we always take the last token and we use it as input so we don't append it to the query we just use it directly as query but because the query needs to access all the previous tokens we keep the keys and the values so we append the last input to the keys and the values but we don't append it to the queries we replace it entirely with the queries let's see with an example for example this is our first step of inferencing so this is just the start of sentence token so we just have one token we multiply it by the transposed of the keys it will result in one by one so we only have one token as output this token we in the previous case was appended to the queries so it became in the next step it became a matrix of Dimension 2 by 4096 but in our case at the time Step 2 we don't append it we only append it to the end of the keys and the values and we only keep the queries here if we do this product again now we will see that this row here is the only one we are interested in so the one that was not violent in the previous diagram and if we do this dot product it will result in only the last token the one we are interested in and every time we keep doing this job we will see the key and the values grow the queries will be always the last token but the number of dot products that we are doing during the inferencing is much less we don't need to do all those dot products that we did before so compare this is time step four this is four dot products compare it with the previous time step four so here here we have 16 dot products so we reduce it by a factor of four and so it's that's why it's much faster to do inferencing with the KV cache and uh let's review again so here as you can see the Matrix kill K and every time we add the token to the queue rows right but all these previous values with the KV cache we are not Computing it again so that's why this is much faster we only compute the one we need and we only get one token as output if this mechanism is not clear please watch my previous video about lamin which if I describe it in much more detail and also with much more visualizations now let's go build this um there is another thing actually I want to show you before we go to build it which is the grouped query attention this one here so I call it grouped multiple attention because it's actually the successive version of the multi-core rotation it's something in between the multi-head attention and the multi-query attention but actually the real name is grouped query attention in the paper also it's called grouped query attention now the reason we introduced the grouped query attention is first of all we had the multi-query attention the multi-query attention basically were introduced to solve one problem that is we first had the multi-head attention we introduced the KV cache with the multi-head attention just the one we just saw the problem was that with the multi-header tension we were doing too many dot products with the multi-head with the KV cache we do less dot products this resulted in a lot less computation but it also resulted in a new bottleneck for the algorithm so the bottleneck was no longer the number of computations but how many memory access we were performing to access this sensors because in the GPU um the GPU is much faster at doing computations that it is at moving tensors around in its memory so when we are optimizing algorithm we know we not only need to consider how many operations we are doing but also how many tensors we are accessing and where are these tensors located so it's not a good idea to keep copying tensor from one place to another because the GPU is much slower at copying memory from one place to another than it is at Computing operations and this is can be visualized on the datasheet of the GPU you can see for example that Computing operations is a 19.5 Tera floating Point operations per second and while the memory bandwidth so how fast it can move memory is 40 times slower so it's uh we are we need to optimize algorithms also for managing how many tensors we access and how we move them around the memory this is why we added the we introduced the multi-query attention the multi-query attention basically means that we have many heads for the queries but we only have one head for the key and the values this resulted in a new algorithm that was all much more efficient than the algorithm just with the KV cache because the KV cache yeah it reduced the number of dot products but it had a new bottleneck that is the number of memory access with this algorithm we also May optimize the memory access but we lose some quality because we are reducing the number of heads for the key and the values so we are reducing the number of parameters in the model and this way um the model we because we are reducing the number of parameters involved in the attention mechanism of course the model will degrade in quality and but we saw that practically it degraded the quality not so much so actually the quality was not bad and this was in this paper so they showed that the quality degradation was very little so from 26.7 to 26.5 but the performance gain were very important we went from 48 microseconds per token to 5 microsecond or 6 microsecond per token so a lot faster now let's introduce the grouped query attention or the group multiple attention in the multi-head attention we had and N heads for the queries and heads for the keys and N heads for the values in the multi-query attention we have n heads for the kids but only one head for the keys and the values in the group to multi-query attention or the grouped query attention we have and less number of heads for the keys and values so every two head for the queries in this case for example we will have one head for the keys and the values and this is a good balance between quality and speeds because of course the fastest one is this one because you have less heads but of course the better the best one from a quality point of view is this one but this is a good compromise between the two so you don't lose quality but at the same time you also optimize the speed compared to the multi-head attention so now that we have reviewed all this concept let's go build it so please again if you did not understand very much in detail this part is better you go to review my other video about llama in which I explain all this part much better otherwise if I have to repeat the same content of the previous video this would be the current video would become 10 hours so let's go build it okay we need to save some things compared to the original code from llama from Facebook from metab I actually removed the parallelization first of all because I cannot test it I don't have multiple gpus I don't have a very powerful GPU actually and then so I simplify the code a lot and KV had some indicates the number of heads for the the keys and the values because they can be different than the number of heads for the queries and this is why we also have an ads queue this value here is represents the ratio between the number of heads for the query and the number of heads or for the keys and the values we will use it later when we calculate the attention so let me write some comments so this is and then we have a cell Dot head Dimension which is this indicates the the part of the embedding that will be visualized by each head because as you know the embedding is split into multiple heads so each head will watch the full sentence but a part of the embedding of each word then we have the W matrices WQ wkwv and W O just like in the normal vanilla Transformer and they don't have any bias Oops why did I write true and then we create a cache we will see later how it's used I just now created one for the keys and one with the for the values foreign [Music] the forward method which is the Salient part here so tell X is to simplify the the code for you I will write the for each operation I will write the dimensions of the tensor that is involved in the operation and also the resulting tensor from each operation this start position indicates just the position of the token inside of the sentence and these are the frequencies that we have computed okay let's start by extracting like size B sequence length and dimension but the sequence length we know it's one so Dimension yeah then what we do is we multiply just like in the original Transformer we take the query the key and values we multiply it by then the WQ W K and W key Matrix so X cube is equal to self.w Q this means going from B one dimension to B one head Dimension so the number of heads for the query multiplied by the dimension because we are in this case we are this is actually equal to dim so the number of heads multiply the head Dimension as you can see from here so we are not changing the the shape in this case however we may change the shape of the because the number of heads for the KV may be smaller than Q so this Matrix may have a last Dimension that is smaller than xq and the same is uh for XP so here let me write some comment apply the WQ WK and WV mattress two queries keys and values which are the same because it's a self-attention so the query key and value is always X we then divide them into their corresponding number of heads so xq is equal to X Cube Dot View that size we keep it like this sequence length is one so we divide B1 HQ multiplied by head Dimension into B one add HQ and head dimension so we divide them into the H heads for the query and then we do the same for the key and the values and the same for the B now we have multiplied okay we have the X input we multiplied by the WQ W K and W key why let's go check the code here as you remember we take the input we multiplied by WQ w k and WV this will result in this matrices here we then divide them into the number of heads but in the case of grouped orientation they may be different so this may be four heads and this may be two heads and this may be two heads so they are not the same number the next thing we are going to do and this is present in the here we need to apply the rotary positional encodings to the query and the keys but not the values let's do it and this is how we apply the original encodings this will not change the size of the vectors you can see that here because at the end we have the same shape as the original input vector okay now now comes the KV cache part let's watch again these lights as we can see here every time we have an input output token so for example the attention to here it suppose it's the token number two we append it at the end of the keys and the values and this is exactly what we are going to do so what we do here we keep a cache of the keys and the values because they are will be used for the next iterations because at every iteration in X we only receive the latest token that was output from the previous iteration we append it to the K and the V and then we compute the attention between all the K all the V but only the single token as query let's do it so first replace this is the position of the token this should be 1 because sequence length is actually one always but I try to keep this code the same as the one from llama from meta this is basically it means that if we have one token from many batches I mean we have one token for every batch we we replace them for for forever because we can use we can send the we can process multiple batches so we replace the entry for this particular position for every batch okay now we replace it only for this position here but when we compute the attention using the KV cache that's cool watch again we need to calculate the product the dot product between the only one token but all the keys and then we will need to multiply with all the values and this will result in only one token as output so we need to extract from this cache all the tokens as keys and all the tokens as values up to this position here the one we are passing so this is equal to all so starting from 0 up to start pause plus sequence length and the values of length now now what happens is that let me write also some sizes here we have B sequence length of K and V because the sequence length of the input is always one we know that but the sequence length of the cache means all the um all the cached case keys and values which are up to start position so this sequence length is actually equal to start position and actually start position plus one next Dimension is the number of heads for the K and B and then the dimension of each head now the number of heads for the keys and values may not correspond to the number of heads we of the um queries so how do we compute in the original code from Lama what they did was basically let's go check code for here so in the grouped query attention we have that the number of heads for the keys and the values is not the same as the number of heads for the queries so there are two ways one is to make an optimized algorithm that actually takes this into consideration the other ways is to just copy this single head into multiple um into multiple heads such that we arrive to this situation here and then we just compute it just like a multi-head this is not optimized solution but is the one used by The Code by llama and it's also the one I will I will be sticking to because I don't have any way of testing other codes because the only model that that supports the grouped query attention is the biggest one from La mouth so with the 70 billion parameters but my computer will never be able to load that model and so I don't have any way of testing it so that's why I also didn't optimize the the code for actually Computing the grouped qualitation but I will just replicate this single head multiple times such that we arrive to this situation here so I will also repeat okay this function here repeat KV just repeats the the the keys until we reach the number of for this number of times so n-wrap what is this this is the ratio of the number of heads of the queries by the number of heads of the keys so if the number of heads of the keys is 4 and the number of heads for the queries is eight we need means we need to repeat twice each head so let's build also this method since we are here okay we don't need to repeat it so there is only one repetition we just return the basic the answer otherwise we repeat it 10 times so the first thing we do is we add a new dimension and we can do like this but sequence length number of heads then nothing and then this will add this uh new um dimension in this position then we expand it and then we reshape it basically we introduce a new dimension we repeat this all the sequence this Dimension number of times so and wrap along this Dimension and wrap number of times and then we just flatten it so we remove remove again this dimension and this is how we repeat this keys and also the values now we can repeat okay now we just proceed just like with the standard the standard calculation for the multi-head attention that is we first move the head Dimension before the sequence Dimension because each head will watch all the sequence but a part of the embedding of each token so what we are doing is batch one because one is the sequence length of the queries which the number of heads of the queries and head dimension patch head sequence length and head dimension we do the same for the keys and the values then we do this standard formula for queries multiplied by the transpose of the keys divided by the square root of the dimension of each head so X Cube so the query is multiplied by the transpose of the keys all of this and divide by the square root of the dimension of each head then we apply the soft Max and this one will result in a shape of queries one add dimension multiplied by KV the soft Max doesn't change the dimension then we multiply it by the values so the formula is the queries multiplied by the transpose of the keys and the output will then we do the soft Max then the output is multiplied by the values so this will result in um B and then we multiplied by the output Matrix but before we remove all the heads so we concatenate again this is what we did also here so here we take the output of all the heads then we concatenate them together and then we multiply it by the w o Matrix and this will result in a B1 team V1 dim this one is bhq one head Dimension into B1 HQ head dimension because of the transposition and then we remove the dimension for the head so B1 Dimension and this is our self attention with KV cache so let's review what we have done here I think I made some mistake because uh self that's why it's color differently okay let's review what we have done when we calculate the self-attention because we are inferencing so this code will only work for inferencing we can use the KV cache the KV cache allow us to save a number of dot products that we don't need why because every time we are com in the original Transformer we were Computing a lot of those products for tokens output tokens that we don't care about in this case we simplified the mechanism to Output only one token as you can see the output of the self attention is B so batch one token only with its embedding size which is 4096 so we are only outputting one token not uh many tokens we input only one token and we output one token and but because we need to relate that single token with all the previous tokens we keep a cache of the keys and the values every time we have a token we put it into the cache like here then we retrieve all the previous save the tokens from the cache and then we calculate the attention between all the previous tokens so the keys and the values and the single token as input of as queries the output is the only token we care about this is the idea behind um KV cache and the grouped query attention is the fact that we have a different number of heads for the keys and values but in our case we do have we do have a number of different number of heads for the keys and the queries but we just repeat the one that we are missing to calculate the attention so the direction is calculated just like the previous Transformer like a normal multi-head attention foreign but by repeating the missing keys and values heads instead of actually optimizing the algorithm this has also been done by meta in its official implementation and I also did it here the biggest reason is because I cannot test any other modification I cannot test the another algorithm that actually tries to optimize this calculation so if I find another implementation that I know is working I will share it with you guys otherwise I will try to run it on collab and see if I can come up with a better solution but for now we just repeat it but at least we got the concept of the grouped query attention that is we have less number of heads and it's something that is in between with between the multi-query attention and the multi-head attention that doesn't sacrifice quality but improves speed now the last thing that we didn't Implement is the feed forward layer for the feed layer forward layer the only thing that we need to review is this video activation function that we can see here and this activation function has been changed compared to the previous activation function used in the vanilla Transformer which was the relu function and the only reason we replaced it is because this one performs better and as I showed in my previous video there is no outside there is we cannot prove why it works better because in such a big model with 70 billion parameters it's difficult to explain why something a little modification works better than another we just know that some things that work better in practice for that kind of model or for that kind of application and this is actually not my opinion this is actually written in the paper so as you can see here in the conclusion of the paper they say that we offer no explanation as to why these architectures seem to work we attribute their success as all else to Divine benevolence so it means that when you have a such a big model and you change a little thing and it works better you cannot always come up with a pattern to describe why it is working better you just take it for granted that it works better and you use it because it works better in practice so to implement this wiggler function we need to apply this this is the formula from the original Transformer so the we have two matrices here so this is the relu function of the first linear layer and the second linear layer in the Llama we use the Ziggler function which involves the three matrices here because they incremented the number of parameters here and also they were experimenting with the grouped query attention the architecture of Lama has a multi says some more parameters to adjust the number of parameters of this feed forward layer so as it respects some constraints and this is actually used in deep learning research whenever we modify them Transformer model and this reduces the number of parameters or increases the number of parameters the first thing the researchers do they adjust the numbers of parameters of the feed forward layer so that when they make comparison between two models they have the same number of parameters so I will also of course use the same structure because otherwise I cannot load the weight from the pre-trained model so let's do it the hidden size is calculated like this so four times the dimension then they do the two thirds of this dimension and then they also have a multiplier if it's specified then they say round the hidden oops by using this modification so to concluding The Hiding Dimension like this it may not be the case that this hidden Dimension is a multiple of this number here so maybe they want the F the size of the Hidden layer to be a multiple of 256 path by calculating it like this it may not be so what they do is they make it round up to the next multiple of the multiple of parameter this is a way to do it okay let me give example it's easy to show with an example than to actually write it so suppose you have a the hidden size is equal to let's say seven but you want it to multiple off is equal to five so you want the hidden size to be a multiple of five so how do we do well what we do is basically we do hidden plus four in this case so we do seven plus four which is 11. we divide it by 5 which is equal to two and then we multiply these two by five so it will result in two by five is equal to 10 it will result in the first multiple that is bigger or equal to this number here that's the idea then we have this matrices for this wiggler function it's very easy and we just follow the formula for this vehicle function which is here so w the switch of what is which um this wish is The Silo function because the switch with the beta is equal to one is actually the zero function which has this graph here and then we multiply it with another parameter Matrix here and then we apply it to another linear layer W2 so in total we have three matrices W1 we call it W2 and W3 and they don't have bias oops the hidden dimension okay now we'll implement the forward method the first thing we do is we calculate this which function then we calculate so we are calculating let me show you we are calculating this one x w switch of x w then we calculate this XV then we multiply them together just like in the formula so which multiplied by X V and then we apply the last linear layer which is W2 results in a multiplication by the W2 matrix by the way and then return X and this is the field forward layer now that we have all the building blocks we need to go to the inferencing let's start building the inference code so inference dot Pi the first code we will be first we will build the code to load the model and then we will build the code to inference the the model I will actually also show all the inference techniques that are out there and which one we will apply and why so let's start by building first the code for loading the model so let's import the stuff we need we need the Json to load the parameters and then we need the sentence piece to load the tokenizer because the sentence piece is the tokenizer that have been used and it's a library from Google okay our model import model arcs and the Transformer class we define a class llama which is our model takes a Transformer a tokenizer which is a sentence piece processor and then the model arguments oops arcs model arcs yeah model or okay now we built a static method static method and we call it build just like in the original code from Lama in which we pass the directory which we where the checkpoints are saved in this case the director name is llama27b in my case but it depends on which size of the model you have downloaded then the tokenizer path which is the the path to the tokenizer this is the file of the tokenizer that I downloaded then we have a load model Max sequence length max box size and we have device okay this is only for displaying how much time it takes to load the model if we want to load the model we will also load the checkpoint so checkpoints is equal to sorted the global method allows you to find all the files that match this filter okay we see that we are loading checkpoint this one and then we actually load it and we save it on the CPU we can show how much time it takes to load the model I in my computer usually takes her 10 to 20 seconds then previous time we rewrite it so we can also show how much time it takes to load all the parameters of the model then we load the parameters so the Json file we read it open it as read only file okay okay let me build the arguments maximum sequence length is the one we have specified and then we have the max batch size is the max pad size the device is the one we have specified and then all the parameters loaded from the Json file then we loaded the tokenizer then we by using the tokenizer we can populate the vocab size of the model arcs the vocabulary size is actually not the number of tokens inside the tokenizer now this is also the default tensor uh Fork for pi torch so whenever pytorch wants to create a new tensor what kind of type it should use it's defined this is by meta so they want for Cuda to use this type that I show you here default sensor type torch dot Cuda half tensor this changes the Precision that the tensor supports so how much space it occupies in memory otherwise then we created the actual model okay when we load a checkpoint actually the checkpoint is a list of key and values each key is a matrix in the model so the weight for example of a linear layer or the bias of a linear layer or something like this and the the names that we have using used for uh the variable names and the the matrices here for example wqwk match actually the name that are present in the checkpoint here except for one name so to make sure that I have used the right names I will load the the checkpoint with the strict equal true strict equal true means that if there is at least one name that doesn't match it will throw an error so if load model dot node State ticked tricked equal true so if there is at least one name in The Loaded file that doesn't match the the name in the classes that I have created here in the model it will throw an error but I know that there is one key that we don't need which are the frequencies for the rotary positional embeddings which we actually are Computing every time we create the tensor so we are creating them here by using this function so we don't need to load them from the model so we can remove it from the model from the checkpoint so because the checkpoint is a dictionary we can just remove this it's called rope.frex and then we can print how much time it took to load the model and then there's a llama model tokenizer and model arcs now before we proceed further let me test if the model can be successfully loaded so let's do it with name first I will set the manual CD to zero so later we use it for inferencing then I don't want to use Cuda because my my GPU doesn't support it so I say hello could I equal four then device is equal to storage dot Cuda dot is available and along Cuda Health CPU next time if you want to load the model with Cuda I just set this variable to true but in my case I will always leave it to false because I don't want to load Cuda as you can select I set it to 1024 Max batch size let's say three and device device now let's run it and hopefully it will not crash wow already not then sorry but tensor so let's run it again there is always a lot of typos when you write code another problem here ah not storage but tensor it should be tensor B for 16 tensor yeah let's try again hidden hidden what hidden dimension of course and let's try again yeah all okay okay wonderful it means that at least it's doing something and it's not crashing which is always a good news so our next step is actually to build the inferencing code so what we want to do is actually we want to be able to give some prompts to the model and then check the output for this prompt so let's define some prompts we will Define some prompts here and here we pass for example the size of the prompts and then we want to you know we want to inference inference the model so before we start inferencing the model we need to build the code for inferencing the model because we need to find a strategy for selecting the next token etc etc so let's review how the inferencing works and what are the various strategies for inferencing okay so when we are dealing with the next token prediction task when we want to inference we usually give the prompt and then we want to predict the tokens but we give one token at a time and every time we give one more token the model will output one more token as output and we only keep the last one but with the KV cache actually we always give one token at a time the KV cache will keep the cache for the keys and the values and with all only output one token okay the point is we need to find strategies for selecting this token among all the tokens that we have in the vocabulary and this is the job of the logits and the soft Max so let's review how they work now imagine I give you the following task as human so complete the following sentence I think nuclear power is and then you have to choose a word now you as human may have thought of the possible next tokens which may be clean dangerous cheap expensive safe difficult or something else the choice of the next token in your head depends on your education on your experience with nuclear power and your opinion on the matter large language models also face the same problem when we give them a prompt then the model has to choose the next word the model the uncertainty of the choice derives entirely from their training process and the strategy that we use to select the next token there are many strategies for example we have the greedy strategy the beam search temperature is a parameter random sampling top K top p in this video we will review all these strategies and how they work but first we need to understand what are the logits let's look at the Transformer model from llama so the output of the self-attention is a sequence in the case of the KV cache is only one token we then run it through a linear layer so after normalization we run it through a linear layer the linear layer will transform the embedding that is output from the self-attention here into a into a list of numbers that represent the kind of the probability they are not really a probability but we can think of it as a probability of that token in the vocabulary so if our vocabulary is made of let's say 100 tokens this linear layer will output 100 numbers and after we apply the soft Max this 100 numbers will become the probability of that token being the next more probable token for the prompt given to the input so given an input a prompt the model comes up with a probabilities probabilities for which token to choose next and so what is the job of the linear layer and what is the job of the software X the linear layer converts the embedding of a token into a list of numbers such that each numbers represent a score that later with the soft Max represents the probability of that particular token in the vocabulary the softmax job is just to scale the logistics in such a way that they sum up to one so that's why we can we can talk about probabilities with the soft Max but with the logits so the output of the softmax is just a probability distribution over all the words in the vocabulary that is each word in the vocabulary will have a probability Associated to it but now given these words each one with their probability how do we choose the next token there are many strategies and the easiest one is the greedy the greedy strategy basically says we just select the token with the maximum probability so imagine we are inferencing and the time step is the first time step in the greedy strategy The Prompt is Celia you're breaking my heart you're shaking my okay this is a line from a very famous song from Simon Simon ergon carfunkel and the next word for those who know will be confidence so Celia you're breaking my heart you are shaking my confidence suppose the so output of the soft Max is this distribution here so we have 40 probability for this word 20 for this word the 15 for this word and 10 for this word with a greedy strategy we always choose the token with the maximum probability then we append it to the input so the input as the next inference step becomes Celia you're breaking my heart you're shaking my confidence and then the model has to come up with the next word which if you know this song is daily if we use the greedy strategy we select the one with the highest probability so in this case it's daily and it's also the correct one so this is how the grid strategy works at every step we choose the token with the maximum probability which is then appended to the input to generate the next token and so on but if the initial token happens to be the wrong one so not only the initial but the initial two three tokens happen to be the long wrong ones it's very likely that all the next tokens will also be wrong because we are giving a wrong prompt to the model so imagine at the time step one we don't use confidence but somehow the model came up with a high score for Lever so you are shaking my liver but then the next word the model will not be able to come up with a reasonable next word because there is no song that says you are shaking my liver so if we make a mistake in the early stage of the greedy all the next token very probably will also be wrong but it's very easy to implement and however it performs poorly in practice that it's very it's not used so much another strategy is the beam search in beam search we have a parameter which is called K which means that at every step we not only choose the top ones but the top k at every step and we always keep the top two best performing tokens so in this case for example imagine we are time step one so cilia you're breaking my heart you are shaking my and the top two words are pizza and confidence Pizza somehow has a higher uh has a higher probability because maybe the model has never seen this song before so it doesn't know that the next word is confidence so maybe the model outputs these probabilities but we choose the the two top most the two tokens with the highest probabilities at the next time step we make two prompts one in case we choose the first one so the first token and one in case we choose the second token and then we see what are the next possible choices if we use the first token and what are the next choices if we use the second token so we check the model output for the first prompt and for the second prompt and in case we use for example the first prompt the model will output this probabilities and if we use the second prompt the model will output these probabilities what we do then is we calculate the cumulative score for each possible path so for pizza for example the probability was 40 percent but after Pizza the model produced the probability for the margarita for example is 0.04 so for this part pizza margarita it's 0.004 the probability is 0.4 percent Pizza anchovies it's going to be 0.2 percent or 0.002 however with confidence we get a new uh next token that can be either daily or monthly with daily we get a cumulative score of 0.16 and with monthly of 0.02 so as we can see at the time Step 2 even if at the time step one pizza was the most probable word because we kept the second choice alive so we didn't kill it just like we did it with greedy let me use the laser we can see that the confidence then produces a next token that is very probable because now the model has more prompt and so it can come up with more specific choices for the next tokens with a very high confidence so we we can't compute the the cumulative score of all these parts and we keep the two parts that have the top choices so now the pizza part has been killed because it's later we chose pizza at the beginning because somehow the model thought it was pizza but then it couldn't find the model was not so confident about the next words but in the case of this token here the model was very confident about the second score so we killed all this part here and we kept this one until we arrived to the last token in which we just selected the path with the highest score and that's the output of our inferencing strategy with beam search and repeat the steps of the last type of the last slide for all the successive tokens until we'll arrive to the last one and then with the beam search at every step we keep alive the top key parts and all the other are killed it increases inferencing time because at every step we must explore key possible options but generally it performs better than the greedy strategy for the reason that I have just shown another thing that is interesting in inferencing is the temperature because the idea of the temperature is that we want to make the we can make the model more confident or less confident so for example when we calculate the logits which are not the probabilities so they are what will become the probabilities after we apply the soft Max so before we apply the softmax we can scale the logics so that if we use a for example like this so for example you have this logits here I choose the negative numbers so the soft Max probabilities are reasonable numbers and so these are the logits and if we divide these Logics before applying the soft Max by a number that is low so low temperature it's called additional number is called the temperature it will make the model more confident because it will make bigger probabilities bigger and smaller probabilities smaller so the gap between the low and high probability increases so for example you can see that with without applying any temperature the highest logic gets 80 percent probability but applying a 0.4 temperature the highest logic becomes 98 probability and if we apply a high temperature it makes the model less confident so the gap between the low and high probability reduces the temperature is important if we want to increase the confidence of the model or not because it can be used in conjunction with other strategies like for example the greedy or the top key or the top width that we will see later another strategy is the random sampling so as we saw the logits are not a probability distribution but after we apply the soft Max they become a distribution so what we do because it's a distribution we can also sample from this distribution for example we in this distribution here that becomes from these Logics here we have one token that can be chosen with the 12 probability one can be chosen with seven percent probability and one that can be chosen with 80 probability if we flip a coin by eighty percent of the time we will choose this token and 12 of the time we will choose this token and seven percent of the time we will choose this token so this means sample from this distribution it means sample take a number from distribute this distribution according to its weight to its probability now there is a problem with this sampling strategy here that with very very little probability it may happen that we choose tokens that are total crap for example in this uh scenario here for example before with the greedy strategy or with Brim search for example this token here if we use a random sampling we will choose the word pizza with 40 probability the word confidence with the 20 probability but with a very little probability it may happen that we will choose the word Pokemon with 10 probability of course the probability is low so the the probability of us making a bad choice is low but it's there is this probability so this is a problem with random sampling the next strategy is the top K in top K what we do is to avoid selecting the crappy tokens we just remove them so we stick all the logits we sort them and then we just keep the highest K so that's the crappy one we just removed them from this um distribution and then we calculate the distribution for the rest so we also apply this aftermax only to the one that survive the problem is also here given the following these two distributions the low probability tokens can still make their way into the top K because it all depends on the distribution to which we apply the top K let me give you a graphical example imagine we have a distribution that is very flat suppose this distribution here so some more some words these are this is our vocabulary this is the probability of each word so the word number one word number two word number three word number four etc etc etc but more or less all the words have the same probability so imagine we take the top 10 words so it will select all these tokens right okay so it will select the token number one token number two token number three talking number four up to whatever token here is here imagine we have another distribution that is made like this so we still have a vocabulary we still have a probability distribution and the distribution is made like this so because it's sorted we have a distribution that is very rescued with the because we still keep the top 10 as you can see we will select this token this token this token this token this but these tokens here are very crappy compared to the this one here so they will still make their way into our selection and this is not something that we want we want to avoid selecting crappy tokens but we still want to have some Randomness so we don't want to be totally greedy because sometimes the tokens that are in the top and tokens maybe they are very they're all reasonable also sometimes the prompt can be quite ambiguous so we don't know which even as humans we may not know what is the next word to be chosen so we we want some Randomness but we also don't want the very low probability tokens but with this top case strategy the top probably the low probability tokens can still make their way into our selection and this problem is solved with top p with the top P we only keep the tokens with the highest probability such that the cumulative probability is greater than or equal to the parameter p what does this mean it means that if we have the previous distributions so one that is a quite flat for example and the one that is has a mode so for example this one so this one is nearly 90 and the other one are 0.000 but this is more or less all of them are like 0.2 percent and then they go down in the case imagine p is equal to let's say 0.5 in this case we will select all the tokens such that the area under the the curve is equal to 0.5 but here because this first token is already 0.9 we will actually only select one token and all the crappy ones will not be selected because this area under the under the curve is already 0.9 and this is the idea behind the top P so to when the model is when the distribution distribution is more flat we select more tokens because it means that we are more uncertain about which token to choose but when we have a big mode we select uh the fewer tokens and um this way we we re-avoid getting the low probability ones so now that we reviewed all the strategies for selecting the token we will implement it and in the case of Lama also in the official code they actually implement the top P strategy in my case I think that the beam search is a reasonable choice so in my another video maybe I will make a how to implement the beam search but for now I will implement the top so let's go build it so implement the method let's call it text completion which is the same name that's used in the original code from llama given prompt a temperature that is 0.6 and so 0.6 means that we want to make the model more confident top p means that we want all the tokens such that their cumulative probability is at least 0.9 so 90 percent okay I think here should be lowercase okay so if we didn't specify the max generation length then we just generate the maximum token so dot Max args just generate as mass token as we can up to the sequence length and then we first of all convert each token of the prompter so each each prompt actually into tokens using the tokenizer mm-hmm then we as we saw before we need to add the beginning of sentence when we um when we we get that we we pass the input to the model for inferencing but not the end of sentence okay because we specify the max batch also for the model when we built it for the KV cache so we need to make sure that the batch size of the prompts is not too large and then Max prompt land is the maximum prompt length that we have in the prompt I'm not writing any message even if you should but okay for us it's just a basic debugging then the total length is how much we want to uh how many tokens we want to get from the model okay now we create the list that will contain the generated token this means create um a tensor of shape batch size by total length in which each item is actually the padding token and then we fill the initial tokens with the prompt tokens okay we also need this variable that tells if we reach the end of sentence in any of the prompts this indicates is um if the token in this position is a padding token or not so true if the token is a prompt token false otherwise and then we can finally make the for Loop to generate the tokens picture take radium range from 1 to now we generate one token at a time so the logits come from the model so set dot model dot forward we need to pass one token at a time so which token the one currently we want to Output so current minus one plus it's only one token and we also tell the model what is the position of this token because for the KV cache and if we use the temperature we apply it as you can see every time when we want to inference we always select the last token but because we are using the carry cache actually our model will only output one token at a time so the next token will be selected according to our topped P strategy so we have the probabilities now we use the apply the soft p in the top p i just Define it here so sample top B and then we implement it if we didn't buy any specify any temperature we just use the greedy okay now we have the next token according to this strategy or this degree then we only replace the token if it is a adding token so the problem is we already have some tokens that come from The Prompt but we don't want to but we still need to give the prompt to the model but we are only giving one token at a time to the model to build the initial cache so we will give in the first prompt tokens will be given to the model not because we care about what the model will output for those tokens um but only because we want the KV cache to be built for those positions and after we give the last token of the prompt then we care about what is the model outputting so for only replace the next token if it is a padding token and which one is the padding token the one that was not an initial prompt token because here we build tokens full of paddings but then we replace the pad The Prompt tokens with the the padding tokens with the prompt tokens for the for the one with the for the initial tokens all the other have to be inferred by the model so token this means basically check this mask what is this mask if it's true if the token is a prompt token so if it is a prompt token replace it with this one and if it's not a prompt token just keep it the current one okay then since we do not care about what the model outputs for the initial prompt tokens but only for the last prompt token we don't care if we find an end of sentence position for those tokens so and then end of sentence is only reached if we find it for one of the tokens that we actually want to inference not the one that we send to the model just to build the KV cache okay this basically means the the end of sentence for a particular prompt is reached only if it was a padding token so only if it was a padding token so it was not a prompt token this means not and we actually found an endos token at the from the model output if all of the prompts have reached the end of an sentence token then we stop this for Loop we don't need to inference anymore now we prepare the output this means that if we found the endos and end of sentence token for one of the prompts we just cut the prompt output to there the our model output at that particular token we don't care about what it output next this is the output text and then we output the tokens and the text hopefully I didn't make too many typos and mistakes so now we need to build the sample top P so we have the logits that are the output of the model we transform them into probabilities by using the soft Max but given these probabilities we need to use the sample top P strategy to select all the tokens such that their cumulative probability is equal to top P which in our case is 0.9 so 90 percent there be okay the first thing we do is we sort these probabilities in descending order we then calculated the cumulative sum then we create the mask that says which tokens we want to keep and which one we don't want to keep so mask is equal to probability sum minus probability sort more than P why do we do a minus probability sort because we want to shift let me show you on the slides here you can see here for example the cumulative probabilities so the probabilities are this one 0.4 for forty four forty four percent forty percent six percent four percent and three percent then we calculated the cumulative that means up to here it's 44 then this one plus this one is 85 percent and this one plus this one plus this one is 91 percent this one plus this one plus this one plus this one it's 96 etc etc but imagine we have a 0.99 90 probability or 0.55 probability uh we need to keep up to this token here because this one is not enough it's zero so we need to up to this one so the first number that is less than or equal to p and it's this in case it's this one so that's why we shift it we want also this token inclusive and this is why we do this minus probability sort then all the one that we didn't select we zero them out zero repeat and then we redistribute the probabilities because of course if we remove some items from here they don't sum up to one anymore so we need to redistribute the probabilities and this is very easy okay then the next token is basically suppose we keep the first two tokens and then what we do is we want to sample from them so the first token is 0.44 probability the second token is 0.40 probability but after we redistribute their probabilities actually this one will be a little higher and this one will be a little higher than 40 percent and um and then we sample it means that the first token will have a slightly better chances of being Chosen and the second token will have slightly less chance of being selected we want one sample because we want one token and it's not text token and then next token because this indicates which index to select then we need to map that index to the actual number in the vocabulary but because we already changed the order of these numbers because we sorted it so initially the logits were built in such a way that the first logic corresponded to the first number of the vocabulary the second log it corresponded to the second number of the vocabulary but because we sorted it by descending order this order has been gone so we don't know now just given the the the token selected we don't know which number it Maps back into the vocabulary that's why the sort method returns two arguments one is the sorted numbers and one is the indexes that it's changed so it it will tell you for each position what was the original position item in that position so this is why we actually query using gather together allow us to retrieve from an element what was the original one given this index here and then we return the next token and this will map back into their vocabulary directly and this should be it so now let's create some prompts and let's run the code I have some prompts here that I copied and pasted so now let's build the the inference code so alt tokens alt text we want to generate maximum 64 tokens we assert that the line of the output text is actually equal equal to line of prompts it should be so we are in range hopefully the model will work and then we print the output text for each prompt so let's run the code and let's hope for the best okay self-attention is missing the required forward function let's see why oops it's forward it should be forward let's run again some received this is wrong because it should be dimension not div but should be dim let's run again for B float 16. so let's see why EOS token let me check okay now it's training I just changed this tensor from capital T to small t I will investigate why well we have an output so let's check first of all let's check the prompt simply put the theory of relatively relativity states that says that one time is relative to The Observer mass is relative to The Observer speed is relative to The Observer energy is relative to that so looks like it's not bad this suppose the second problem says if Google was an Italian company founded immunity to put it would be listed on the Milan Stock Exchange as the middle Stock Exchange is the largest in Italy but since Google is a U.S company it is listed on the NASDAQ Stock Exchange so it avoided actually answering the question let's try the few shot prompt so this is how you copied actually from the Llama code so they ask to translate from English to French and we after cheese after cheese we expect to find fromage onion onion Etc so looks correct and we can also see that the the spaces have been kept so these spaces were not added by me but actually by the model so it keeps the output aligned with what was the prompt and then I created the zero shot prompt so tell me if the following person is actually the rhyme on these guys that's human so the name is Umar Jamil and the decision is he is a hero in every sense of the world he's a hero in every sense of the world I'm very happy llama actually okay this is the output of the model with the manual seat zero if I think I changed the seed to something other number and run on the model again the the output will be totally different or maybe slightly different I hope not but it may be different anyway thanks for watching my video guys I tried to convey the idea of what made what is the architecture inside Lama and even if I didn't build the training code because actually to build the training code is rather complicated we need a big Corpus of text we need to tokenize it and it's going to take a long long time but I hope to make another video in the future on how to train a language model maybe with a smaller data set and with a lighter architecture and I try to convey all the math behind all the choices and also the how the inner workings of the KV cache and the the grouped query attention if you have any questions please write in the comments I also will share the repository with the code that I have previously Built For This and which has much more comments than the one I have written here it's more much more in detail so everyone can understand step by step all the dimensions involved here I try to write the most important Dimensions but because of time I didn't write all of them so thank you again guys for watching it was a long journey but I can assure you that you learned a lot hopefully and I hope you will visit again my channel for more videos about deep learning about pytorch about coding and about everything that we love in AI thank you for watching guys have a nice day

Original Description

Full coding of LLaMA 2 from scratch, with full explanation, including Rotary Positional Embedding, RMS Normalization, Multi-Query Attention, KV Cache, Grouped Query Attention (GQA), the SwiGLU Activation function and more! I explain the most used inference methods: Greedy, Beam Search, Temperature Scaling, Random Sampling, Top K, Top P I also explain the math behind the Rotary Positional Embedding, with step by step proofs. Repository with PDF slides: https://github.com/hkproj/pytorch-llama Download the weights from: https://github.com/facebookresearch/llama Prerequisites: 1) Transformer explained: https://www.youtube.com/watch?v=bCz4OMemCcA 2) LLaMA explained: https://www.youtube.com/watch?v=Mn_9W1nCFLo Chapters 00:00:00 - Introduction 00:01:20 - LLaMA Architecture 00:03:14 - Embeddings 00:05:22 - Coding the Transformer 00:19:55 - Rotary Positional Embedding 01:03:50 - RMS Normalization 01:11:13 - Encoder Layer 01:16:50 - Self Attention with KV Cache 01:29:12 - Grouped Query Attention 01:34:14 - Coding the Self Attention 02:01:40 - Feed Forward Layer with SwiGLU 02:08:50 - Model weights loading 02:21:26 - Inference strategies 02:25:15 - Greedy Strategy 02:27:28 - Beam Search 02:31:13 - Temperature 02:32:52 - Random Sampling 02:34:27 - Top K 02:37:03 - Top P 02:38:59 - Coding the Inference
Watch on YouTube ↗ (saves to browser)
Sign in to unlock AI tutor explanation · ⚡30

Related Reads

Chapters (20)

Introduction
1:20 LLaMA Architecture
3:14 Embeddings
5:22 Coding the Transformer
19:55 Rotary Positional Embedding
1:03:50 RMS Normalization
1:11:13 Encoder Layer
1:16:50 Self Attention with KV Cache
1:29:12 Grouped Query Attention
1:34:14 Coding the Self Attention
2:01:40 Feed Forward Layer with SwiGLU
2:08:50 Model weights loading
2:21:26 Inference strategies
2:25:15 Greedy Strategy
2:27:28 Beam Search
2:31:13 Temperature
2:32:52 Random Sampling
2:34:27 Top K
2:37:03 Top P
2:38:59 Coding the Inference
Up next
5 Levels of AI Agents - From Simple LLM Calls to Multi-Agent Systems
Dave Ebbelaar (LLM Eng)
Watch →