pretrain core

config-0.json

@@ -16,7 +16,7 @@
   "model_type": "llama",
   "num_attention_heads": 16,
   "num_hidden_layers": 32,
-  "num_key_value_heads": 8,
+  "num_key_value_heads": 4,
   "pretraining_tp": 1,
   "rms_norm_eps": 1e-05,
   "rope_scaling": null,

scripts/pretrain_core_model_0.yaml

@@ -12,7 +12,7 @@ model_config:
   n_layer: 32
   n_head: 16
   n_embd: 512
-  n_query_groups: 8
+  n_query_groups: 4
   rotary_percentage: 1.0
   parallel_residual: False
   bias: False