Updated for new model

model.py

@@ -1,5 +1,7 @@
 import torch 
 from torch import nn
+import bitsandbytes as bnb
+from torch.utils.checkpoint import checkpoint
 
 class TransformerBlock(nn.Module):
     def __init__(self, device, n_heads, embed_dim, ffn_dim, dropout):
@@ -28,7 +30,7 @@ class GPTModel(nn.Module):
         super().__init__()
         self.max_seq_len = max_seq_len
         self.device = device
-        self.embedding = nn.Embedding(n_vocab, embed_dim)
+        self.embedding = bnb.nn.StableEmbedding(n_vocab, embed_dim)
         nn.init.normal_(self.embedding.weight, mean=0.0, std=0.02)
         self.positional_embedding = nn.Embedding(max_seq_len, embed_dim)
 
@@ -42,8 +44,11 @@ class GPTModel(nn.Module):
         input_embed = self.embedding(input_tokens)
         positions = torch.arange(0, input_tokens.size(1), device=input_tokens.device).unsqueeze(0)
         input_embed = input_embed + self.positional_embedding(positions)
+        x = input_embed
+
+        for block in self.transformer_blocks:
+            x = checkpoint(block, x, use_reentrant=False)
 
-        x = self.transformer(input_embed)
         x = self.layer_norm(x)
         x = self.output_projection(x)
         return x

requirements.txt

@@ -1,4 +1,5 @@
 torch
 tiktoken
 huggingface_hub
-safetensors
+safetensors
+bitsandbytes