File size: 1,168 Bytes
8e70045 | 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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | import torch
from transformers import GPT2Config, GPT2LMHeadModel
ckpt = torch.load("ckpt.pt", map_location="cpu")
state = ckpt["model"] if "model" in ckpt else ckpt
# Remove DDP prefix if present
state = {
k.replace("_orig_mod.", ""): v
for k, v in state.items()
}
config = GPT2Config(
vocab_size=100277,
n_positions=64,
n_ctx=64,
n_embd=128,
n_layer=4,
n_head=4,
bos_token_id=100257,
eos_token_id=100257,
)
model = GPT2LMHeadModel(config)
new_state = {}
transpose = [
"attn.c_attn.weight",
"attn.c_proj.weight",
"mlp.c_fc.weight",
"mlp.c_proj.weight",
]
for k, v in state.items():
hf = k
hf = hf.replace("transformer.wte", "transformer.wte")
hf = hf.replace("transformer.wpe", "transformer.wpe")
hf = hf.replace("transformer.h", "transformer.h")
hf = hf.replace("transformer.ln_f", "transformer.ln_f")
if any(hf.endswith(x) for x in transpose):
v = v.t()
new_state[hf] = v
missing, unexpected = model.load_state_dict(new_state, strict=False)
print("Missing:", missing)
print("Unexpected:", unexpected)
model.save_pretrained("hf_model", safe_serialization=True)
|