diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index a954469..f8b216b 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -4,10 +4,12 @@ on: push: branches: - master - paths: ['.github/workflows/**', '**/Makefile', '**/*.c', '**/*.h'] + paths: ['.github/workflows/**', '**/Makefile', '**/*.c', '**/*.h', '**/*.py'] pull_request: types: [opened, synchronize, reopened] - paths: ['**/Makefile', '**/*.c', '**/*.h'] + paths: ['**/Makefile', '**/*.c', '**/*.h', '**/*.py'] + # for manual triggering + workflow_dispatch: env: BRANCH_NAME: ${{ github.head_ref || github.ref_name }} @@ -15,7 +17,7 @@ env: jobs: # check basic builds to avoid breaking changes ubuntu-focal-make: - runs-on: ubuntu-20.04 + runs-on: ubuntu-latest steps: - name: Clone @@ -28,6 +30,16 @@ jobs: sudo apt-get update sudo apt-get install build-essential -y + - name: Set up Python 3.10 + uses: actions/setup-python@v3 + with: + python-version: "3.10" + + - name: Pip setup + run: | + python -m pip install --upgrade pip + if [ -f requirements.txt ]; then pip install -r requirements.txt; fi + - name: Build id: make_build run: | @@ -38,6 +50,10 @@ jobs: run: | make runfast + - name: Test with pytest + run: | + pytest + macOS-latest-make: runs-on: macos-latest diff --git a/README.md b/README.md index f34dc4b..8d76a5d 100644 --- a/README.md +++ b/README.md @@ -4,9 +4,11 @@ Cute Llama

-With the code in this repo you can train the Llama 2 LLM architecture from scratch in PyTorch, then export the weights to a binary file, and load that into one ~simple 500-line C file ([run.c](run.c)) that inferences the model. Alternatively, you can load, finetune, and inference Meta's Llama 2 (but this is still being actively fleshed out). Hence, this repo is a "fullstack" train + inference solution for Llama 2 LLM, with a focus on minimalism and simplicity. You might think that you need many billion parameter LLMs to do anything useful, but in fact very small LLMs can have surprisingly strong performance if you make the domain narrow enough. I recommend looking at the [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories) paper for inspiration. +Train the Llama 2 LLM architecture in PyTorch then inference it with one simple 700-line C file ([run.c](run.c)). You might think that you need many billion parameter LLMs to do anything useful, but in fact very small LLMs can have surprisingly strong performance if you make the domain narrow enough (ref: [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories) paper). This repo is a "fullstack" train + inference solution for Llama 2 LLM, with focus on minimalism and simplicity. -Please note that this started recently as just a fun weekend project: I took my earlier [nanoGPT](https://github.com/karpathy/nanoGPT), tuned it to implement the Llama-2 architecture instead of GPT-2, and the meat of it was writing the C inference engine in [run.c](run.c). So the project is young and moving quickly. Hat tip to the awesome [llama.cpp](https://github.com/ggerganov/llama.cpp) for inspiring this project. I wanted something super minimal so I chose to hard-code the Llama 2 architecture, stick to fp32, and just roll one inference file of pure C with no dependencies. +As the architecture is identical, you can also load and inference Meta's Llama 2 models. However, the current code only inferences models in fp32, so you will most likely not be able to productively load models larger than 7B. Work on model quantization is currently ongoing. + +Please note that this repo started recently as a fun weekend project: I took my earlier [nanoGPT](https://github.com/karpathy/nanoGPT), tuned it to implement the Llama-2 architecture instead of GPT-2, and the meat of it was writing the C inference engine in [run.c](run.c). So the project is young and moving quickly. Hat tip to the awesome [llama.cpp](https://github.com/ggerganov/llama.cpp) for inspiring this project. Compred to llama.cpp, I wanted something super simple, minimal, and educational so I chose to hard-code the Llama 2 architecture and just roll one inference file of pure C with no dependencies. ## feel the magic @@ -56,7 +58,9 @@ You can also prompt the model with a prefix or a number of additional command li > One day, Lily met a Shoggoth. He was very shy, but was also very generous. Lily said “Hello Shoggy! Can I be your friend?” Shoggy was happy to have a friend and said “Yes, let’s explore the universe together!” So they set off on a journey to explore the universe. As they travelled, Shoggy was happy to explain to Lily about all the wonderful things in the universe. At the end of the day, Lily and Shoggy had gathered lots of wonderful things from the universe, and they both felt very proud. They promised to explore the universe as one big pair and to never stop being generous to each other. -There is also an even better 110M param model available, see [models](#models). Quick note on sampling, the recommendation for good results is to use `-t 1.0 -p 0.9`, i.e. top-p sampling at 0.9 with temperature 1.0 (this is the default). To control the diversity of samples use either the temperature (i.e. vary `-t` between 0 and 1 and keep top-p off with `-p 0`) or the top-p value (i.e. vary `-p` between 0 and 1 and keep `-t 1`), but not both. Nice explainers on LLM sampling strategies include [this](https://peterchng.com/blog/2023/05/02/token-selection-strategies-top-k-top-p-and-temperature/), [this](https://docs.cohere.com/docs/controlling-generation-with-top-k-top-p) or [this](https://huggingface.co/blog/how-to-generate). +There is also an even better 110M param model available, see [models](#models). + +Quick note on sampling, the recommendation for ~best results is to sample with `-t 1.0 -p 0.9`, i.e. temperature 1.0 (default) but also top-p sampling at 0.9 (default). Intuitively, top-p ensures that tokens with tiny probabilities do not get sampled, so we can't get "unlucky" during sampling, and we are less likely to go "off the rails" afterwards. More generally, to control the diversity of samples use either the temperature (i.e. vary `-t` between 0 and 1 and keep top-p off with `-p 0`) or the top-p value (i.e. vary `-p` between 0 and 1 and keep `-t 1`), but not both. Nice explainers on LLM sampling strategies include [this](https://peterchng.com/blog/2023/05/02/token-selection-strategies-top-k-top-p-and-temperature/), [this](https://docs.cohere.com/docs/controlling-generation-with-top-k-top-p) or [this](https://huggingface.co/blog/how-to-generate). ## Meta's Llama 2 models @@ -83,11 +87,12 @@ base models... ¯\\_(ツ)_/¯. Since we can inference the base model, it should For the sake of examples of smaller, from-scratch models, I trained a small model series on TinyStories. All of these trained in a few hours on my training setup (4X A100 40GB GPUs). The 110M took around 24 hours. I am hosting them on huggingface hub [tinyllamas](https://huggingface.co/karpathy/tinyllamas), both in the original PyTorch .pt, and also in the llama2.c format .bin: -| model | dim | n_layers | n_heads | max context length | parameters | val loss | download -| --- | --- | --- | --- | --- | --- | --- | --- | -| OG | 288 | 6 | 6 | 256 | 15M | 1.072 | [stories15M.bin](https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin) | -| 42M| 512 | 8 | 8 | 1024 | 42M | 0.847 | [stories42M.bin](https://huggingface.co/karpathy/tinyllamas/resolve/main/stories42M.bin) | -| 110M| 768 | 12 | 12 | 1024 | 110M | 0.760 | [stories110M.bin](https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.bin) | +| model | dim | n_layers | n_heads | n_kv_heads | max context length | parameters | val loss | download +| --- | --- | --- | --- | --- | --- | --- | --- | --- | +| 260K | 64 | 5 | 8 | 4 | 512 | 260K | 1.297 | [stories260K](https://huggingface.co/karpathy/tinyllamas/tree/main/stories260K) +| OG | 288 | 6 | 6 | 6 | 256 | 15M | 1.072 | [stories15M.bin](https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin) | +| 42M| 512 | 8 | 8 | 8 | 1024 | 42M | 0.847 | [stories42M.bin](https://huggingface.co/karpathy/tinyllamas/resolve/main/stories42M.bin) | +| 110M| 768 | 12 | 12 | 12 | 1024 | 110M | 0.760 | [stories110M.bin](https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.bin) | You'll notice that the 110M model is equivalent to GPT-1 in size. Alternatively, this is also the smallest model in the GPT-2 series (`GPT-2 small`), except the max context length is only 1024 instead of 2048. The only notable changes from GPT-1/2 architecture is that Llama uses RoPE relatively positional embeddings instead of absolute/learned positional embeddings, a bit more fancy SwiGLU non-linearity in the MLP, RMSNorm instead of LayerNorm, bias=False on all Linear layers, and is optionally multiquery (but this is not yet supported in llama2.c). @@ -130,15 +135,53 @@ Watch the tokens stream by, fun! We can also run the PyTorch inference script fo ```bash wget https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.pt -P out15M -mv out15M/stories15M.pt out15M/ckpt.pt # sorry the sample script current assumes this directory structure / filename... -python sample.py --out_dir=out15M +python sample.py --checkpoint=out15M/stories15M.pt ``` -Which gives the same results. More detailed testing will be done in `test_all.py`. Currently you will need two files to test or sample: both the .bin file, and the .ckpt file inside a directory (see `test_all.py` for details). Sorry this is a bit janky right now, I have to think through running the tests without having to download 200MB of data. But run the tests with pytest: +Which gives the same results. + +## custom tokenizers + +In everything above, we've assumed the custom Lllama 2 tokenizer with 32,000 tokens. However, in many boutique LLMs, using vocabulary this big might be an overkill. If you have a small application you have in mind, you might be much better off training your own tokenizers. This can make everything nicer - with smaller vocabs your model has fewer parameters (because the token embedding table is a lot smaller), the inference is faster (because there are fewer tokens to predict), and your average sequence length per example could also get smaller (because the compression is a lot more efficient on your data). So let's see how we train a custom tokenizer. + +By default, to pretokenize the tinystories dataset we had to run, in order: -```bash -$ pytest ``` +python tinystories.py download +python tinystories.py pretokenize +``` + +The `pretokenize` stage here loads the Llama 2 tokenizer (vocab size 32,000) and uses it to convert the downloaded text into integers, and saves that to file. We now change this as follows, to train an example 4096-token tokenizer: + +``` +python tinystories.py download +python tinystories.py train_vocab --vocab_size=4096 +python tinystories.py pretokenize --vocab_size=4096 +``` + +The `train_vocab` stage will call the `train_vocab.sh` script, which calls the `sentencepiece` library to train the tokenizer, storing it in a new file `data/tok4096.model`. I tried to reproduce as well as I could the settings that (I think) Meta used to train their vocabulary. This uses the Byte Pair Encoding algorithm that starts out with raw utf8 byte sequences of the text data and then iteratively merges the most common consecutive pairs of tokens to form the vocabulary. Inspect the `tinystories.py` file - the custom tokenizers are stored in a special directory structure indexed by the vocab size. + +A quick note of interest is that vocab size of 4096 trained specifically on tinystories creates integer sequences with about the same sequence length per example as the default Llama 2 tokenizer of 32000 tokens! This means that our custom, tailored tokenizer is a lot better adapted to our specific text, and can compress it very effectively. So our trained models are smaller and faster. + +Now that we have pretokenized the dataset with our custom tokenizer, we can train the model. The training script `train.py` doesn't care about the exact tokens, it only cares about the vocabulary size so it can correctly initialize the model. So when training your model, make sure to pass in + +``` +python train.py --vocab_source=custom --vocab_size=4096 +``` + +(The defaults are `llama2` and `32000` respectively, which indicates the default Llama 2 tokenizer). This trains the model. Finally we are ready to run inference with our `run.c` script. For that we need two things. Number one, we have to export our tokenizer in the `.bin` format, do that with: + +``` +python tokenizer.py --tokenizer-model=data/tok4096.model +``` + +This writes the tokenizer to `data/tok4096.bin`. Now we can run inference, pointing it to this tokenizer using the `-z` flag: + +``` +./run out/model.bin -z data/tok4096.bin +``` + +This should print the samples. If you leave out the `-z` flag, it will use the default Llama 2 tokenizer, which would generate a good sequence of integers, but they would get translated using a different vocabulary to text, so it would look like gibberish. ## performance @@ -184,6 +227,17 @@ On **Centos 7**, **Amazon Linux 2018** use `rungnu` Makefile target: `make rungn On **Mac**, use clang from brew for openmp build. Install clang as `brew install llvm` and use the installed clang binary to compile with openmp: `make runomp CC=/opt/homebrew/opt/llvm/bin/clang` +## tests + +You can run tests simply with pytest: + +```bash +$ pip install pytest +$ pytest +``` + +This will currently invoke two tests inside `test_all.py`, which forward the model in both C and Python for 200 steps and check the output against a known good expected output. The tests currently run in only a few seconds, but will have to download and cache the stories260K models in a temporary `test` directory (only ~2MB download). + ## ack I trained the llama2.c storyteller models on a 4X A100 40GB box graciously provided by the excellent [Lambda labs](https://lambdalabs.com/service/gpu-cloud), thank you. @@ -216,6 +270,7 @@ If your candidate PRs have elements of these it doesn't mean they won't get merg - [llama2.rs](https://github.com/gaxler/llama2.rs) by @[gaxler](https://github.com/gaxler): a Rust port of this project - [llama2.rs](https://github.com/leo-du/llama2.rs) by @[leo-du](https://github.com/leo-du): A Rust port of this project - [llama2-rs](https://github.com/danielgrittner/llama2-rs) by @[danielgrittner](https://github.com/danielgrittner): a Rust port of this project + - [llama2.rs](https://github.com/lintian06/llama2.rs) by @[lintian06](https://github.com/lintian06): A Rust port of this project - Go - [go-llama2](https://github.com/tmc/go-llama2) by @[tmc](https://github.com/tmc): a Go port of this project - [llama2.go](https://github.com/nikolaydubina/llama2.go) by @[nikolaydubina](https://github.com/nikolaydubina): a Go port of this project @@ -228,6 +283,7 @@ If your candidate PRs have elements of these it doesn't mean they won't get merg - [llama2.cpp](https://github.com/leloykun/llama2.cpp) by @[leloykun](https://github.com/leloykun): a C++ port of this project - JavaScript - [llama2.js](https://github.com/epicure/llama2.js) by @[epicure](https://github.com/epicure): a JavaScript port of this project + - [llama2.ts](https://github.com/wizzard0/llama2.ts) by @[oleksandr_now](https://twitter.com/oleksandr_now): a TypeScript port of this project. Full Llama2-7B capable. - [llama2.c-emscripten](https://github.com/gohai/llama2.c-emscripten) by @[gohai](https://github.com/gohai): Emscripten (JavaScript) port, based on @ggerganov's initial prototype - Zig - [llama2.zig](https://github.com/cgbur/llama2.zig) by @[cgbur](https://github.com/cgbur): A Zig port of this project @@ -245,16 +301,16 @@ If your candidate PRs have elements of these it doesn't mean they won't get merg - [llama2.py](https://github.com/tairov/llama2.py) by @[tairov](https://github.com/tairov): a simple one file pure Python port of this project with zero dependencies - C# - [llama2.cs](https://github.com/trrahul/llama2.cs) by @[trrahul](https://github.com/trrahul): a C# port of this project +- WebAssembly + - [icpp-llm](https://github.com/icppWorld/icpp-llm): LLMs for the Internet Computer - [llama2.c - Llama 2 Everywhere](https://github.com/trholding/llama2.c) by @[trholding](https://github.com/trholding): Standalone, Bootable & Portable Binary Llama 2 ## unsorted todos -- add multiquery support into run.c -- add custom bpe training code and the ability to train a smaller vocabulary (32K is to much) +- make it easier to add a new dataset with not too much pain - should calculate freq_cis online in the script run.c instead of loading them - int4/8 quantization - export the model in a more sensible output format with a proper header, etc. -- train a tiny Llama test model (committed to repo) and use it as reference in unit tests - support Llama 2 7B Chat models and tune run.c to Chat UI/UX - llama2.cu investigate and merge - (LoRA) finetuning and export of Llama 2 models diff --git a/model.py b/model.py index f7edbb6..c8c82a9 100644 --- a/model.py +++ b/model.py @@ -11,12 +11,13 @@ from torch import nn @dataclass class ModelArgs: + # default hyperparameters for the Llama 7B model dim: int = 4096 n_layers: int = 32 n_heads: int = 32 n_kv_heads: Optional[int] = None - vocab_size: int = -1 # defined later by tokenizer - multiple_of: int = 256 # make SwiGLU hidden layer size multiple of large power of 2 + vocab_size: int = 32000 + multiple_of: int = 256 # MLP hidden layer size will be multiple of norm_eps: float = 1e-5 max_seq_len: int = 2048 dropout: float = 0.0 @@ -93,6 +94,7 @@ class Attention(nn.Module): def __init__(self, args: ModelArgs): super().__init__() self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads + assert args.n_heads % self.n_kv_heads == 0 model_parallel_size = 1 self.n_local_heads = args.n_heads // model_parallel_size self.n_local_kv_heads = self.n_kv_heads // model_parallel_size diff --git a/requirements.txt b/requirements.txt index e3f97c4..7187a73 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,7 +2,6 @@ numpy==1.23.5 pytest==7.4.0 Requests==2.31.0 sentencepiece==0.1.99 -tiktoken==0.3.3 torch==2.0.1 tqdm==4.64.1 wandb==0.15.5 diff --git a/run.c b/run.c index 9f4a1b2..df95e6f 100644 --- a/run.c +++ b/run.c @@ -39,11 +39,11 @@ typedef struct { // weights for rmsnorms float* rms_att_weight; // (layer, dim) rmsnorm weights float* rms_ffn_weight; // (layer, dim) - // weights for matmuls - float* wq; // (layer, dim, dim) - float* wk; // (layer, dim, dim) - float* wv; // (layer, dim, dim) - float* wo; // (layer, dim, dim) + // weights for matmuls. note dim == n_heads * head_size + float* wq; // (layer, dim, n_heads * head_size) + float* wk; // (layer, dim, n_kv_heads * head_size) + float* wv; // (layer, dim, n_kv_heads * head_size) + float* wo; // (layer, n_heads * head_size, dim) // weights for ffn float* w1; // (layer, hidden_dim, dim) float* w2; // (layer, dim, hidden_dim) @@ -82,6 +82,7 @@ typedef struct { void malloc_run_state(RunState* s, Config* p) { // we calloc instead of malloc to keep valgrind happy + int kv_dim = (p->dim * p->n_kv_heads) / p->n_heads; s->x = calloc(p->dim, sizeof(float)); s->xb = calloc(p->dim, sizeof(float)); s->xb2 = calloc(p->dim, sizeof(float)); @@ -93,8 +94,8 @@ void malloc_run_state(RunState* s, Config* p) { s->att = calloc(p->n_heads * p->seq_len, sizeof(float)); s->logits = calloc(p->vocab_size, sizeof(float)); s->probindex = calloc(p->vocab_size, sizeof(ProbIndex)); - s->key_cache = calloc(p->n_layers * p->seq_len * p->dim, sizeof(float)); - s->value_cache = calloc(p->n_layers * p->seq_len * p->dim, sizeof(float)); + s->key_cache = calloc(p->n_layers * p->seq_len * kv_dim, sizeof(float)); + s->value_cache = calloc(p->n_layers * p->seq_len * kv_dim, sizeof(float)); // ensure all mallocs went fine if (!s->x || !s->xb || !s->xb2 || !s->hb || !s->hb2 || !s->q || !s->k || !s->v || !s->att || !s->logits || !s->key_cache @@ -124,19 +125,20 @@ void free_run_state(RunState* s) { // initialization: read from checkpoint void checkpoint_init_weights(TransformerWeights *w, Config* p, float* f, int shared_weights) { + int head_size = p->dim / p->n_heads; float* ptr = f; w->token_embedding_table = ptr; ptr += p->vocab_size * p->dim; w->rms_att_weight = ptr; ptr += p->n_layers * p->dim; w->wq = ptr; - ptr += p->n_layers * p->dim * p->dim; + ptr += p->n_layers * p->dim * (p->n_heads * head_size); w->wk = ptr; - ptr += p->n_layers * p->dim * p->dim; + ptr += p->n_layers * p->dim * (p->n_kv_heads * head_size); w->wv = ptr; - ptr += p->n_layers * p->dim * p->dim; + ptr += p->n_layers * p->dim * (p->n_kv_heads * head_size); w->wo = ptr; - ptr += p->n_layers * p->dim * p->dim; + ptr += p->n_layers * (p->n_heads * head_size) * p->dim; w->rms_ffn_weight = ptr; ptr += p->n_layers * p->dim; w->w1 = ptr; @@ -148,7 +150,6 @@ void checkpoint_init_weights(TransformerWeights *w, Config* p, float* f, int sha w->rms_final_weight = ptr; ptr += p->dim; w->freq_cis_real = ptr; - int head_size = p->dim / p->n_heads; ptr += p->seq_len * head_size / 2; w->freq_cis_imag = ptr; ptr += p->seq_len * head_size / 2; @@ -158,12 +159,6 @@ void checkpoint_init_weights(TransformerWeights *w, Config* p, float* f, int sha // ---------------------------------------------------------------------------- // neural net blocks -void accum(float *a, float *b, int size) { - for (int i = 0; i < size; i++) { - a[i] += b[i]; - } -} - void rmsnorm(float* o, float* x, float* weight, int size) { // calculate sum of squares float ss = 0.0f; @@ -218,6 +213,8 @@ void transformer(int token, int pos, Config* p, RunState* s, TransformerWeights* // a few convenience variables float *x = s->x; int dim = p->dim; + int kv_dim = (p->dim * p->n_kv_heads) / p->n_heads; + int kv_mul = p->n_heads / p->n_kv_heads; // integer multiplier of the kv sharing in multiquery int hidden_dim = p->hidden_dim; int head_size = dim / p->n_heads; @@ -237,29 +234,33 @@ void transformer(int token, int pos, Config* p, RunState* s, TransformerWeights* // qkv matmuls for this position matmul(s->q, s->xb, w->wq + l*dim*dim, dim, dim); - matmul(s->k, s->xb, w->wk + l*dim*dim, dim, dim); - matmul(s->v, s->xb, w->wv + l*dim*dim, dim, dim); + matmul(s->k, s->xb, w->wk + l*dim*kv_dim, dim, kv_dim); + matmul(s->v, s->xb, w->wv + l*dim*kv_dim, dim, kv_dim); // RoPE relative positional encoding: complex-valued rotate q and k by freq_cis in each head for (int i = 0; i < dim; i+=2) { float q0 = s->q[i]; float q1 = s->q[i+1]; - float k0 = s->k[i]; - float k1 = s->k[i+1]; float fcr = freq_cis_real_row[(i % head_size) / 2]; float fci = freq_cis_imag_row[(i % head_size) / 2]; s->q[i] = q0 * fcr - q1 * fci; s->q[i+1] = q0 * fci + q1 * fcr; + } + for (int i = 0; i < kv_dim; i+=2) { + float k0 = s->k[i]; + float k1 = s->k[i+1]; + float fcr = freq_cis_real_row[(i % head_size) / 2]; + float fci = freq_cis_imag_row[(i % head_size) / 2]; s->k[i] = k0 * fcr - k1 * fci; s->k[i+1] = k0 * fci + k1 * fcr; } // save key,value at this time step (pos) to our kv cache - int loff = l * p->seq_len * dim; // kv cache layer offset for convenience - float* key_cache_row = s->key_cache + loff + pos * dim; - float* value_cache_row = s->value_cache + loff + pos * dim; - memcpy(key_cache_row, s->k, dim*sizeof(*key_cache_row)); - memcpy(value_cache_row, s->v, dim*sizeof(*value_cache_row)); + int loff = l * p->seq_len * kv_dim; // kv cache layer offset for convenience + float* key_cache_row = s->key_cache + loff + pos * kv_dim; + float* value_cache_row = s->value_cache + loff + pos * kv_dim; + memcpy(key_cache_row, s->k, kv_dim * sizeof(*key_cache_row)); + memcpy(value_cache_row, s->v, kv_dim * sizeof(*value_cache_row)); // multihead attention. iterate over all heads int h; @@ -272,7 +273,7 @@ void transformer(int token, int pos, Config* p, RunState* s, TransformerWeights* // iterate over all timesteps, including the current one for (int t = 0; t <= pos; t++) { // get the key vector for this head and at this timestep - float* k = s->key_cache + loff + t * dim + h * head_size; + float* k = s->key_cache + loff + t * kv_dim + (h / kv_mul) * head_size; // calculate the attention score as the dot product of q and k float score = 0.0f; for (int i = 0; i < head_size; i++) { @@ -291,7 +292,7 @@ void transformer(int token, int pos, Config* p, RunState* s, TransformerWeights* memset(xb, 0, head_size * sizeof(float)); for (int t = 0; t <= pos; t++) { // get the value vector for this head and at this timestep - float* v = s->value_cache + loff + t * dim + h * head_size; + float* v = s->value_cache + loff + t * kv_dim + (h / kv_mul) * head_size; // get the attention weight for this timestep float a = att[t]; // accumulate the weighted value into xb @@ -305,7 +306,9 @@ void transformer(int token, int pos, Config* p, RunState* s, TransformerWeights* matmul(s->xb2, s->xb, w->wo + l*dim*dim, dim, dim); // residual connection back into x - accum(x, s->xb2, dim); + for (int i = 0; i < dim; i++) { + x[i] += s->xb2[i]; + } // ffn rmsnorm rmsnorm(s->xb, x, w->rms_ffn_weight + l*dim, dim); @@ -329,7 +332,9 @@ void transformer(int token, int pos, Config* p, RunState* s, TransformerWeights* matmul(s->xb, s->hb, w->w2 + l*dim*hidden_dim, hidden_dim, dim); // residual connection - accum(x, s->xb, dim); + for (int i = 0; i < dim; i++) { + x[i] += s->xb[i]; + } } // final rmsnorm @@ -465,17 +470,24 @@ int sample_topp(float* probabilities, int n, float topp, ProbIndex* probindex) { // tokens that exceed probability topp. This way we never sample tokens that // have very low probabilities and are less likely to go "off the rails". + int n0 = 0; // quicksort indices in descending order of probabilities + // values smaller than (1 - topp) / (n - 1) cannot be part of the result + // so for efficiency we crop these out as candidates before sorting + const float cutoff = (1.0f - topp) / (n - 1); for (int i = 0; i < n; i++) { - probindex[i].index = i; - probindex[i].prob = probabilities[i]; + if (probabilities[i] >= cutoff) { + probindex[n0].index = i; + probindex[n0].prob = probabilities[i]; + n0++; + } } - qsort(probindex, n, sizeof(ProbIndex), compare); + qsort(probindex, n0, sizeof(ProbIndex), compare); // truncate the list where cumulative probability exceeds topp float cumulative_prob = 0.0f; - int last_idx = 0; - for (int i = 0; i < n; i++) { + int last_idx = n0 - 1; // in case of rounding errors consider all elements + for (int i = 0; i < n0; i++) { cumulative_prob += probindex[i].prob; if (cumulative_prob > topp) { last_idx = i; @@ -504,10 +516,11 @@ void error_usage() { fprintf(stderr, "Example: run model.bin -n 256 -i \"Once upon a time\"\n"); fprintf(stderr, "Options:\n"); fprintf(stderr, " -t temperature, default 1.0\n"); - fprintf(stderr, " -p p value in top-p (nucleus) sampling. default 0.9, 0 = off\n"); + fprintf(stderr, " -p p value in top-p (nucleus) sampling. default 0.9\n"); fprintf(stderr, " -s random seed, default time(NULL)\n"); fprintf(stderr, " -n number of steps to run for, default 256. 0 = max_seq_len\n"); fprintf(stderr, " -i input prompt\n"); + fprintf(stderr, " -z optional path to custom tokenizer\n"); exit(EXIT_FAILURE); } @@ -515,8 +528,9 @@ int main(int argc, char *argv[]) { // default inits char *checkpoint = NULL; // e.g. out/model.bin + char *tokenizer = "tokenizer.bin"; float temperature = 1.0f; // 0.0 = greedy deterministic. 1.0 = original. don't set higher - float topp = 0.9f; // top-p in nucleus sampling + float topp = 0.9f; // top-p in nucleus sampling. 1.0 = off. 0.9 works well, but slower rng_seed = 0; // seed rng with time by default int steps = 256; // number of steps to run for char *prompt = NULL; // prompt string @@ -534,6 +548,7 @@ int main(int argc, char *argv[]) { else if (argv[i][1] == 's') { rng_seed = atoi(argv[i + 1]); } else if (argv[i][1] == 'n') { steps = atoi(argv[i + 1]); } else if (argv[i][1] == 'i') { prompt = argv[i + 1]; } + else if (argv[i][1] == 'z') { tokenizer = argv[i + 1]; } else { error_usage(); } } if(rng_seed == 0) { rng_seed = (unsigned int)time(NULL);} @@ -567,13 +582,13 @@ int main(int argc, char *argv[]) { // right now we cannot run for more than config.seq_len steps if (steps <= 0 || steps > config.seq_len) { steps = config.seq_len; } - // read in the tokenizer.bin file + // read in the tokenizer .bin file char** vocab = (char**)malloc(config.vocab_size * sizeof(char*)); float* vocab_scores = (float*)malloc(config.vocab_size * sizeof(float)); unsigned int max_token_length; { - FILE *file = fopen("tokenizer.bin", "rb"); - if (!file) { fprintf(stderr, "couldn't load tokenizer.bin\n"); return 1; } + FILE *file = fopen(tokenizer, "rb"); + if (!file) { fprintf(stderr, "couldn't load %s\n", tokenizer); return 1; } if (fread(&max_token_length, sizeof(int), 1, file) != 1) { fprintf(stderr, "failed read\n"); return 1; } int len; for (int i = 0; i < config.vocab_size; i++) { @@ -623,7 +638,7 @@ int main(int argc, char *argv[]) { // apply softmax to the logits to get the probabilities for next token softmax(state.logits, config.vocab_size); // we sample from this distribution to get the next token - if (topp <= 0) { + if (topp <= 0 || topp >= 1) { // simply sample from the predicted probability distribution next = sample(state.logits, config.vocab_size); } else { diff --git a/sample.py b/sample.py index 040bc14..b26e277 100644 --- a/sample.py +++ b/sample.py @@ -5,17 +5,19 @@ import os import pickle from contextlib import nullcontext import torch -import tiktoken from model import ModelArgs, Transformer from tokenizer import Tokenizer +from tinystories import get_tokenizer_model_path + # ----------------------------------------------------------------------------- -out_dir = 'out' # ignored if init_from is not 'resume' +checkpoint = 'out/ckpt.pt' start = "" # or "<|endoftext|>" or etc. Can also specify a file, use as: "FILE:prompt.txt" num_samples = 1 # number of samples to draw max_new_tokens = 100 # number of tokens generated in each sample temperature = 1.0 # 1.0 = no change, < 1.0 = less random, > 1.0 = more random, in predictions top_k = 300 # retain only the top_k most likely tokens, clamp others to have 0 probability +tokenizer = "" # override the tokenizer model path seed = 1337 device = 'cuda' if torch.cuda.is_available() else 'cpu' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1', etc. #dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' # 'float32' or 'bfloat16' or 'float16' @@ -33,11 +35,10 @@ ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torc ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype) # init from a model saved in a specific directory -ckpt_path = os.path.join(out_dir, 'ckpt.pt') -checkpoint = torch.load(ckpt_path, map_location=device) -gptconf = ModelArgs(**checkpoint['model_args']) +checkpoint_dict = torch.load(checkpoint, map_location=device) +gptconf = ModelArgs(**checkpoint_dict['model_args']) model = Transformer(gptconf) -state_dict = checkpoint['model'] +state_dict = checkpoint_dict['model'] unwanted_prefix = '_orig_mod.' for k,v in list(state_dict.items()): if k.startswith(unwanted_prefix): @@ -50,8 +51,12 @@ if compile: print("Compiling the model...") model = torch.compile(model) # requires PyTorch 2.0 (optional) -# load the tokenizer -enc = Tokenizer() +# load the tokenizer, either provided, or attempt to find it +if tokenizer: + tokenizer_model = tokenizer +else: + tokenizer_model = get_tokenizer_model_path(vocab_size=gptconf.vocab_size) +enc = Tokenizer(tokenizer_model=tokenizer_model) # encode the beginning of the prompt if start.startswith('FILE:'): diff --git a/test_all.py b/test_all.py index 8563614..625af44 100644 --- a/test_all.py +++ b/test_all.py @@ -4,37 +4,67 @@ $ pytest """ import os import pytest # pip install pytest +import requests import subprocess + import torch from model import ModelArgs, Transformer +from tokenizer import Tokenizer -def test_argmax_inference(): - """ - Only the simplest test for now: run inference with temperature 0 - (for determinism) in both C and PyTorch, and see that the sampled tokens - are the same. - """ - test_ckpt_dir = "out" # TODO create a dummy test checkpoint for this? +# ----------------------------------------------------------------------------- +# test utilities - # run C version - model_path = os.path.join(test_ckpt_dir, "model.bin") - command = ["./run", model_path, "0.0"] - proc = subprocess.Popen(command, stdout=subprocess.PIPE) - c_tokens = [] - for line in proc.stdout: - token = int(line.decode('utf-8').strip()) - c_tokens.append(token) - proc.wait() - #print(c_tokens) +test_ckpt_dir = "test" - # run PyTorch version - device = "cuda" if torch.cuda.is_available() else "cpu" - ckpt_path = os.path.join(test_ckpt_dir, "ckpt.pt") - checkpoint = torch.load(ckpt_path, map_location=device) - gptconf = ModelArgs(**checkpoint['model_args']) +def download_file(url, filename): + print(f"Downloading {url} to {filename}") + response = requests.get(url, stream=True) + response.raise_for_status() # Raise an HTTPError on bad status code + with open(filename, 'wb') as file: + for chunk in response.iter_content(chunk_size=8192): + file.write(chunk) + +def attempt_download_files(): + os.makedirs(test_ckpt_dir, exist_ok=True) + root_url = "https://huggingface.co/karpathy/tinyllamas/resolve/main/stories260K" + need = ["stories260K.bin", "stories260K.pt", "tok512.bin", "tok512.model"] + for file in need: + url = os.path.join(root_url, file) + filename = os.path.join(test_ckpt_dir, file) + if not os.path.exists(filename): + download_file(url, filename) + +expected_stdout = b'Once upon a time, there was a little girl named Lily. She loved to play outside in the park. One day, she saw a big, red ball. She wanted to play with it, but it was too high.\nLily\'s mom said, "Lily, let\'s go to the park." Lily was sad and didn\'t know what to do. She said, "I want to play with your ball, but I can\'t find it."\nLily was sad and didn\'t know what to do. She said, "I\'m sorry, Lily. I didn\'t know what to do."\nLily didn\'t want to help her mom, so she' + +# ----------------------------------------------------------------------------- +# actual tests + +def test_runc(): + """ Forwards a model against a known-good desired outcome in run.c for 200 steps""" + attempt_download_files() + + model_path = os.path.join(test_ckpt_dir, "stories260K.bin") + tokenizer_path = os.path.join(test_ckpt_dir, "tok512.bin") + command = ["./run", model_path, "-z", tokenizer_path, "-t", "0.0", "-n", "200"] + proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + + stdout, stderr = proc.communicate() + + # strip the very last \n that is added by run.c for aesthetic reasons + stdout = stdout[:-1] + assert stdout == expected_stdout + +def test_python(): + """ Forwards a model against a known-good desired outcome in sample.py for 200 steps""" + attempt_download_files() + + device = "cpu" # stories260K is small enough to just breeze through it on CPU + checkpoint = os.path.join(test_ckpt_dir, "stories260K.pt") + checkpoint_dict = torch.load(checkpoint, map_location=device) + gptconf = ModelArgs(**checkpoint_dict['model_args']) model = Transformer(gptconf) - state_dict = checkpoint['model'] + state_dict = checkpoint_dict['model'] unwanted_prefix = '_orig_mod.' for k,v in list(state_dict.items()): if k.startswith(unwanted_prefix): @@ -44,10 +74,12 @@ def test_argmax_inference(): model.to(device) x = torch.tensor([[1]], dtype=torch.long, device=device) # 1 is BOS with torch.inference_mode(): - y = model.generate(x, max_new_tokens=gptconf.max_seq_len, temperature=0.0) + y = model.generate(x, max_new_tokens=200, temperature=0.0) pt_tokens = y[0].tolist() - pt_tokens = pt_tokens[1:] # remove BOS - #print(pt_tokens) - # compare - assert c_tokens == pt_tokens + tokenizer_model = os.path.join(test_ckpt_dir, "tok512.model") + enc = Tokenizer(tokenizer_model=tokenizer_model) + text = enc.decode(pt_tokens) + text = text.encode('ascii') # turn into bytes + + assert text == expected_stdout \ No newline at end of file diff --git a/tinyshakespeare.py b/tinyshakespeare.py deleted file mode 100644 index 602624c..0000000 --- a/tinyshakespeare.py +++ /dev/null @@ -1,140 +0,0 @@ -""" -Download, preprocess and serve the TinyShakespeare dataset as a DataLoader. - -Follows the same interface as the TinyStories dataset. -""" - -import argparse -import os -import random - -import numpy as np -import requests -import torch -import torch.distributed as dist -from tqdm import tqdm - -from tokenizer import Tokenizer - -DATA_CACHE_DIR = "data" - -def download_file(url: str, fname: str, chunk_size=1024): - """Helper function to download a file from a given url""" - resp = requests.get(url, stream=True) - total = int(resp.headers.get("content-length", 0)) - with open(fname, "wb") as file, tqdm( - desc=fname, - total=total, - unit="iB", - unit_scale=True, - unit_divisor=1024, - ) as bar: - for data in resp.iter_content(chunk_size=chunk_size): - size = file.write(data) - bar.update(size) - - -def download(): - """Downloads the dataset to disk.""" - os.makedirs(DATA_CACHE_DIR, exist_ok=True) - - # download the TinyShakespeare dataset, unless it's already downloaded - data_url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt" - data_filename = os.path.join(DATA_CACHE_DIR, "tinyshakespeare.txt") - if not os.path.exists(data_filename): - print(f"Downloading {data_url} to {data_filename}...") - download_file(data_url, data_filename) - else: - print(f"{data_filename} already exists, skipping download...") - - print("Download done.") - -def pretokenize(): - enc = Tokenizer() - - data_file = os.path.join(DATA_CACHE_DIR, "tinyshakespeare.txt") - - all_tokens = [] - with open(data_file, "r") as f: - for line in f: - text = line.strip() - tokens = enc.encode(text, bos=True, eos=False) - all_tokens.extend(tokens) - all_tokens = np.array(all_tokens, dtype=np.uint16) - print(f"Total tokens: {len(all_tokens)}") - with open(data_file.replace(".txt", ".bin"), "wb") as f: - f.write(all_tokens.tobytes()) - print(f"Saved {data_file.replace('.txt', '.bin')}") - print("Done.") - - -class PretokDataset(torch.utils.data.IterableDataset): - """Loads pretokenized examples from disk and yields them as PyTorch tensors.""" - - def __init__(self, split, max_seq_len): - super().__init__() - self.split = split - self.max_seq_len = max_seq_len - - def __iter__(self): - # get worker info within a DataLoader - worker_info = torch.utils.data.get_worker_info() - worker_id = worker_info.id if worker_info else 0 - # get DDP rank info - rank = dist.get_rank() if dist.is_initialized() else 0 - # combine the worker_id and worker_rank to create a unique seed for rng - seed = 42 + worker_id + 1337 * rank - rng = random.Random(seed) - print(f"Created a PretokDataset with rng seed {seed}") - data_file = os.path.join(DATA_CACHE_DIR, "tinyshakespeare.bin") - m_all = np.memmap(data_file, dtype=np.uint16, mode="r") - - # split out 10% of the data for validation - split_ix = int(len(m_all) * 0.9) - if self.split == "train": - m = m_all[:split_ix] - else: - m = m_all[split_ix:] - - num_batches = len(m) // self.max_seq_len - num_batches -= 1 # drop the last partial batch - assert num_batches > 0, "this split is way too small? investigate." - - while True: - ixs = list(range(num_batches)) - rng.shuffle(ixs) - for ix in ixs: - start = ix * self.max_seq_len - end = start + self.max_seq_len + 1 - # calling .astype will copy the data into a new numpy array, now in RAM - chunk = torch.from_numpy((m[start:end]).astype(np.int64)) - x = chunk[:-1] - y = chunk[1:] - yield x, y - - -class ShakespeareTask: - - @staticmethod - def iter_batches(split, batch_size, max_seq_len, device, num_workers=0): - ds = PretokDataset(split, max_seq_len) - dl = torch.utils.data.DataLoader( - ds, batch_size=batch_size, pin_memory=True, num_workers=num_workers - ) - for x, y in dl: - x = x.to(device, non_blocking=True) - y = y.to(device, non_blocking=True) - yield x, y - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("stage", type=str, choices=["download", "train_tokenizer", "pretokenize"]) - args = parser.parse_args() - - # depending on the stage call the appropriate function - fun = { - "download": download, - "pretokenize": pretokenize, - } - fun[args.stage]() \ No newline at end of file diff --git a/tinystories.py b/tinystories.py index 419e0d5..690cb02 100644 --- a/tinystories.py +++ b/tinystories.py @@ -9,6 +9,7 @@ import os import random from typing import List from concurrent.futures import ProcessPoolExecutor +from functools import partial import numpy as np import requests @@ -37,7 +38,7 @@ def download_file(url: str, fname: str, chunk_size=1024): def download(): - """Downloads the dataset to disk.""" + """Downloads the TinyStories dataset to DATA_CACHE_DIR""" os.makedirs(DATA_CACHE_DIR, exist_ok=True) # download the TinyStories dataset, unless it's already downloaded @@ -66,10 +67,61 @@ def download(): print(f"Number of shards: {len(shard_filenames)}") print(f"Example story:\n{data[0]}") +def train_vocab(vocab_size): + """ + Trains a custom sentencepiece tokenizer on the TinyStories dataset. + The custom tokenizer files will be saved in DATA_CACHE_DIR/tok{N} directories, + where N is the vocab size. This is also where the pretok .bin files will go. + """ + assert vocab_size > 0, "Vocab size must be positive" -def process_shard(args): + # output file prefix path for sentencepiece + prefix = os.path.join(DATA_CACHE_DIR, f"tok{vocab_size}") + + # how many shards we'll use for vocab training, kept low for efficiency + num_shards = 10 + + # 1) export a large chunk of text as a single text file tiny.txt + tiny_file = os.path.join(DATA_CACHE_DIR, "tiny.txt") + data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data") + shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.json"))) + + print(f"Writing temporary file {tiny_file} with {num_shards} shards...") + with open(tiny_file, "w") as of: + for shard in tqdm(shard_filenames[:num_shards]): + with open(shard, "r") as f: + data = json.load(f) + for example in data: + text = example["story"] + text = text.strip() + of.write(text + "\n") + print(f"Size is: {os.path.getsize(tiny_file) / 1024 / 1024:.2f} MB") + + # 2) run the train_vocab.sh script that trains the sentencepiece model + print("Will now train the vocab with:") + cmd = f"bash train_vocab.sh {tiny_file} {prefix} {vocab_size}" + print(cmd) + print("OK? [y/N] ") + dec = input() + if dec.lower() != "y": + print("Exiting...") + return + os.system(cmd) + + # 3) optional cleanup, ask the user if they'd like to delete tiny.txt + dec = input(f"Delete the temporary file {tiny_file}? [y/N] ") + if dec.lower() == "y": + os.remove(tiny_file) + print(f"Deleted {tiny_file}") + + print(f"Trained tokenizer is in {prefix}.model") + print("Done.") + + +def process_shard(args, vocab_size): shard_id, shard = args - enc = Tokenizer() + tokenizer_model = get_tokenizer_model_path(vocab_size) + enc = Tokenizer(tokenizer_model) with open(shard, "r") as f: data = json.load(f) all_tokens = [] @@ -80,31 +132,49 @@ def process_shard(args): all_tokens.extend(tokens) # convert to uint16 nparray all_tokens = np.array(all_tokens, dtype=np.uint16) - # write to disk - tokenized_filename = shard.replace(".json", ".bin") + # calculate the output filename + if vocab_size == 0: + # if we're using Llama 2, just save the tokenized file in the same dir + tokenized_filename = shard.replace(".json", ".bin") + else: + # save .bin files into a new tok{N} directory + bin_dir = os.path.join(DATA_CACHE_DIR, f"tok{vocab_size}") + shard_basename = os.path.basename(shard) + bin_basename = shard_basename.replace(".json", ".bin") + tokenized_filename = os.path.join(bin_dir, bin_basename) + # write the bytes with open(tokenized_filename, "wb") as f: f.write(all_tokens.tobytes()) - print(f"Saved {tokenized_filename}") + # calculate the average sequence length (they are separated by BOS=1) + avg_seq_len = all_tokens.size / ((all_tokens == 1).sum()) + print(f"Saved {tokenized_filename}, average seqlen: {avg_seq_len:.2f}") -def pretokenize(): +def pretokenize(vocab_size): # iterate the shards and tokenize all of them one by one data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data") shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.json"))) + if vocab_size > 0: + # .bin files will be saved into tok{N} directory, create it once here + bin_dir = os.path.join(DATA_CACHE_DIR, f"tok{vocab_size}") + os.makedirs(bin_dir, exist_ok=True) # process all the shards in a process pool + fun = partial(process_shard, vocab_size=vocab_size) with ProcessPoolExecutor() as executor: - executor.map(process_shard, enumerate(shard_filenames)) + executor.map(fun, enumerate(shard_filenames)) print("Done.") class PretokDataset(torch.utils.data.IterableDataset): """Loads pretokenized examples from disk and yields them as PyTorch tensors.""" - def __init__(self, split, max_seq_len): + def __init__(self, split, max_seq_len, vocab_size, vocab_source): super().__init__() self.split = split self.max_seq_len = max_seq_len + self.vocab_size = vocab_size + self.vocab_source = vocab_source def __iter__(self): # get worker info within a DataLoader @@ -116,8 +186,14 @@ class PretokDataset(torch.utils.data.IterableDataset): seed = 42 + worker_id + 1337 * rank rng = random.Random(seed) print(f"Created a PretokDataset with rng seed {seed}") - data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data") - shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.bin"))) + if self.vocab_source == "llama2": + # the .bin files are right along the .json files + bin_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data") + shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin"))) + elif self.vocab_source == "custom": + # the .bin files are in tok{N} directory + bin_dir = os.path.join(DATA_CACHE_DIR, f"tok{self.vocab_size}") + shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin"))) # train/test split. let's use only shard 0 for test split, rest train shard_filenames = shard_filenames[1:] if self.split == "train" else shard_filenames[:1] while True: @@ -139,12 +215,25 @@ class PretokDataset(torch.utils.data.IterableDataset): y = chunk[1:] yield x, y +# ----------------------------------------------------------------------------- +# public interface functions + +def get_tokenizer_model_path(vocab_size): + """ + Returns path to the sentencepiece tokenizer model for a given vocab size + vocab_size = 0 designates the default Llama 2 tokenizer, in that case + None is returned. + """ + if vocab_size == 0: + return None + else: + return os.path.join(DATA_CACHE_DIR, f"tok{vocab_size}.model") class Task: @staticmethod - def iter_batches(split, batch_size, max_seq_len, device, num_workers=0): - ds = PretokDataset(split, max_seq_len) + def iter_batches(batch_size, device, num_workers=0, **dataset_kwargs): + ds = PretokDataset(**dataset_kwargs) dl = torch.utils.data.DataLoader( ds, batch_size=batch_size, pin_memory=True, num_workers=num_workers ) @@ -153,16 +242,33 @@ class Task: y = y.to(device, non_blocking=True) yield x, y +# ----------------------------------------------------------------------------- +# CLI for constructing the dataset if __name__ == "__main__": + """ + These stages are designed to be run in order. + + To tokenize data with the Llama 2 tokenizer: + python tinystories.py download + python tinystories.py pretokenize + + To tokenize data with a custom tokenizer we train ourselves with sentencepiece, e.g.: + python tinystories.py download + python tinystories.py train_vocab --vocab_size=2048 + python tinystories.py pretokenize --vocab_size=2048 + """ parser = argparse.ArgumentParser() - parser.add_argument("stage", type=str, choices=["download", "train_tokenizer", "pretokenize"]) + parser.add_argument("stage", type=str, choices=["download", "pretokenize", "train_vocab"]) + parser.add_argument("--vocab_size", type=int, default=0, help="pretokenization vocab size. 0 = use Llama 2 tokenizer.") args = parser.parse_args() # depending on the stage call the appropriate function - fun = { - "download": download, - "pretokenize": pretokenize, - } - fun[args.stage]() - + if args.stage == "download": + download() + elif args.stage == "train_vocab": + train_vocab(vocab_size=args.vocab_size) + elif args.stage == "pretokenize": + pretokenize(vocab_size=args.vocab_size) + else: + raise ValueError(f"Unknown stage {args.stage}") diff --git a/tokenizer.py b/tokenizer.py index 35eee20..bc2a35a 100644 --- a/tokenizer.py +++ b/tokenizer.py @@ -4,20 +4,19 @@ import os import struct -from logging import getLogger +import argparse from typing import List from sentencepiece import SentencePieceProcessor TOKENIZER_MODEL = "tokenizer.model" # the llama sentencepiece tokenizer model -TOKENIZER_BIN = "tokenizer.bin" # binary version of the tokenizer for inference in C class Tokenizer: - def __init__(self): - model_path = TOKENIZER_MODEL + def __init__(self, tokenizer_model=None): + model_path = tokenizer_model if tokenizer_model else TOKENIZER_MODEL assert os.path.isfile(model_path), model_path self.sp_model = SentencePieceProcessor(model_file=model_path) - #print(f"Loaded SentencePiece model from {model_path}") + self.model_path = model_path # BOS / EOS token IDs self.n_words: int = self.sp_model.vocab_size() @@ -59,17 +58,23 @@ class Tokenizer: tokens.append(b) scores.append(s) - + # record the max token length max_token_length = max(len(t) for t in tokens) # write to a binary file - with open(TOKENIZER_BIN, 'wb') as f: + # the tokenizer.bin file is the same as .model file, but .bin + tokenizer_bin = self.model_path.replace('.model', '.bin') + with open(tokenizer_bin, 'wb') as f: f.write(struct.pack("I", max_token_length)) for bytes, score in zip(tokens, scores): f.write(struct.pack("fI", score, len(bytes))) f.write(bytes) if __name__ == "__main__": - t = Tokenizer() + parser = argparse.ArgumentParser() + parser.add_argument("-t", "--tokenizer-model", type=str, help="optional path to custom tokenizer ") + args = parser.parse_args() + + t = Tokenizer(args.tokenizer_model) t.export() diff --git a/train.py b/train.py index dbf0b24..b1972dc 100644 --- a/train.py +++ b/train.py @@ -29,7 +29,6 @@ from torch.distributed import destroy_process_group, init_process_group from torch.nn.parallel import DistributedDataParallel as DDP from tinystories import Task -from tinyshakespeare import ShakespeareTask # ----------------------------------------------------------------------------- # I/O @@ -47,11 +46,13 @@ wandb_run_name = "run" + datetime.now().strftime("%Y_%m_%d_%H_%M_%S") # data batch_size = 128 # if gradient_accumulation_steps > 1, this is the micro-batch size max_seq_len = 256 -dataset = "tinystories" # tinystories|tinyshakespeare +vocab_source = "llama2" # llama2|custom; use Lllama 2 vocab from Meta, or custom trained +vocab_size = 32000 # the Llama 2 tokenizer has 32K tokens # model dim = 288 n_layers = 6 n_heads = 6 +n_kv_heads = 6 multiple_of = 32 dropout = 0.0 # adamw optimizer @@ -83,6 +84,10 @@ config = {k: globals()[k] for k in config_keys} # will be useful for logging lr_decay_iters = max_iters # should be ~= max_iters per Chinchilla min_lr = 0.0 # minimum learning rate, should be ~= learning_rate/10 per Chinchilla +# validating checks +assert vocab_source in ["llama2", "custom"] +assert vocab_source == "custom" or vocab_size == 32000, "The vocab from Meta has 32K tokens" + # various inits, derived attributes, I/O setup ddp = int(os.environ.get("RANK", -1)) != -1 # is this a ddp run? if ddp: @@ -123,11 +128,12 @@ ctx = ( ) # task-specific setup -task = {'tinystories': Task, 'tinyshakespeare': ShakespeareTask}[dataset] iter_batches = partial( - task.iter_batches, + Task.iter_batches, batch_size=batch_size, max_seq_len=max_seq_len, + vocab_size=vocab_size, + vocab_source=vocab_source, device=device, num_workers=0, ) @@ -141,8 +147,8 @@ model_args = dict( dim=dim, n_layers=n_layers, n_heads=n_heads, - n_kv_heads=n_heads, - vocab_size=32000, + n_kv_heads=n_kv_heads, + vocab_size=vocab_size, multiple_of=multiple_of, max_seq_len=max_seq_len, dropout=dropout, @@ -206,7 +212,7 @@ def estimate_loss(): out = {} model.eval() for split in ["train", "val"]: - batch_iter = iter_batches(split) + batch_iter = iter_batches(split=split) losses = torch.zeros(eval_iters) # keep on CPU for k in range(eval_iters): X, Y = next(batch_iter) @@ -238,7 +244,7 @@ if wandb_log and master_process: wandb.init(project=wandb_project, name=wandb_run_name, config=config) # training loop -train_batch_iter = iter_batches("train") +train_batch_iter = iter_batches(split="train") X, Y = next(train_batch_iter) # fetch the very first batch t0 = time.time() local_iter_num = 0 # number of iterations in the lifetime of this process diff --git a/train_vocab.sh b/train_vocab.sh new file mode 100755 index 0000000..7803af8 --- /dev/null +++ b/train_vocab.sh @@ -0,0 +1,126 @@ +#!/bin/bash + +# Trains a sentencepiece tokenizer model on a bunch of given data, my best +# effort attempt to replicate how Meta trained their Llama 2 tokenizer. + +# usage: $ train_vocab.sh +# example: +# ./train_vocab.sh tiny.txt tokenizer_tiny 1024 +# requirements: +# install https://github.com/google/sentencepiece + +# check if the correct number of arguments are provided +if [ $# -ne 3 ]; then + echo "Usage: $0 " + exit 1 +fi + +# assign command-line arguments to variables +input=$1 +model_prefix=$2 +vocab_size=$3 + +# check if input file exists +if [ ! -f "$input" ]; then + echo "Usage: $0 " + echo "input '$input' not found." + exit 1 +fi + +# check if vocab_size is a positive integer +if ! [[ "$vocab_size" =~ ^[0-9]+$ ]] || [ "$vocab_size" -lt 1 ]; then + echo "Usage: $0 " + echo "vocab_size size must be a positive integer." + exit 1 +fi + +# Print the processed inputs +echo "Input: $input" +echo "Model Prefix: $model_prefix" +echo "Vocabulary Size: $vocab_size" + +# train a sentencepiece tokenizer model +# Llama 2 config can be printed as follows: + +# import sentencepiece.sentencepiece_model_pb2 +# mp = sentencepiece.sentencepiece_model_pb2.ModelProto() +# mp.ParseFromString(open("tokenizer.model", "rb").read()) +# print(mp.trainer_spec) +# print(mp.normalizer_spec) + +# this gives: + +# trainer_spec { +# input: "/large_experiments/theorem/datasets/MERGED/all.test1.merged" +# model_prefix: "spm_model_32k_200M_charcov099995_allowWSO__v2" +# model_type: BPE +# vocab_size: 32000 +# self_test_sample_size: 0 +# input_format: "text" +# character_coverage: 0.9999499917030334 +# input_sentence_size: 200000000 +# seed_sentencepiece_size: 1000000 +# shrinking_factor: 0.75 +# num_threads: 80 +# num_sub_iterations: 2 +# max_sentence_length: 4192 +# shuffle_input_sentence: true +# max_sentencepiece_length: 16 +# split_by_unicode_script: true +# split_by_whitespace: true +# split_by_number: true +# treat_whitespace_as_suffix: false +# split_digits: true +# allow_whitespace_only_pieces: true +# vocabulary_output_piece_score: true +# hard_vocab_limit: true +# use_all_vocab: false +# byte_fallback: true +# required_chars: "" +# unk_id: 0 +# bos_id: 1 +# eos_id: 2 +# pad_id: -1 +# unk_surface: " \342\201\207 " +# unk_piece: "" +# bos_piece: "" +# eos_piece: "" +# pad_piece: "" +# train_extremely_large_corpus: false +# enable_differential_privacy: false +# differential_privacy_noise_level: 0.0 +# differential_privacy_clipping_threshold: 0 +# } +# normalizer_spec { +# name: "identity" +# precompiled_charsmap: "" +# add_dummy_prefix: true +# remove_extra_whitespaces: false +# normalization_rule_tsv: "" +# } + +# let's now use spm_train to train this exact model +# options docs: https://github.com/google/sentencepiece/blob/master/doc/options.md + +# we'll depart on a few settings: +# character_coverage -> 1.0 + +# other important notes: +# --split-digits = true, per the paper +# --allow_whitespace_only_pieces is true, default in spm is false +# --byte_fallback is true, default in spm is false +# --normalization_rule_name is identity, default in spm is nmt_nfkc + +spm_train --input="$input" \ + --model_prefix="$model_prefix" \ + --model_type=bpe \ + --vocab_size="$vocab_size" \ + --self_test_sample_size=0 \ + --input_format="text" \ + --character_coverage=1.0 \ + --num_threads="$(nproc)" \ + --split_digits=true \ + --allow_whitespace_only_pieces=true \ + --byte_fallback=true \ + --unk_surface=" \342\201\207 " \ + --normalization_rule_name=identity \