Upload folder using huggingface_hub

config.json

@@ -0,0 +1,24 @@
+{
+    "architectures": [
+      "LlamaForCausalLM"
+    ],
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "hidden_act": "silu",
+    "hidden_size": 2048,
+    "initializer_range": 0.02,
+    "intermediate_size": 5504,
+    "max_position_embeddings": 2048,
+    "model_type": "llama",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 27,
+    "num_key_value_heads": 32,
+    "pretraining_tp": 1,
+    "rms_norm_eps": 1e-05,
+    "rope_scaling": null,
+    "tie_word_embeddings": false,
+    "torch_dtype": "bfloat16",
+    "transformers_version": "4.31.0.dev0",
+    "use_cache": true,
+    "vocab_size": 32768
+  }

convert_weights.py

@@ -0,0 +1,81 @@
+import torch
+from tqdm import tqdm
+
+
+input_dir_path = "/scratch/project_465000144/dasamuel/looped-lm/mimir_core_cold/global_step64000"
+output_dir_path = "/scratch/project_465000144/dasamuel/looped-lm/mimir_core_cold_hf"
+
+n_hidden = 4096
+n_ffn_hidden = 14336
+n_heads = 32
+n_kv_heads = 8
+n_layers = 32
+n_tp = 2
+
+
+weights = {}
+
+# embedding
+embedding_weights = []
+for i in range(n_tp):
+    path = f"{input_dir_path}/layer_01-model_0{i}-model_states.pt"
+    checkpoint = torch.load(path)
+
+    embedding_weights.append(checkpoint["word_embeddings.weight"].bfloat16())
+
+weights[f"model.embed_tokens.weight"] = torch.cat(embedding_weights, dim=0)
+del embedding_weights
+
+lm_head_weights = []
+for i in range(n_tp):
+    path = f"{input_dir_path}/layer_{n_layers + 5}-model_0{i}-model_states.pt"
+    checkpoint = torch.load(path)
+
+    lm_head_weights.append(checkpoint["lm_head.weight"].bfloat16())
+
+weights[f"lm_head.weight"] = torch.cat(lm_head_weights, dim=0)
+del lm_head_weights
+
+
+# transformer layers
+for layer in tqdm(range(n_layers)):
+    q_weights, k_weights, v_weights, o_weights = [], [], [], []
+    up_weights, gate_weights, down_weights = [], [], []
+
+    for i in range(n_tp):
+        path = f"{input_dir_path}/layer_{layer+3:02d}-model_0{i}-model_states.pt"
+        checkpoint = torch.load(path)
+
+        weights[f"model.layers.{layer}.input_layernorm.weight"] = checkpoint["input_layernorm.weight"].bfloat16()
+        weights[f"model.layers.{layer}.post_attention_layernorm.weight"] = checkpoint["post_attention_layernorm.weight"].bfloat16()
+
+        kv_weight = checkpoint["self_attention.key_value.weight"].bfloat16()
+        k_weight, v_weight = torch.chunk(kv_weight, 2, dim=0)
+        k_weights.append(k_weight)
+        v_weights.append(v_weight)
+
+        q_weights.append(checkpoint["self_attention.query.weight"].bfloat16())
+        o_weights.append(checkpoint["self_attention.dense.weight"].bfloat16())
+        down_weights.append(checkpoint["mlp.dense_4h_to_h.weight"].bfloat16())
+
+        up_gate_weight = checkpoint["mlp.dense_h_to_4h.weight"].bfloat16()
+        up_weight, gate_weight = torch.chunk(up_gate_weight, 2, dim=0)
+        up_weights.append(up_weight)
+        gate_weights.append(gate_weight)
+
+    weights[f"model.layers.{layer}.self_attn.q_proj.weight"] = torch.cat(q_weights, dim=0)
+    weights[f"model.layers.{layer}.self_attn.k_proj.weight"] = torch.cat(k_weights, dim=0)
+    weights[f"model.layers.{layer}.self_attn.v_proj.weight"] = torch.cat(v_weights, dim=0)
+    weights[f"model.layers.{layer}.self_attn.o_proj.weight"] = torch.cat(o_weights, dim=1)
+    weights[f"model.layers.{layer}.mlp.up_proj.weight"] = torch.cat(up_weights, dim=0)
+    weights[f"model.layers.{layer}.mlp.gate_proj.weight"] = torch.cat(gate_weights, dim=0)
+    weights[f"model.layers.{layer}.mlp.down_proj.weight"] = torch.cat(down_weights, dim=1)
+
+
+# output layer norm
+path = f"{input_dir_path}/layer_{n_layers + 4}-model_00-model_states.pt"
+checkpoint = torch.load(path)
+
+weights[f"model.norm.weight"] = checkpoint["weight"].bfloat16()
+
+torch.save(weights, f"{output_dir_path}/pytorch_model.bin")

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:33e72a1eb422f5a2770d4718149ce61365f6aff1b92597964be9ec3b8f78ad19
+size 14496143545

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

tokenizer_config.json

@@ -0,0 +1,51 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "bos_token_id": 1,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "eos_token_id": 2,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "pad_token_id": 3,
+  "padding_side": "right",
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "unk_token_id": 0,
+  "use_default_system_prompt": false
+}