|
import torch |
|
from tqdm import tqdm |
|
|
|
|
|
input_dir_path = "/scratch/project_462000086/norwegian_gpt/Megatron-DeepSpeed-fixed/mistral-7b-from-scratch-2nd-run/global_step30000" |
|
output_dir_path = "/scratch/project_462000086/norwegian_gpt/Megatron-DeepSpeed-fixed/hf_mistral_from_scratch_60k" |
|
|
|
n_hidden = 4096 |
|
n_ffn_hidden = 14336 |
|
n_heads = 32 |
|
n_kv_heads = 8 |
|
n_layers = 32 |
|
n_tp = 2 |
|
|
|
|
|
weights = {} |
|
|
|
|
|
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 |
|
|
|
|
|
|
|
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) |
|
|
|
|
|
|
|
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") |
|
|
