move to lit gpt llama impl

Dockerfile

@@ -36,6 +36,7 @@ RUN echo "export PATH=\"/opt/conda/bin:/root/.cargo/bin:\$PATH\"" >> /root/.bash
 WORKDIR open_diloco
 RUN pip install --pre torchdata --index-url https://download.pytorch.org/whl/nightly/cpu
 RUN pip install flash-attn>=2.5.8
+
 COPY requirements.txt requirements.txt
 RUN pip install --no-cache-dir -r requirements.txt
 COPY requirements-dev.txt requirements-dev.txt

open_diloco/configs/config_150m.json

@@ -1,13 +1,10 @@
-{
-    "architectures": [
-      "LlamaForCausalLM"
-    ],
-    "model_type": "llama",
-    "hidden_size": 1024,
-    "intermediate_size": 2688,
-    "num_attention_heads": 16,
-    "num_hidden_layers": 12,
-    "use_cache": false,
-    "rms_norm_eps": 1e-05
-}
 
+{
+  "name": "llama150m",
+  "n_embd": 1024,
+  "intermediate_size": 4096,
+  "n_head": 16,
+  "n_layer": 12,
+  "vocab_size": 32000,
+  "block_size": 1024
+}

open_diloco/configs/config_1b.json

@@ -1,13 +1,10 @@
 {
-  "architectures": [
-    "LlamaForCausalLM"
-  ],
-  "model_type": "llama",
-  "hidden_size": 2048,
+  "name": "llama1b",
+  "n_embd": 2048,
   "intermediate_size": 5632,
-  "num_attention_heads": 32,
-  "num_hidden_layers": 22,
-  "use_cache": false,
-  "rms_norm_eps": 1e-05,
-  "num_key_value_heads": 4
+  "n_head": 32,
+  "n_layer": 22,
+  "n_query_groups": 4,
+  "vocab_size": 32000,
+  "block_size": 1024
 }

open_diloco/configs/config_2m.json

@@ -1,15 +1,10 @@
 {
-    "architectures": [
-      "LlamaForCausalLM"
-    ],
-    "model_type": "llama",
-    "hidden_size": 64,
+    "name": "llama_2m",
+    "n_embd": 64,
     "intermediate_size": 256,
-    "num_attention_heads": 2,
-    "num_hidden_layers": 2,
-    "rms_norm_eps": 1e-05,
-    "use_cache": false,
+    "n_head": 2,
+    "n_layer": 2,
     "vocab_size": 1024
-  }
+}
 
 

open_diloco/configs/config_7b.json

@@ -0,0 +1,10 @@
+{
+    "name": "llama7b",
+    "n_embd": 4096,
+    "intermediate_size": 11008,
+    "n_head": 32,
+    "n_layer": 32,
+    "n_query_groups": null,
+    "vocab_size": 32000,
+    "block_size": 1024
+}