pretrain core

scripts/pretrain_core_model_0.yaml

@@ -10,7 +10,7 @@ model_config:
   vocab_size: 131072
   padded_vocab_size: 131072
   n_layer: 32
-  n_head: 32
+  n_head: 16
   n_embd: 512
   n_query_groups: 8
   rotary_percentage: 1.0

scripts/requirements.in

@@ -14,7 +14,7 @@ mergekit @ git+https://github.com/arcee-ai/mergekit.git
 torchao
 # bitsandbytes
 # grokadamw
-# sophia-opt
+sophia-opt
 dolphinflow @ git+https://github.com/cognitivecomputations/dolphinflow-optimizer.git
 # unsloth
 lm_eval[ifeval,math]