Delete convert_weights.py

convert_weights.py

@@ -1,84 +0,0 @@
-import torch
-from safetensors.torch import save_file
-from tqdm import tqdm
-
-
-input_dir_path = "/scratch/project_465000144/dasamuel/normistral/normistral-11b-masked-post/global_step10000"
-output_dir_path = "/scratch/project_465000144/dasamuel/normistral/normistral-11b-masked-post-hf-60000"
-
-n_hidden = 5120
-n_ffn_hidden = 14336
-n_heads = 32
-n_kv_heads = 8
-n_layers = 40
-head_dim = 128
-n_tp = 2
-n_pp = 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")
-save_file(weights, f"{output_dir_path}/model.safetensors")