| 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"}) | |