update to resume training

.gitignore

@@ -0,0 +1,3 @@
+wandb
+__pycache__
+.ipynb_checkpoints

app.py

@@ -1,3 +1,4 @@
+# pip install accelerate datasets transformers huggingface_hub wandb gated_state_spaces_pytorch
 import torch
 import torch.nn as nn
 from torch.optim import AdamW
@@ -20,9 +21,9 @@ if __name__ == '__main__':
         entity="naxalpha",
     )
     
-    gpt_2 = GPT2LMHeadModel.from_pretrained('gpt2-xl')
-    gpt_2.requires_grad_(False)
-    gpt_2 = gpt_2.cuda()
+    # gpt_2 = GPT2LMHeadModel.from_pretrained('gpt2-xl')
+    # gpt_2.requires_grad_(False)
+    # gpt_2 = gpt_2.cuda()
 
     f_emb = 1600
     model = AutoregressiveWrapper(
@@ -34,19 +35,20 @@ if __name__ == '__main__':
     )
     wandb.watch(model)
 
-    emb = gpt_2.state_dict()['transformer.wte.weight']
+    # emb = gpt_2.state_dict()['transformer.wte.weight']
 
     model.net.token_emb.weight.requires_grad_(False)
-    model.net.token_emb.weight.copy_(emb)
+    # model.net.token_emb.weight.copy_(emb)
 
     model.net.to_logits.weight.requires_grad_(False)
-    model.net.to_logits.weight.copy_(emb)
+    # model.net.to_logits.weight.copy_(emb)
     
     model.net.to_logits = nn.Sequential(
         nn.LayerNorm(f_emb),
         model.net.to_logits,
     )
-
+    
+    model.load_state_dict(torch.load('model.pt'))
     model = model.cuda()
     optim = AdamW(model.parameters(), 2e-5)
 
@@ -65,18 +67,7 @@ if __name__ == '__main__':
 
     for i, batch in enumerate(prog):
         batch = batch.cuda()
-        if i % 2 == 0:  # distil
-            batch = batch[:, :-1]
-            with torch.no_grad():
-                logits = gpt_2(batch).logits
-                probs = logits.softmax(dim=-1)
-            out = model.net(batch)
-            los = F.cross_entropy(
-                out.flatten(0,1),
-                probs.flatten(0,1),
-            )
-        else:  # scratch
-            los = model(batch)
+        los = model(batch)
 
         (los / k).backward()
         if (i+1) % k == 0: