set to gpt2 hyprs

train.py

@@ -11,11 +11,11 @@ import log
 import sys
 
 hypr = {
-    "embed_size":  512,
-    "n_heads":     8,
-    "n_blocks":    6,
-    "block_size":  256,
-    "batch_size":  16,
+    "embed_size":  768,
+    "n_heads":     12,
+    "n_blocks":    12,
+    "block_size":  512,
+    "batch_size":  8,
     "starting_lr": 6e-4,
     "minimum_lr":  6e-5,
     "warmup":      1_000,
@@ -25,6 +25,7 @@ hypr = {
     "subset":      "cosmopedia-v2",
     "chat_dataset": "yahma/alpaca-cleaned",
     "chat_subset":  None,
+    "half":         False,
 }
 
 print(Device.DEFAULT)
@@ -32,6 +33,8 @@ chat = len(sys.argv) > 1
 if(chat):
     hypr["dataset"] = hypr["chat_dataset"]
     hypr["subset"] = hypr["chat_subset"]
+    hypr["starting_lr"] *= 0.1
+    hypr["minimum_lr"] *= 0.1
 
 #for loging
 loger = ThreadPoolExecutor(max_workers=2)
@@ -49,6 +52,9 @@ batch = data.startDataWorker(dataset,encoding,hypr["batch_size"],hypr["block_siz
 model = Transformer(hypr["vocab_size"],hypr["embed_size"],hypr["n_heads"],hypr["n_blocks"],hypr["block_size"])
 if (chat):
     load_state_dict(model,safe_load(sys.argv[1]))
+if hypr["half"]:
+    from tinygrad import dtypes
+    model = model.cast(dtypes.float16)
 params = nn.state.get_parameters(model)
 optimizer = optm.llmOptimizer(params,hypr["steps"],hypr["starting_lr"],hypr["minimum_lr"])