good benchmark. pepe sad

gptj.py

@@ -5,11 +5,6 @@ from time import perf_counter, perf_counter_ns
 import numpy as np
 from tqdm import tqdm
 from contextlib import contextmanager
-from transformers import (
-    AutoModelForCausalLM,
-    GPTNeoForCausalLM,
-    AutoConfig,
-)
 #replicating timeit magic function of ipython
 def timeit(func, r=1, n=5, quiet=False, function=None, do_tqdm=False, first=True):
     precision = 'ns'
@@ -64,13 +59,12 @@ def timeit(func, r=1, n=5, quiet=False, function=None, do_tqdm=False, first=True
 
 
 with torch.no_grad():
-    model = load_gpt_j().cuda().half().eval()
+    model = load_gpt_j().cuda().half()
     x = torch.zeros(1, 2048).cuda().long()
-    our = model(x)
-    print(our.shape)
-    del model
-    model = no_init(lambda: AutoModelForCausalLM.from_pretrained('/home/xuser/models/j6b_ckpt_14001')).cuda().half().eval()
-    hf = model(x, output_hidden_states=True)["hidden_states"][-1]
-    print(our[0, 2047, 1000:1020])
-    print(hf[0, 2047, 1000:1020])
-    print(hf.shape)
+    print(model(x).shape)
+    print("PyTorch Eager")
+    timeit(r=1, n=100, func=lambda: model(x), do_tqdm=False, first=False)
+    module = torch.jit.trace(model, torch.zeros((1, 2048)).long().cuda())
+    torch.jit.optimize_for_inference(module)
+    print("PyTorch JIT")
+    timeit(r=1, n=100, func=lambda: module(x), do_tqdm=False, first=False)
\ No newline at end of file
\ No newline at end of file

hfport.py

 from main import *
 
-state_dict = SplitCheckpoint("/home/xuser/models/j6b_ckpt_14001", device="cpu")
+state_dict = SplitCheckpoint("j6b_vanilla", device="cpu")
 
 
 # ORIGINAL
@@ -62,4 +62,4 @@ def save(state_dict, path):
         torch.save(x[1], f"{path}/b{i}.pt")
     torch.save(checkpoint, f"{path}/m.pt")
 
-save(new_state_dict, "models/6b")
+save(new_state_dict, "models/6b_vanilla")
\ No newline at end of file