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