File size: 7,304 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 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | """
Script to convert GPT2 models from llm.c binary format to Hugging Face
It can optinally upload to your account on Hugging Face if you have the CLI:
pip install -U "huggingface_hub[cli]"
huggingface-cli login
Export to a local HF model:
python export_hf.py --input input_file.bin --output output_dir
Export to a local HF model and also push to your account on Hugging Face:
python export_hf.py --input input_file.bin --output output_dir --push true
"""
import numpy as np
import torch
import argparse, sys
from transformers import GPT2Config, GPT2Tokenizer, GPT2LMHeadModel
# -----------------------------------------------------------------------------
# Tensor functions for both bfloat16 (from int16) and normal float32
# Both return float32 tensors
def tensor_bf16(data_int16, transpose=False):
if transpose:
data_int16 = data_int16.transpose(1,0)
return torch.tensor(data_int16).view(torch.bfloat16).to(torch.float32)
def tensor_fp32(data_float32, transpose=False):
if transpose:
data_float32 = data_float32.transpose(1,0)
return torch.tensor(data_float32).view(torch.float32)
# -----------------------------------------------------------------------------
# Main conversion function
def convert(filepath, output, push_to_hub=False, out_dtype="bfloat16"):
print(f"Converting model {filepath} to {output} in {out_dtype} format and pushing to Hugging Face: {push_to_hub}")
f = open(filepath, 'rb')
# Read in our header, checking the magic number and version
# version 3 = fp32, padded vocab
# version 5 = bf16, padded vocab
model_header = np.frombuffer(f.read(256*4), dtype=np.int32)
if model_header[0] != 20240326:
print("ERROR: magic number mismatch in the data .bin file!")
exit(1)
version = model_header[1]
if not version in [3, 5]:
print("Bad version in model file")
exit(1)
# Load in our model parameters
maxT = model_header[2].item() # max sequence length
V = model_header[3].item() # vocab size
L = model_header[4].item() # num layers
H = model_header[5].item() # num heads
C = model_header[6].item() # channels
Vp = model_header[7].item() # padded vocab size
print(f"{version=}, {maxT=}, {V=}, {Vp=}, {L=}, {H=}, {C=}")
# Define the shapes of our parameters
shapes = {
'wte': (Vp, C),
'wpe': (maxT, C),
'ln1w': (L, C),
'ln1b': (L, C),
'qkvw': (L, 3 * C, C),
'qkvb': (L, 3 * C),
'attprojw': (L, C, C),
'attprojb': (L, C),
'ln2w': (L, C),
'ln2b': (L, C),
'fcw': (L, 4 * C, C),
'fcb': (L, 4 * C),
'fcprojw': (L, C, 4 * C),
'fcprojb': (L, C),
'lnfw': (C,),
'lnfb': (C,),
}
# Load in our weights given our parameter shapes
dtype = np.float32 if version == 3 else np.int16
w = {}
for key, shape in shapes.items():
num_elements = np.prod(shape)
data = np.frombuffer(f.read(num_elements * np.dtype(dtype).itemsize), dtype=dtype)
w[key] = data.reshape(shape)
# The binary file saves the padded vocab - drop the padding back to GPT2 size
if shape[0] == Vp:
w[key] = w[key].reshape(shape)[:(V-Vp), :]
# Ensure the file is fully read and then close
assert f.read() == b''
f.close()
# Map to our model dict, the tensors at this stage are always fp32
mk_tensor = {
3 : tensor_fp32,
5 : tensor_bf16,
}[version]
model_dict = {}
model_dict['transformer.wte.weight'] = mk_tensor(w['wte'])
model_dict['transformer.wpe.weight'] = mk_tensor(w['wpe'])
model_dict['lm_head.weight'] = model_dict['transformer.wte.weight'] # Tie weights
for i in range(L):
model_dict[f'transformer.h.{i}.ln_1.weight'] = mk_tensor(w['ln1w'][i])
model_dict[f'transformer.h.{i}.ln_1.bias'] = mk_tensor(w['ln1b'][i])
model_dict[f'transformer.h.{i}.attn.c_attn.weight'] = mk_tensor(w['qkvw'][i], True)
model_dict[f'transformer.h.{i}.attn.c_attn.bias'] = mk_tensor(w['qkvb'][i])
model_dict[f'transformer.h.{i}.attn.c_proj.weight'] = mk_tensor(w['attprojw'][i], True)
model_dict[f'transformer.h.{i}.attn.c_proj.bias'] = mk_tensor(w['attprojb'][i])
model_dict[f'transformer.h.{i}.ln_2.weight'] = mk_tensor(w['ln2w'][i])
model_dict[f'transformer.h.{i}.ln_2.bias'] = mk_tensor(w['ln2b'][i])
model_dict[f'transformer.h.{i}.mlp.c_fc.weight'] = mk_tensor(w['fcw'][i], True)
model_dict[f'transformer.h.{i}.mlp.c_fc.bias'] = mk_tensor(w['fcb'][i])
model_dict[f'transformer.h.{i}.mlp.c_proj.weight'] = mk_tensor(w['fcprojw'][i], True)
model_dict[f'transformer.h.{i}.mlp.c_proj.bias'] = mk_tensor(w['fcprojb'][i])
model_dict['transformer.ln_f.weight'] = mk_tensor(w['lnfw'])
model_dict['transformer.ln_f.bias'] = mk_tensor(w['lnfb'])
# Create a GPT-2 model instance, in the requested dtype
config = GPT2Config(vocab_size = V,
n_positions = maxT,
n_ctx = maxT,
n_embd = C,
n_layer = L,
n_head = H)
model = GPT2LMHeadModel(config)
if out_dtype == "bfloat16":
model = model.to(torch.bfloat16)
# Set the model dict and save
model.load_state_dict(model_dict)
model.save_pretrained(output, max_shard_size="5GB", safe_serialization=True)
# Copy over a standard gpt2 tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokenizer.save_pretrained(output)
if push_to_hub:
print(f"Uploading {output} to Hugging Face")
model.push_to_hub(output)
tokenizer.push_to_hub(output)
def spin(output):
print("Taking the exported model for a spin...")
print('-'*80)
from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(output)
model = AutoModelForCausalLM.from_pretrained(output, attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16, device_map='cuda')
model.eval()
tokens = tokenizer.encode("During photosynthesis in green plants", return_tensors="pt")
tokens = tokens.to('cuda')
output = model.generate(tokens, max_new_tokens=64, repetition_penalty=1.3)
samples = tokenizer.batch_decode(output)
for sample in samples:
print('-'*30)
print(sample)
# -----------------------------------------------------------------------------
if __name__== '__main__':
parser=argparse.ArgumentParser()
parser.add_argument("--input", "-i", help="The name of the llm.c model.bin file", type=str, required=True)
parser.add_argument("--output","-o", help="The Hugging Face output model directory", type=str, required=True)
parser.add_argument("--dtype", "-d", help="Output as either float32 or bfloat16 (default)", type=str, default="bfloat16")
parser.add_argument("--push", "-p", help="Push the model to your Hugging Face account", type=bool, default=False)
parser.add_argument("--spin", "-s", help="Take the model for a spin at the end?", type=bool, default=True)
args = parser.parse_args()
convert(args.input, args.output, args.push, args.dtype)
if args.spin:
spin(args.output)
|