Added new architecture supporting GQA

app.py

@@ -4,7 +4,7 @@ from my_gpt import my_gpt
 from tokenizer.tokenizer import BPE
 
 ##Load model
-model = my_gpt.load_pretrained("model/model_1000_cpu.bin")
+model = my_gpt.load_pretrained("model/model_1000_.bin")
 # model.to(torch.device("cpu"))
 # model.save_pretrained("model/model_1000_cpu.bin")
 # exit()
@@ -21,7 +21,7 @@ def generate(input_text):
 iface = gr.Interface(fn=generate, 
                      inputs="text", 
                      outputs="text",
-                     title="GPT - 1000 steps",
-                     description="""This model is trained for 1000 steps only. It is not
-                     able to generate perfect sentences/words. However, it has learnt a gist of the English language""")
+                     title="NoobGPT - 1000 steps",
+                     description="""This 13M param model is trained for 1000steps only and has seen only 1M tokens. It is not
+                     able to generate perfect sentences/words but has acquired a rudimentary understanding of the English language""")
 iface.launch()

model/{model_1000_cpu.bin → model_1000_.bin}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f88eff25b6947e11a832f96e2bc914c6818989045539c327438c3e490b184cc9
-size 56951293
+oid sha256:5877a72287e65e61deab89188115afa2eb7dade01cbde49c3103fa40b468a1c8
+size 56607625

my_gpt.py

@@ -4,13 +4,14 @@ from torch.nn import functional as F
 import json
 import logging 
 
-
-block_size = 256
+block_size = 128
 vocab_size = 500
 n_embed = 384
 dropout = 0.2
 n_head = 6
 n_layer = 6
+kv_heads = 3
+max_position_embeddings = 128
 
 class Head(nn.Module):
     def __init__(self, head_size=16):
@@ -40,18 +41,102 @@ class Head(nn.Module):
 
         return out
     
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
 class MultiHeadAttention(nn.Module):
-    def __init__(self,num_heads, head_size) :
+    def __init__(self,num_heads, head_dim) :
         super().__init__()
+        assert num_heads%kv_heads == 0
+        self.n_embed = n_embed
+        self.num_attn_heads = num_heads
+        self.head_dim = head_dim
+        self.kv_heads = kv_heads
+        # self.kv_out_proj = head_dim *  self.kv_heads #Doubt
+        self.num_kv_groups = self.num_attn_heads // self.kv_heads
 
-        self.heads = nn.ModuleList(Head(head_size=head_size) for _ in range(num_heads))
-        self.proj = nn.Linear(head_size * num_heads, n_embed)
+
+        self.heads = nn.ModuleList(Head(head_size=head_dim) for _ in range(num_heads))
+        ##Only self attention
+
+        #For num_attn_heads number of heads
+        self.Wq = nn.Linear(self.n_embed, self.num_attn_heads*self.head_dim)
+        #For kv_heads number of heads
+        self.Wk = nn.Linear(self.n_embed, self.kv_heads * self.head_dim)
+        self.Wv = nn.Linear(self.n_embed, self.kv_heads * self.head_dim)
+
+        self.o_proj = nn.Linear(self.head_dim * self.num_attn_heads, self.n_embed)
         self.dropout = nn.Dropout(dropout)
 
-    def forward(self, x):
-        out = torch.cat([h(x) for h in self.heads], dim=-1)
-        out = self.dropout(self.proj(out))
-        return out
+        # self.attention_mask = torch.zeros((bsz, self.num_attn_heads, qlen, qlen))
+        # self.attention_mask[:, :, :, qlen:] = float('-inf')  # Mask out positions beyond the key sequence length
+
+
+    def forward(self, x, attn_mask= None):
+        """
+        Parameters:
+            x (bsz, qlen, embed) : input
+        """
+        # out = torch.cat([h(x) for h in self.heads], dim=-1)
+        # attn_output = self.dropout(self.o_proj(out))
+
+        # ################ Experiment
+
+
+        bsz, qlen, embed = x.size()
+        # print("input size", x.size())
+
+        q = self.Wq(x) ##(B,T,head_dim * num_heads)
+        k = self.Wk(x) ##(B,T,head_dim * kv_heads)
+        v = self.Wv(x) ##(B,T,head_dim * kv_heads)
+
+
+
+        q = q.view(bsz, qlen, self.num_attn_heads, self.head_dim).transpose(2,1)  ##(B,T,head_dim * num_heads)
+        k = k.view(bsz, qlen, self.kv_heads, self.head_dim).transpose(2,1) ##(B,T,head_dim * kv_heads)
+        v = v.view(bsz, qlen, self.kv_heads, self.head_dim).transpose(2,1)  ##(B,T,head_dim * kv_heads)
+
+        # print("k-shape before ",k.shape)
+        k = repeat_kv(k, self.num_kv_groups) ##(B, n_kvheads * nrep, qlen, head_dim)
+        v = repeat_kv(v, self.num_kv_groups)
+
+        attn_scores = q @ k.transpose(-1,-2)/torch.sqrt(torch.tensor(self.n_embed)) ##(B, T, block_size)
+
+        ################
+        # print("Q-shape ", q.shape)
+        # print("k-shape ",k.shape)
+        # print(k.shape[-2])
+        # print(attn_scores.shape)
+        
+        if attn_mask is not None:
+            # causal_mask = attn_mask[:, :, :, : k.shape[-2]]
+            # attn_scores = attn_scores + causal_mask
+            attn_scores = attn_scores.masked_fill(
+                attn_mask[None, None, :qlen, :qlen]==0 , float("-inf")
+            )
+
+
+        attn_scores = F.softmax(attn_scores, dim=-1)
+        attn_scores = F.dropout(attn_scores) ##Why this dropout is required??
+
+        attn_output = torch.matmul(attn_scores, v) ##(B, n_heads, qlen, hidden_size)
+        attn_output = attn_output.transpose(1,2).contiguous()    
+        attn_output = attn_output.view(bsz, qlen, self.n_embed)
+
+        attn_output = self.o_proj(attn_output)
+        return attn_output
+
 
 class FeedForward(nn.Module):
     def __init__(self,n_embed) -> None:
@@ -68,26 +153,33 @@ class FeedForward(nn.Module):
         return x
 
 class decoder_block(nn.Module):
-    def __init__(self, n_embed, n_heads):
+    def __init__(self, n_embed, n_heads, attn_mask=None):
         super().__init__()
-        self.sa = MultiHeadAttention(n_heads,n_embed//n_heads)
+        # Assume 0 for allowed positions and -inf for masked positions
+        
+        self.sa = MultiHeadAttention(n_heads,n_embed//n_head)
         self.ln1 = nn.LayerNorm(n_embed)
         self.ln2 = nn.LayerNorm(n_embed)
         self.ffwd = FeedForward(n_embed)
+        # self.causal_mask = torch.tril(torch.ones(block_size,block_size))
+        self.register_buffer('causal_mask',torch.tril(torch.ones(block_size,block_size)))
+
+
 
     def forward(self, x):
-        x = x + self.sa(self.ln1(x))
+        x = x + self.sa(self.ln1(x), attn_mask = self.causal_mask)
         x = x + self.ffwd(self.ln2(x))
         return x
 
 
 
 class my_gpt(nn.Module):
-    def __init__(self, block_size = 256):
+    def __init__(self, device='cpu', block_size = 128):
         super().__init__()
+        self.device = device
         self.block_size = block_size ##context window size
         self.token_embed = nn.Embedding(vocab_size, n_embed)
-        self.pos_embed = nn.Embedding(vocab_size, n_embed)
+        self.pos_embed = nn.Embedding(max_position_embeddings, n_embed)
         self.lm_head = nn.Linear(n_embed, vocab_size)
         self.sa_head = Head(vocab_size)
         self.d_blocks = nn.Sequential(*[decoder_block(n_embed=n_embed,n_heads=n_head) for _ in range(n_layer)])
@@ -103,19 +195,20 @@ class my_gpt(nn.Module):
         elif isinstance(module, nn.Embedding):
             torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
 
-    def forward(self, idx, targets = None):
+    def forward(self, x, targets = None):
         """
         Args:
-            idx: int(B,T) Token ids
+            x: int(B,T) Token ids
             targets :
 
         Returns:
             logits
         """
         # print("idx ", idx)
-        B, T = idx.shape ##
-        tok_emd = self.token_embed(idx) ##(B,T,C)
-        pos_emd = self.pos_embed(idx)
+        B, T = x.size() ##
+        tok_emd = self.token_embed(x) ##(B,T,C)
+        position_ids = torch.arange(T, device = self.device )
+        pos_emd = self.pos_embed(position_ids)
 
         
         x = tok_emd + pos_emd
@@ -154,6 +247,7 @@ class my_gpt(nn.Module):
         for _ in range(max_new_tokens): 
             ##Take only last allowed number of tokens
             idx_tokens = context[:, -self.block_size:]
+            # print(f"idx tokens {idx_tokens.shape}")
 
             ##generate the next token
             logits, loss = self(idx_tokens)