File size: 5,753 Bytes
4d5ab1c | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """
Convert a HuggingFace GPT-2-architecture model to llm.c binary format (version 5, bf16).
Usage:
python import_hf.py --input <hf_model_dir> --output <output.bin>
The output is compatible with train_gpt2cu's -e flag.
"""
import argparse
import math
import struct
import numpy as np
import torch
from transformers import GPT2LMHeadModel, AutoConfig, AutoModelForCausalLM
def fp32_to_bf16_int16(tensor: torch.Tensor) -> np.ndarray:
"""Convert float32 tensor to bfloat16 stored as int16 (same byte layout)."""
bf16 = tensor.to(torch.bfloat16)
return bf16.view(torch.int16).cpu().numpy()
def write_model(model: GPT2LMHeadModel, output_path: str) -> None:
cfg = model.config
V = cfg.vocab_size
maxT = cfg.n_positions
L = cfg.n_layer
H = cfg.n_head
C = cfg.n_embd
# Pad vocab to next multiple of 128
Vp = math.ceil(V / 128) * 128
print(f"V={V}, Vp={Vp}, maxT={maxT}, L={L}, H={H}, C={C}")
# Build header: 256 int32 values
header = np.zeros(256, dtype=np.int32)
header[0] = 20240326 # magic
header[1] = 5 # version 5 = bf16 + padded vocab
header[2] = maxT
header[3] = V
header[4] = L
header[5] = H
header[6] = C
header[7] = Vp
sd = model.state_dict()
def get(key):
return sd[key].float()
with open(output_path, "wb") as f:
f.write(header.tobytes())
# wte: (Vp, C) — pad vocab rows with zeros
wte = get("transformer.wte.weight") # (V, C)
pad = torch.zeros(Vp - V, C, dtype=torch.float32)
wte_padded = torch.cat([wte, pad], dim=0) # (Vp, C)
f.write(fp32_to_bf16_int16(wte_padded).tobytes())
# wpe: (maxT, C)
f.write(fp32_to_bf16_int16(get("transformer.wpe.weight")).tobytes())
# Per-layer weights
for i in range(L):
# ln1w, ln1b: (C,)
f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.ln_1.weight")).tobytes())
for i in range(L):
f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.ln_1.bias")).tobytes())
# qkvw: (L, 3C, C) — HF stores c_attn.weight as (C, 3C), transpose to (3C, C)
for i in range(L):
w = get(f"transformer.h.{i}.attn.c_attn.weight") # (C, 3C)
f.write(fp32_to_bf16_int16(w.T.contiguous()).tobytes()) # (3C, C)
# qkvb: (L, 3C)
for i in range(L):
f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.attn.c_attn.bias")).tobytes())
# attprojw: (L, C, C) — HF stores c_proj.weight as (C, C), transpose
for i in range(L):
w = get(f"transformer.h.{i}.attn.c_proj.weight") # (C, C)
f.write(fp32_to_bf16_int16(w.T.contiguous()).tobytes())
# attprojb: (L, C)
for i in range(L):
f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.attn.c_proj.bias")).tobytes())
# ln2w, ln2b
for i in range(L):
f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.ln_2.weight")).tobytes())
for i in range(L):
f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.ln_2.bias")).tobytes())
# fcw: (L, 4C, C) — HF c_fc.weight is (C, 4C), transpose to (4C, C)
for i in range(L):
w = get(f"transformer.h.{i}.mlp.c_fc.weight") # (C, 4C)
f.write(fp32_to_bf16_int16(w.T.contiguous()).tobytes())
# fcb: (L, 4C)
for i in range(L):
f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.mlp.c_fc.bias")).tobytes())
# fcprojw: (L, C, 4C) — HF c_proj.weight is (4C, C), transpose to (C, 4C)
for i in range(L):
w = get(f"transformer.h.{i}.mlp.c_proj.weight") # (4C, C)
f.write(fp32_to_bf16_int16(w.T.contiguous()).tobytes())
# fcprojb: (L, C)
for i in range(L):
f.write(fp32_to_bf16_int16(get(f"transformer.h.{i}.mlp.c_proj.bias")).tobytes())
# lnfw, lnfb
f.write(fp32_to_bf16_int16(get("transformer.ln_f.weight")).tobytes())
f.write(fp32_to_bf16_int16(get("transformer.ln_f.bias")).tobytes())
size_mb = __import__("os").path.getsize(output_path) / 1e6
print(f"Saved {output_path} ({size_mb:.1f} MB)")
def main():
parser = argparse.ArgumentParser("HF GPT-2 → llm.c bf16 binary")
parser.add_argument("--input", "-i", required=True, help="HF model directory")
parser.add_argument("--output", "-o", required=True, help="Output .bin path")
args = parser.parse_args()
print(f"Loading model from {args.input} ...")
# Support loading when config and weights are in different snapshot dirs.
# If args.input has no config.json, look for one in sibling snapshot dirs.
import os, glob
cfg_path = os.path.join(args.input, "config.json")
if not os.path.exists(cfg_path):
snapshots_dir = os.path.dirname(args.input)
candidates = sorted(glob.glob(os.path.join(snapshots_dir, "*/config.json")))
if not candidates:
raise FileNotFoundError(f"No config.json found near {args.input}")
cfg_path = candidates[0]
print(f"Using config from: {cfg_path}")
from transformers import GPT2Config
config = GPT2Config.from_pretrained(os.path.dirname(cfg_path))
model = GPT2LMHeadModel(config)
from safetensors.torch import load_file
weights_file = os.path.join(args.input, "model.safetensors")
if os.path.exists(weights_file):
sd = load_file(weights_file)
model.load_state_dict(sd, strict=False)
print(f"Loaded safetensors weights from {weights_file}")
else:
model = GPT2LMHeadModel.from_pretrained(args.input, config=config)
model.eval()
write_model(model, args.output)
if __name__ == "__main__":
main()
|