A Practical Guide to Building GPT-2 with PyTorch (Part 2)
Source: tai_blog

This is the second part of the GPT-2 from scratch project. If you havent read the first part yet  I highly recommend getting familiar with the language model basics before continuing.   Build and Train GPT-2 (Part 1)Final Loss:   4. Implement GPT-2 architectureIn this section  we will add the GPT-2 parts one by one and then train & evaluate how the model performs in each stage. Heres how it goes:   a. Positional Encoding + Fully Connected Layer (NN)   b. (Masked) Self-Attention + Normalization   c. (Masked) Multi-Head Attention   d. Multiple GPT Decoder Blocks   e. Improving Tokenizer   f. Final GPT-2 Training   To recall from previous part  our model looks like below:   Code:   import torch.nn as nn import torch.nn.functional as F # used to define size of embeddings d_model = vocab_size class GPT(nn.Module): def __init__(self  vocab_size  d_model): super().__init__() self.wte = nn.Embedding(vocab_size  d_model) # word token embeddings def forward(self  inputs  targets = None): logits = self.wte(inputs) # dim -> batch_size  sequence_length  d_model loss = None if targets != None: batch_size  sequence_length  d_model = logits.shape # to calculate loss for all token embeddings in a batch # kind of a requirement for cross_entropy logits = logits.view(batch_size * sequence_length  d_model) targets = targets.view(batch_size * sequence_length) loss = F.cross_entropy(logits  targets) return logits  loss def generate(self  inputs  max_new_tokens): # this will store the model outputs along with the initial input sequence # make a copy so that it doesn't interfare with model for _ in range(max_new_tokens): # we only pass targets on training to calculate loss logits  _ = self(inputs) # for all the batches  get the embeds for last predicted sequence logits = logits[:  -1  :] probs = F.softmax(logits  dim=1) # get the probable token based on the input probs idx_next = torch.multinomial(probs  num_samples=1) inputs = torch.cat([inputs  idx_next]  dim=1) # as the inputs has all model outputs + initial inputs  we can use it as final output return [decode(out.tolist()) for out in inputs] m = GPT(vocab_size=vocab_size  d_model=d_model).to(device)Now  lets add Positional Encoding into our model:   The Output Embedding (in our case its the input embedding wte) is added with Positional Encoding and then passed into the further network. To understand what PE is  lets recall token embedding which stores d_model dimension of vector for each character in our vocabulary. It represents different properties of the character based on how and where it appeared while training.   Similar to this  the Positional Encoding stores the order/positional signal of every character in the context_length. It is only calculated once using sine and cosine functions and doesnt need training. This means the positional vector of each character in the sequence will be same for all the data in training set.   So  when we add them both together we get the property + position of the characters in a sequence which will help model learn better.   I will only show the things I added in the code blocks so that you can add them accordingly. If theres any modifications  I will change lines to + for added lines and  for removed lines.   This is how you can add PE to the model:   # define our PE Class class PositionalEncoding(nn.Module): def __init__(self  context_length  d_model) -> None: super().__init__() # Create a matrix of shape (context_length  d_model) to store the positional encodings pe = torch.zeros(context_length  d_model) # Create a vector with positions [0  1  2  ...  context_length-1] of shape (context_length  1) position = torch.arange(0  context_length  dtype=torch.float).unsqueeze(1) # Create a vector with the divisor terms based on the dimension div_term = torch.exp(torch.arange(0  d_model  2).float() * (-math.log(10000.0) / d_model)) # Compute the positional encodings using sine and cosine functions pe[:  0::2] = torch.sin(position * div_term) pe[:  1::2] = torch.cos(position * div_term) pe = pe.unsqueeze(0) # Shape: (1  context_length  d_model) # Register pe as a buffer  so it is not considered a parameter but is part of the module's state self.register_buffer('pe'  pe) def forward(self  x: torch.Tensor) -> torch.Tensor: # Add the positional encodings to the input embeddings return x + self.pe[: :x.size(1)  :] class GPT(nn.Module): def __init__(self  vocab_size  d_model): ... # initialize positional encodings self.wpe = PositionalEncoding(context_length  d_model) def forward(self  inputs  targets = None): logits = self.wte(inputs) # pass logits to the PE logits = self.wpe(logits) ... return logits  loss ...Now  if you try to train the model and generate a sequence  you would get an error like below:   This basically means we tried generating 1000 tokens one by one and passing previous n tokens to model for getting next token. But now that we have a PositionalEmbedding layer  it only expects token of size less than or equal to the context_length which is 256 in our case.   Lets modify our generate function to accommodate the context_length:   def generate(self  inputs  max_new_tokens): # this will store the model outputs along with the initial input sequence # make a copy so that it doesn't interfare with model output = inputs.clone() for _ in range(max_new_tokens): current_seq_length = inputs.size(1) # Truncate inputs if it exceeds context_length if current_seq_length > context_length: inputs = inputs[:  -context_length:] ... output = torch.cat([output  idx_next]  dim=1) return [decode(out.tolist()) for out in output]We can already train our model and observe improvements  but before jumping into that  lets add one more layer of mapping. Recall how we are currently obtaining different representations of characters and feeding them into the model. How beneficial would it be if we had additional networks to combine this information and learn more complex representations of the embedding?   That would be the Fully Connected Networks. Lets add PyTorchs Linear Layer in our model:   Code:   class GPT(nn.Module): def __init__(self  vocab_size  d_model): ... self.fcn = nn.Sequential( nn.Linear(d_model  4 * d_model)  nn.GELU()  nn.Linear(4 * d_model  d_model) ) def forward(self  inputs  targets = None): ... logits = self.fcn(logits) ... return logits  lossThats it  simple as that !! Now lets train and evaluate the performance of our model. Im setting the epochs to 5000 and learning rate to 1e-3 for this run.   Maybe not much of an improvement but its now starting to form correct words which it is learning through the position of the characters. Lets keep going shall we?   b. (Masked) Self-Attention + NormalizationU+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525U+1F525   ATTENTION HERE   This is the most interesting part of a transformer model: Self-Attention. To make this concept more clear  refer to the visuals from Jay Alammar. In simple terms  Self-Attention defines which next token the model should pay more attention to given the current and previous n tokens.   It does this by assigning scores to the embedding of each character (in our case) and combines them based on different contexts using Queries  Keys  and Values.   Now enough of the theory  lets get into coding:   Heres how you can add Self-Attention to your model:   class SelfAttention(nn.Module): def __init__(self  d_model: int): super().__init__() self.query = nn.Linear(d_model  d_model) self.key = nn.Linear(d_model  d_model) self.value = nn.Linear(d_model  d_model) self.fc_out = nn.Linear(d_model  d_model) self.dropout = nn.Dropout(0.2) def forward(self  inputs: torch.Tensor): B  seq_length  d_model = inputs.shape # Project the input embeddings into Q  K  and V Q = self.query(inputs) K = self.key(inputs) V = self.value(inputs) # Compute attention scores attention_scores = torch.matmul(Q  K.transpose(-2  -1)) # Apply mask to prevent attention to future tokens mask = torch.triu(torch.ones(seq_length  seq_length)  diagonal=1).bool().to(inputs.device) attention_scores = attention_scores.masked_fill(mask  float('-inf')) attention_weights = torch.softmax(attention_scores  dim=-1) # Compute the weighted sum of the values attention_output = torch.matmul(attention_weights  V) # Apply the final linear transformation out = self.fc_out(attention_output) return out class GPT(nn.Module): def __init__(self  vocab_size  d_model): ... self.att = SelfAttention(d_model) def forward(self  inputs  targets = None): ... logits = self.att(logits) logits = self.fcn(logits) ... return logits  lossSimple as that. Now lets train the model and see the outcome:   U+1F60DWOW. Is is only me or do you also think that the model is now starting to understand a lot of word representation and how its put together in a song? Thats pretty impressive. Wait till this layer gets Multi-Head.   Normalization   One thing you may notice if you are training your model along with me is that the losses are decreasing very quickly  and the model is starting to overfit the data. This can happen because the model is becoming too large relative to our limited training data.   To mitigate this  lets add few LayerNorm and Dropout layers to balance out the learning.   Code:   class GPT(nn.Module): def __init__(self  vocab_size  d_model): ... self.ln1 = nn.LayerNorm(d_model) self.ln2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(0.2) def forward(self  inputs  targets = None): ... logits = self.wte(inputs) logits = self.wpe(logits) att_logits = self.att(logits) adn_logits = self.ln1(logits + att_logits) logits = self.dropout(adn_logits) logits = self.fcn(logits) logits = self.ln2(logits + adn_logits) ... return logits  loss ...This will help us train model for a longer period without over-fitting the dataset.   Quick Change   Now as thats done  I want you to remember one thing from the last part where we set the d_model=vocab_size because we only had one layer which is Embedding.   Well  as now we have a proper mapping layers using Linear  we can change our embedding size to desired number and learn more representation of the character. Lets make it 512.   # used to define size of embeddings d_model = 512 # dont forget to add a linear layer which transforms embedding dim(d_model) to vocab_size class GPT(nn.Module): def __init__(self  vocab_size  d_model): ... self.linear1 = nn.Linear(d_model  vocab_size) def forward(self  inputs  targets = None): ... logits = self.ln2(logits + adn_logits) logits = self.linear1(logits) ... return logits  loss ...By doing just this change  we have completed full model of our GPT-2 Transformer decoder architecture:   But were not done yet. Lets continue improving the model..   c. (Masked) Multi-Head AttentionYou might already be familiar with the power of the Self-Attention mechanism and how it enhances a models ability to generalize contextual relationships within texts. But what if I told you theres a way for the model to understand different linguistic properties within the text  such as how words or characters are interconnected and their temporal usage? Imagine the model learning distinctions between consonants and vowels  and when and where to appropriately use them.   Sounds intriguing  doesnt it? While representing the overall sequence context in Self-Attention with d_model  we can now divide d_model into multiple heads. Each head will have its own sets of representations for Query  Key  and Value  enabling the model to learn multiple contextual nuances within the sequence. Let's enhance our attention layer by incorporating multiple heads.   Code:   n_heads = 4 # number of self-attention heads. should be divisible with d_model class MultiHeadAttention(nn.Module): def __init__(self  d_model: int  n_heads: int): super().__init__() self.n_heads = n_heads self.head_dim = d_model // n_heads assert (n_heads * self.head_dim == d_model) self.query = nn.Linear(d_model  d_model) self.key = nn.Linear(d_model  d_model) self.value = nn.Linear(d_model  d_model) self.fc_out = nn.Linear(d_model  d_model) self.dropout = nn.Dropout(0.2) def forward(self  inputs: torch.Tensor): B  seq_length  d_model = inputs.shape # Project the input embeddings into Q  K  and V Q = self.query(inputs).view(B  seq_length  self.num_heads  self.head_dim).permute(0  2  1  3) K = self.key(inputs).view(B  seq_length  self.num_heads  self.head_dim).permute(0  2  1  3) V = self.value(inputs).view(B  seq_length  self.num_heads  self.head_dim).permute(0  2  1  3) # Compute attention scores attention_scores = torch.matmul(Q  K.transpose(-2  -1)) / math.sqrt(self.head_dim) # Apply mask to prevent attention to future tokens mask = torch.triu(torch.ones(seq_length  seq_length)  diagonal=1).bool().to(inputs.device) attention_scores = attention_scores.masked_fill(mask  float('-inf')) attention_weights = torch.softmax(attention_scores  dim=-1) # Compute the weighted sum of the values attention_output = torch.matmul(self.dropout(attention_weights)  V) # Concatenate heads and put them back to the original shape attention_output = attention_output.permute(0  2  1  3).contiguous() attention_output = attention_output.view(B  seq_length  d_model) # Apply the final linear transformation out = self.fc_out(attention_output) return out class GPT(nn.Module): def __init__(self  vocab_size  d_model  n_heads): super().__init__() ... # replace selfattention layer with multiheadattention self.att = MultiHeadAttention(d_model  n_heads) ... m = GPT(vocab_size=vocab_size  d_model=d_model  n_heads=n_heads).to(device)Now  sit back  let the model train and see the magic   You should now see a significant improvement in the model performance and output. All thanks to the Multi-Head attention. You can play around with the head size to see if the model learn any better representations.   d. GPT Decoder BlocksIf you carefully go through the model diagrams presented throughout the project  you might notice I have starting adding few layers inside a rectangular blocks. They are called decoder blocks. And  just like we can add multiple layers of Linear network  we can also add multiple blocks of those group of models. Lets see how:   Well first take out our Attention  Layer Norms and Feed Forward network into a separate module GPTBlock.   class GPTBlock(nn.Module): def __init__(self  d_model  n_heads): super().__init__() self.att = MultiHeadAttention(d_model  n_heads) self.ln1 = nn.LayerNorm(d_model) self.ln2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(0.2) self.fcn = nn.Sequential( nn.Linear(d_model  4 * d_model)  nn.GELU()  nn.Linear(4 * d_model  d_model) ) def forward(self  logits): att_logits = self.att(logits) adn_logits = self.ln1(logits + att_logits) logits = self.dropout(adn_logits) logits = self.fcn(logits) logits = self.ln2(logits + adn_logits) return logitsNow modify our GPT class to incorporate the block in-place of all these layers inside it along with a constructor parameter n_layer to define number of decoder blocks/layers.   n_layers = 2 # number of gpt blocks/layers class GPT(nn.Module): def __init__(self  vocab_size  d_model  n_heads  n_layers): super().__init__() self.wte = nn.Embedding(vocab_size  d_model) # word token embeddings self.wpe = PositionalEncoding(context_length  d_model) # word position encodings self.blocks = nn.ModuleList([GPTBlock(d_model  n_heads) for _ in range(n_layers)]) self.linear1 = nn.Linear(d_model  vocab_size) def forward(self  inputs  targets = None): logits = self.wte(inputs) # dim -> batch_size  sequence_length  d_model logits = self.wpe(logits) for block in self.blocks: logits = block(logits) logits = self.linear1(logits) ... return logits  loss ... m = GPT(vocab_size=vocab_size  d_model=d_model  n_heads=n_heads  n_layers=n_layers).to(device)e. Improving TokenizerNow theres one more fix that I want to do in our code and its the Tokenizer. Yes  the character level tokenizer which has been overloading our model with tons of tokens with a very little information. Lets improve our tokenizer using the tiktoken library which is an official python library by OpenAI for GPT tokenizers.   The library uses Byte Pair Encoding(BPE) algorithm which creates merges of words or different section of words based on how often they appeared on training the tokenizer.   Installation:   pip install tiktokenCode:   import tiktoken tokenizer = tiktoken.get_encoding('gpt2') vocab_size = tokenizer.n_vocabWe have now increased our vocab size to 50257 which means model gets to see many variations of words and sequences.   Now lets encode our data using the new tokenizer. We will modify our data initialization as:   import torch # use cpu or gpu based on your system device = cpu if torch.cuda.is_available(): device = cuda data_dir = data.txt text = open(data_dir  'r').read() # load all the data as simple string # convert our text data into tokenized tensor data = torch.tensor(tokenizer.encode(text)  dtype=torch.long  device=device)Then  replace any calls to your previous encoding (encode) and decoding (decode) functions with tokenizer.encode() and tokenizer.decode()  respectively. This adjustment ensures compatibility with the new tokenizer.   f. Final GPT-2 TrainingU+1F973 We have finally reached towards the end of the project and quite a new learning experience. We just have to made few adjustments so that our model trains faster and better. And then we are good to go.   Lets do few changes. You can adjust these based on your requirements and system support.   context_length = 512 # number of tokens processed in a single batch d_model = 512 n_layers = 1 # number of gpt blocks/layers class GPT(nn.Module): def __init__(self  vocab_size  d_model  n_heads  n_layers): super().__init__() self.wte = nn.Embedding(vocab_size  d_model) # word token embeddings self.wpe = PositionalEncoding(context_length  d_model) # word position encodings self.blocks = nn.ModuleList([GPTBlock(d_model  n_heads) for _ in range(n_layers)]) self.linear1 = nn.Linear(d_model  vocab_size) # parameter sharing + self.wte.weight = self.linear1.weightTo learn more about parameter sharing in GPT-2  learn here.   You can visualize current model structure by just printing the model variable itself:   And just like that  we have built our own 29M GPT-2 Model which will be sufficient for our use case.   Now  before training our model  lets compile it using torch.compile. It ensures that almost all the matrix multiplications and other operations that happens within the model are mapped before hand. And in simple words  the model can directly compute the final stage by merging all the operations instead of going line by line or layer by layer.   m = GPT(vocab_size=vocab_size  d_model=d_model  n_heads=n_heads  n_layers=n_layers).to(device) m = torch.compile(m)Ive also modified our learning rate and training loop as below:   lr = 1e-3 optim = torch.optim.AdamW(m.parameters()  lr=lr  weight_decay=0.1) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optim  T_max=3000  eta_min=lr*0.1)epochs = 3500 eval_steps = 100 # perform evaluation in every n steps # store the losses train_loss = {} # train the model for e in range(epochs): xb  yb = train_loader.get_batch() logits  loss = m(xb  yb) optim.zero_grad(set_to_none=True) loss.backward() # gradient clipping torch.nn.utils.clip_grad_norm_(m.parameters()  max_norm=1) optim.step() scheduler.step() train_loss[e] = loss.item() if e % eval_steps == 0 or e == epochs-1: m.eval() with torch.no_grad(): xvb  yvb = eval_loader.get_batch() _  e_loss = m(xvb  yvb) print(fEpoch: {e+ep}\ttrain_loss: {loss:.4f}\teval_loss: {e_loss:.4f}) m.train() # Back to training modeAs the model has gotten much bigger now  I am using GoogleColab to train the model. You can acce this link to open on Colab.   After training for ~3500 epochs  I got the following training loss curve:   And finally  a song output from model U+1F3B6U+1F3B6:   There it was. A step-by-step guide to building a custom GPT-2 model and training on own data. Feel free to modify the hyper-parameters and layers to feed according to your data needs.   Thats it for this project. But I wont stop here. I have been planning for few new articles regarding improving model performance and training. So stay tuned. Happy learning :).