File size: 919 Bytes
b9be04a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
from transformers import AutoModelForCausalLM
import torch
from safetensors.torch import save_file
model = AutoModelForCausalLM.from_pretrained("microsoft/Phi-3-mini-128k-instruct", trust_remote_code=True)
params = model.state_dict()
params2 = {}
for r in params.keys():
if "gate_up_proj" in r:
(gate, up) = params[r].chunk(2)
params2[r.replace("gate_up_proj", "gate_proj")] = gate
params2[r.replace("gate_up_proj", "up_proj")] = up
elif "qkv_proj" in r:
(q, k, v) = params[r].chunk(3)
params2[r.replace("qkv_proj", "q_proj")] = q
params2[r.replace("qkv_proj", "k_proj")] = k
params2[r.replace("qkv_proj", "v_proj")] = v
else:
params2[r] = params[r]
for r in params2.keys():
params2[r] = torch.tensor(params2[r].clone().detach(), dtype=torch.bfloat16)
save_file(params2, "model-00001-of-00001.safetensors", metadata={"format": "pt"})
|