From c165855cd331ac16de684e6d15507eabdc2f1a8d Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Mon, 15 Apr 2024 01:57:04 +0000 Subject: [PATCH] add optional support for the tokenizer. people have to re-run train_gpt2.py to get this feature. later, we will change the code to demand this and error if it is not found, for now i just don't want to brick people who do a pull. also i only modified the train_gpt2.cu for now, not the .c. it might be time to separate out the dataloader, the tokenizer, and any other common utilities to their own files (?). otherwise it feels a bit silly to copy paste stuff around? not sure --- train_gpt2.cu | 101 +++++++++++++++++++++++++++++++++++++++++++++++--- train_gpt2.py | 18 +++++++++ 2 files changed, 114 insertions(+), 5 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 239782b..0f3ea00 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -4,6 +4,7 @@ GPT-2 Transformer Neural Net trained in raw CUDA #include #include +#include #include #include #include @@ -1158,6 +1159,86 @@ int sample_mult(float* probabilities, int n, float coin) { return n - 1; // in case of rounding errors } +// ---------------------------------------------------------------------------- +// Tokenizer (only supports decoding) + +typedef struct { + uint32_t vocab_size; + char **token_table; + int init_ok; +} Tokenizer; + +void safe_printf(const char *piece) { + // the tokens are raw bytes, and we we only want to print the printable ones + // many bytes can be various control codes, backspace, etc. + if (piece == NULL) { return; } + if (piece[0] == '\0') { return; } + // handle individual byte tokens + // every token is asserted to be at least one byte so doing piece[1] is ok + if (piece[1] == '\0') { + unsigned char byte_val = piece[0]; + if (!(isprint(byte_val) || isspace(byte_val))) { + return; // weird byte, don't print it + } + } + printf("%s", piece); +} + +void tokenizer_init(Tokenizer *tokenizer, const char *filename) { + FILE *file = fopen(filename, "rb"); + if (file == NULL) { + // try to be more helpful as we just added this feature, erase later + printf("---\n"); + printf("WARNING: Failed to open the tokenizer file %s\n", filename); + printf("The Tokenizer is a new feature added April 14 2024.\n"); + printf("Re-run `python train_gpt2.py` to write it\n"); + printf("---\n"); + tokenizer->init_ok = 0; + return; + } + // read in the header + uint32_t header[256]; + fread(header, sizeof(uint32_t), 256, file); + assert(header[0] == 20240328); + assert(header[1] == 1); + tokenizer->vocab_size = header[2]; + // read in all the tokens + unsigned char length; + tokenizer->token_table = (char **)malloc(tokenizer->vocab_size * sizeof(char *)); + for (uint32_t i = 0; i < tokenizer->vocab_size; i++) { + fread(&length, sizeof(unsigned char), 1, file); + assert(length > 0); // every token should be at least one character + char *token_bytes = (char *)malloc(length + 1); + fread(token_bytes, sizeof(char), length, file); + token_bytes[length] = '\0'; // Add null terminator for printing + tokenizer->token_table[i] = token_bytes; + } + // cleanups + fclose(file); + tokenizer->init_ok = 1; +} + +const char *tokenizer_decode(Tokenizer *tokenizer, uint32_t token_id) { + if (tokenizer->init_ok == 0) { + return NULL; + } + if (token_id < tokenizer->vocab_size) { + return tokenizer->token_table[token_id]; + } else { + printf("invalid token id %d!\n", token_id); + return NULL; + } +} + +void tokenizer_free(Tokenizer *tokenizer) { + if (tokenizer->init_ok) { + for (uint32_t i = 0; i < tokenizer->vocab_size; i++) { + free(tokenizer->token_table[i]); + } + free(tokenizer->token_table); + } +} + // ---------------------------------------------------------------------------- // main training loop int main() { @@ -1206,6 +1287,10 @@ int main() { printf("sequence length: %d\n", T); printf("val_num_batches: %d\n", val_num_batches); + // build the Tokenizer + Tokenizer tokenizer; + tokenizer_init(&tokenizer, "gpt2_tokenizer.bin"); + // some memory for generating samples from the model unsigned long long rng_state = 1337; int* gen_tokens = (int*)malloc(B * T * sizeof(int)); @@ -1236,6 +1321,7 @@ int main() { gen_tokens[i] = GPT2_EOT; } // now sample from the model autoregressively + printf("generating:\n---\n"); for (int t = 1; t < genT; t++) { // note that inference is very wasteful here because for each token // we re-calculate the forward pass for all of (B,T) positions from scratch @@ -1252,12 +1338,16 @@ int main() { float coin = random_f32(&rng_state); int next_token = sample_mult(cpu_probs, model.config.vocab_size, coin); gen_tokens[t] = next_token; + // print the generated token, either using the Tokenizer or a fallback + if (tokenizer.init_ok) { + const char* token_str = tokenizer_decode(&tokenizer, next_token); + safe_printf(token_str); + } else { + // fall back to printing the token id + printf("%d ", next_token); + } } - printf("generated: "); - for (int t = 0; t < genT; t++) { - printf("%d ", gen_tokens[t]); - } - printf("\n"); + printf("\n---\n"); } // do a training step @@ -1276,6 +1366,7 @@ int main() { // free dataloader_free(&train_loader); dataloader_free(&val_loader); + tokenizer_free(&tokenizer); gpt2_free(&model); free(cpu_probs); free(gen_tokens); diff --git a/train_gpt2.py b/train_gpt2.py index d9e0073..36c1918 100644 --- a/train_gpt2.py +++ b/train_gpt2.py @@ -11,6 +11,7 @@ https://github.com/huggingface/transformers/blob/main/src/transformers/models/gp import os import math +import struct from dataclasses import dataclass import numpy as np @@ -290,6 +291,21 @@ def write_state(model, x, y, logits, loss, filename): write_tensors(grads, model.config.n_layer, file) print(f"wrote {filename}") +def write_tokenizer(enc, filename): + n = enc.max_token_value + 1 + header = torch.zeros(256, dtype=torch.int32) + header[0] = 20240328 # magic + header[1] = 1 # tokenizer version = 1 + header[2] = n # number of tokens + with open(filename, "wb") as file: + file.write(header.numpy().tobytes()) + for i in range(n): + b = enc.decode_bytes([i]) + length = len(b) + assert length < 256, f"Token length exceeds 255: {length}" + file.write(struct.pack(""}) decode = lambda l: enc.decode(l) + write_tokenizer(enc, "gpt2_tokenizer.bin") + if args.tensorcores: torch.set_float32_matmul_precision('high')