diff --git a/.github/PULL_REQUEST_SUMMARY.md b/.github/PULL_REQUEST_SUMMARY.md new file mode 100644 index 0000000..f170dce --- /dev/null +++ b/.github/PULL_REQUEST_SUMMARY.md @@ -0,0 +1,3 @@ +# FloppyLLM implementation + +This branch adds a self-contained 32-bit Windows TinyStories CLI under the formatted 1.44 MB floppy limit, plus reproducible quantization, size enforcement, configurable token/context defaults, Windows smoke tests, and tagged releases. diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml new file mode 100644 index 0000000..d96a0a8 --- /dev/null +++ b/.github/workflows/build.yml @@ -0,0 +1,99 @@ +name: Build floppy executable + +on: + push: + branches: [main] + tags: ["v*"] + pull_request: + workflow_dispatch: + inputs: + default_tokens: + description: Default maximum new tokens + required: true + default: "128" + default_context: + description: Default context allocation (0 uses model's 256) + required: true + default: "0" + +permissions: + contents: write + +env: + LLVM_MINGW_VERSION: "20260616" + LLVM_MINGW_SHA256: "a1f7968b48ba8d949194d6dee6c76f3cd0f61cba91658599af2c2c834a55ab87" + DEFAULT_TOKENS: ${{ inputs.default_tokens || '128' }} + DEFAULT_CONTEXT: ${{ inputs.default_context || '0' }} + +jobs: + build: + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install converter dependency + run: python -m pip install numpy + + - name: Build verified Q8 model + run: make model PYTHON=python + + - name: Install pinned Windows cross-compiler + run: | + url="https://github.com/mstorsjo/llvm-mingw/releases/download/${LLVM_MINGW_VERSION}/llvm-mingw-${LLVM_MINGW_VERSION}-msvcrt-ubuntu-22.04-x86_64.tar.xz" + curl -L --fail -o /tmp/llvm-mingw.tar.xz "$url" + echo "${LLVM_MINGW_SHA256} /tmp/llvm-mingw.tar.xz" | sha256sum -c - + mkdir /tmp/llvm-mingw + tar -xJf /tmp/llvm-mingw.tar.xz -C /tmp/llvm-mingw --strip-components=1 + + - name: Build single-file Windows CLI + run: | + make windows \ + PYTHON=python \ + WINDOWS_CC=/tmp/llvm-mingw/bin/i686-w64-mingw32-clang \ + DEFAULT_TOKENS="${DEFAULT_TOKENS}" \ + DEFAULT_CONTEXT="${DEFAULT_CONTEXT}" + + - name: Record checksum + run: sha256sum dist/floppy-llm.exe | tee dist/floppy-llm.exe.sha256 + + - uses: actions/upload-artifact@v4 + with: + name: floppy-llm-windows + path: | + dist/floppy-llm.exe + dist/floppy-llm.exe.sha256 + + - name: Publish tagged release + if: startsWith(github.ref, 'refs/tags/v') + env: + GH_TOKEN: ${{ github.token }} + run: | + gh release create "$GITHUB_REF_NAME" \ + dist/floppy-llm.exe \ + dist/floppy-llm.exe.sha256 \ + --title "FloppyLLM $GITHUB_REF_NAME" \ + --generate-notes + + smoke-test: + needs: build + runs-on: windows-latest + steps: + - uses: actions/download-artifact@v4 + with: + name: floppy-llm-windows + path: dist + + - name: Run Windows executable + shell: pwsh + run: | + .\dist\floppy-llm.exe --version + $text = .\dist\floppy-llm.exe --prompt "Once upon a time" --temperature 0 --tokens 12 + if ($LASTEXITCODE -ne 0 -or [string]::IsNullOrWhiteSpace($text)) { + throw "generation smoke test failed" + } + Write-Output $text + diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..8772ee4 --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +dist/ +model/*.bin +*.pyc +__pycache__/ + diff --git a/BUILD_RESULT.md b/BUILD_RESULT.md new file mode 100644 index 0000000..45269c3 --- /dev/null +++ b/BUILD_RESULT.md @@ -0,0 +1,11 @@ +# Verified local build + +- Windows artifact: `floppy-llm.exe` +- Format: PE32 console executable, Intel 80386 +- Model: tiny1m, 896,256 parameters, Q8_0 group 32 +- Size: 1,148,267 bytes +- Floppy capacity: 1,474,560 bytes +- Remaining: 326,293 bytes +- SHA-256 of the locally verified build: `d9401ff6f14fb92e0691d3f7526dbd86fd226c9718e7a53075445d15aace72e4` + +GitHub Actions reproduces the model from pinned upstream hashes, builds the executable, enforces the size ceiling, and runs it on a Windows runner. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..be6eaa2 --- /dev/null +++ b/LICENSE @@ -0,0 +1,23 @@ +MIT License + +Copyright (c) 2026 FloppyLLM contributors +Copyright (c) 2023 Andrej + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..9af5168 --- /dev/null +++ b/Makefile @@ -0,0 +1,54 @@ +PYTHON ?= python3 +CC ?= cc +WINDOWS_CC ?= i686-w64-mingw32-clang +DEFAULT_TOKENS ?= 128 +DEFAULT_CONTEXT ?= 0 +FLOPPY_BYTES ?= 1474560 + +CFLAGS = -Os -ffunction-sections -fdata-sections -Wall -Wextra +DEFINES = -DDEFAULT_MAX_NEW_TOKENS=$(DEFAULT_TOKENS) -DDEFAULT_CONTEXT_OVERRIDE=$(DEFAULT_CONTEXT) +WINDOWS_LDFLAGS = -Wl,--gc-sections -Wl,-s -Wl,--subsystem,console:4.0 \ + -Wl,--major-os-version,4 -Wl,--minor-os-version,0 \ + -Wl,--disable-dynamicbase -Wl,--disable-nxcompat + +.PHONY: all model native windows check clean + +all: native + +model: model/tiny1m-q8.bin model/tokenizer.bin + +model/model-f32.bin: + mkdir -p model + curl -L --fail -o $@ https://huggingface.co/shibatch/tiny1m/resolve/main/model.bin + printf '%s %s\n' 80ba15bd7feed5f78c13c33bac83992617ff030661f16abbfed8fba0a6e335c1 $@ | sha256sum -c - + +model/tokenizer.bin: + mkdir -p model + curl -L --fail -o $@ https://huggingface.co/shibatch/tiny1m/resolve/main/tokenizer.bin + printf '%s %s\n' 037cb335abb25d1fa9e8ecae30ed2a3a8ace9302862ebcdc05d51a6bbb10c312 $@ | sha256sum -c - + +model/tiny1m-q8.bin: model/model-f32.bin tools/legacy_to_q8.py + $(PYTHON) tools/legacy_to_q8.py $< $@ 32 + printf '%s %s\n' 43c16d0d0702c1f628f841cad10c01bfde4a3a94d7784221a2f6208853ce12c2 $@ | sha256sum -c - + +native: model + mkdir -p dist + $(CC) $(CFLAGS) $(DEFINES) -o dist/floppy-llm-runner src/floppy_llm.c -lm + $(PYTHON) tools/pack_overlay.py dist/floppy-llm-runner model/tiny1m-q8.bin model/tokenizer.bin dist/floppy-llm + chmod +x dist/floppy-llm + $(PYTHON) tools/check_size.py dist/floppy-llm $(FLOPPY_BYTES) + +windows: model + mkdir -p dist + $(WINDOWS_CC) $(CFLAGS) $(DEFINES) $(WINDOWS_LDFLAGS) \ + -o dist/floppy-llm-runner.exe src/floppy_llm.c -lm + $(PYTHON) tools/pack_overlay.py dist/floppy-llm-runner.exe model/tiny1m-q8.bin model/tokenizer.bin dist/floppy-llm.exe + $(PYTHON) tools/check_size.py dist/floppy-llm.exe $(FLOPPY_BYTES) + +check: native + ./dist/floppy-llm --version + ./dist/floppy-llm --prompt "Once upon a time, there was a little robot named Pip." --temperature 0 --tokens 20 + +clean: + $(RM) -r dist + diff --git a/README.md b/README.md index 45105bf..7e4db83 100644 --- a/README.md +++ b/README.md @@ -1 +1,110 @@ -# floppy-llm \ No newline at end of file +# FloppyLLM + +A real 896K-parameter transformer, tokenizer, and native Windows CLI packed +into one executable that fits on a 1.44 MB floppy disk. + +The release file is `floppy-llm.exe`. It does not need Python, a model file, +or an installer beside it. + +## Run it + +On Windows: + +```bat +floppy-llm.exe +``` + +That starts a small interactive story-completion prompt. For one-shot use: + +```bat +floppy-llm.exe -i "Once upon a time, there was a tiny computer" -n 80 +``` + +Useful options: + +```text +-i, --prompt TEXT complete TEXT and exit +-n, --tokens N maximum new tokens; 0 fills the context +-c, --context N runtime context allocation; trained default is 256 +-t, --temperature N 0 is deterministic; default is 0.8 +-p, --top-p N default is 0.9 +-s, --seed N sampling seed + --stats print speed +``` + +`--context` controls the runtime token/cache limit and does not require a +rebuild. A value beyond the trained 256-token context is allowed for +experimentation, but output quality may get worse. + +## What fits + +The checked Windows build is 1,148,267 bytes: + +| Part | Bytes | +| --- | ---: | +| Q8/group-32 model | 1,055,488 | +| tokenizer | 6,227 | +| 32-bit Windows runner | 86,528 | +| overlay footer | 24 | +| **total** | **1,148,267** | +| 1.44 MB formatted floppy capacity | 1,474,560 | +| **free** | **326,293** | + +The executable reads the model and tokenizer appended to itself. At runtime it +needs several megabytes of RAM for model data, activations, and the KV cache; +only the disk footprint is floppy-sized. + +## Why this model + +The build uses [`shibatch/tiny1m`](https://huggingface.co/shibatch/tiny1m), +trained on TinyStories. It has 896,256 parameters and is substantially more +capable than the commonly demonstrated 260K-parameter `stories260K` model, +while its llama2.c-compatible tokenizer and architecture let the whole +Q8/group-32 system remain comfortably below 1.44 MB. In a deterministic +sanity test, the Q8/group-32 build matched the float checkpoint's continuation. + +This is still a deliberately terrible, extremely narrow model. Expect short, +simple children's-story prose—not factual answers or useful coding help. + +## Build + +The repository intentionally does not commit generated model binaries. +`make model` downloads the upstream MIT model, verifies hashes, and produces +the exact Q8 checkpoint. + +Native Linux reference build: + +```sh +python3 -m pip install numpy +make check +``` + +32-bit Windows build from Linux: + +```sh +python3 -m pip install numpy +make windows WINDOWS_CC=/path/to/i686-w64-mingw32-clang +``` + +Change compiled defaults if desired: + +```sh +make windows DEFAULT_TOKENS=256 DEFAULT_CONTEXT=512 +``` + +These only change defaults; users can always pass `-n` and `-c` at runtime. +The GitHub Actions workflow exposes the same two values through +**Run workflow** and refuses to upload anything over 1,474,560 bytes. + +## Old Windows status + +The release is a 32-bit i386 PE console program built against the old +Windows-provided `msvcrt.dll`, with PE OS and subsystem versions set to 4.0. +It is suitable for testing on modern 32/64-bit Windows first. Windows 95/98 +compatibility is an explicit experimental target: the imports are old Win32 +APIs, but this artifact has not yet been executed on physical Win9x hardware. +CPU speed and RAM, rather than disk space, will be the practical constraint. + +## License + +MIT. See `LICENSE` and `THIRD_PARTY_NOTICES.md`. diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000..30cd3e9 --- /dev/null +++ b/THIRD_PARTY_NOTICES.md @@ -0,0 +1,15 @@ +# Third-party notices + +The inference code is derived from Andrej Karpathy's +[`llama2.c`](https://github.com/karpathy/llama2.c), under the MIT License. +The original license is preserved in `licenses/llama2.c-MIT.txt`. + +The bundled model is [`shibatch/tiny1m`](https://huggingface.co/shibatch/tiny1m). +Its model card declares the model and assets MIT licensed. It was trained on +the TinyStories dataset. + +Release builds use the official +[`llvm-mingw`](https://github.com/mstorsjo/llvm-mingw) toolchain to produce a +32-bit PE executable linked against the Windows-provided `msvcrt.dll` and +`kernel32.dll`. + diff --git a/licenses/llama2.c-MIT.txt b/licenses/llama2.c-MIT.txt new file mode 100644 index 0000000..2ad1222 --- /dev/null +++ b/licenses/llama2.c-MIT.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 Andrej + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/model/README.md b/model/README.md new file mode 100644 index 0000000..73b4100 --- /dev/null +++ b/model/README.md @@ -0,0 +1,27 @@ +# Bundled model + +The build downloads [`shibatch/tiny1m`](https://huggingface.co/shibatch/tiny1m), +an MIT-licensed Llama 2 model trained on TinyStories. Its architecture is: + +- 896,256 parameters +- 512-token vocabulary +- 128 hidden width, 4 layers, 2 attention heads +- 352-wide feed-forward layers +- 256-token trained context + +`tools/legacy_to_q8.py` converts the upstream 3,805,724-byte float checkpoint +to llama2.c Q8_0 with groups of 32. The result is 1,055,488 bytes. The +tokenizer is 6,227 bytes. + +Reproducibility hashes: + +| File | SHA-256 | +| --- | --- | +| Upstream `model.bin` | `80ba15bd7feed5f78c13c33bac83992617ff030661f16abbfed8fba0a6e335c1` | +| Upstream `tokenizer.bin` | `037cb335abb25d1fa9e8ecae30ed2a3a8ace9302862ebcdc05d51a6bbb10c312` | +| Converted `tiny1m-q8.bin` | `43c16d0d0702c1f628f841cad10c01bfde4a3a94d7784221a2f6208853ce12c2` | + +The model is a tiny story-completion model, not an instruction-following +assistant. Its value here is that it is a real transformer that fits, runs, +and sometimes produces recognizable simple prose. + diff --git a/src/floppy_llm.c b/src/floppy_llm.c new file mode 100644 index 0000000..405ce69 --- /dev/null +++ b/src/floppy_llm.c @@ -0,0 +1,1238 @@ +/* + * FloppyLLM: a self-contained TinyStories CLI derived from llama2.c. + * The executable can carry its Q8 model and tokenizer as an appended overlay. + */ + +#include +#include +#include +#include +#include +#include +#include +#if defined _WIN32 + #ifndef _WIN32_WINNT + #define _WIN32_WINNT 0x0400 + #endif + #define WIN32_LEAN_AND_MEAN + #include +#else + #include +#endif +// ---------------------------------------------------------------------------- +// Globals +int GS = 0; // group size global for quantization of the weights +int show_stats = 0; + +#define OVERLAY_MAGIC "FLLM0001" + +#ifndef DEFAULT_MAX_NEW_TOKENS +#define DEFAULT_MAX_NEW_TOKENS 128 +#endif + +#ifndef DEFAULT_CONTEXT_OVERRIDE +#define DEFAULT_CONTEXT_OVERRIDE 0 +#endif + +typedef struct { + char magic[8]; + uint32_t model_offset; + uint32_t model_size; + uint32_t tokenizer_offset; + uint32_t tokenizer_size; +} OverlayFooter; + +int find_overlay(const char* path, OverlayFooter* footer) { + FILE* file = fopen(path, "rb"); + if (!file) { return 0; } + if (fseek(file, -(long)sizeof(OverlayFooter), SEEK_END) != 0 || + fread(footer, sizeof(OverlayFooter), 1, file) != 1) { + fclose(file); + return 0; + } + fclose(file); + return memcmp(footer->magic, OVERLAY_MAGIC, 8) == 0; +} + +int executable_path(char* output, size_t output_size, const char* argv0) { +#if defined _WIN32 + DWORD length = GetModuleFileNameA(NULL, output, (DWORD)output_size); + if (length == 0 || length >= output_size) { return 0; } + return 1; +#elif defined __linux__ + ssize_t length = readlink("/proc/self/exe", output, output_size - 1); + if (length > 0) { + output[length] = '\0'; + return 1; + } +#endif + if (!argv0 || strlen(argv0) >= output_size) { return 0; } + strcpy(output, argv0); + return 1; +} + +// ---------------------------------------------------------------------------- +// Transformer model + +typedef struct { + int dim; // transformer dimension + int hidden_dim; // for ffn layers + int n_layers; // number of layers + int n_heads; // number of query heads + int n_kv_heads; // number of key/value heads (can be < query heads because of multiquery) + int vocab_size; // vocabulary size, usually 256 (byte-level) + int seq_len; // max sequence length +} Config; + +typedef struct { + int8_t* q; // quantized values + float* s; // scaling factors +} QuantizedTensor; + +typedef struct { + // token embedding table + QuantizedTensor *q_tokens; // (vocab_size, dim) + float* token_embedding_table; // same, but dequantized + + // weights for rmsnorms + float* rms_att_weight; // (layer, dim) rmsnorm weights + float* rms_ffn_weight; // (layer, dim) + // weights for matmuls. note dim == n_heads * head_size + QuantizedTensor *wq; // (layer, dim, n_heads * head_size) + QuantizedTensor *wk; // (layer, dim, n_kv_heads * head_size) + QuantizedTensor *wv; // (layer, dim, n_kv_heads * head_size) + QuantizedTensor *wo; // (layer, n_heads * head_size, dim) + // weights for ffn + QuantizedTensor *w1; // (layer, hidden_dim, dim) + QuantizedTensor *w2; // (layer, dim, hidden_dim) + QuantizedTensor *w3; // (layer, hidden_dim, dim) + // final rmsnorm + float* rms_final_weight; // (dim,) + // (optional) classifier weights for the logits, on the last layer + QuantizedTensor *wcls; +} TransformerWeights; + +typedef struct { + // current wave of activations + float *x; // activation at current time stamp (dim,) + float *xb; // same, but inside a residual branch (dim,) + float *xb2; // an additional buffer just for convenience (dim,) + float *hb; // buffer for hidden dimension in the ffn (hidden_dim,) + float *hb2; // buffer for hidden dimension in the ffn (hidden_dim,) + QuantizedTensor xq; // quantized x (dim,) + QuantizedTensor hq; // quantized hb (hidden_dim,) + float *q; // query (dim,) + float *k; // key (dim,) + float *v; // value (dim,) + float *att; // buffer for scores/attention values (n_heads, seq_len) + float *logits; // output logits + // kv cache + float* key_cache; // (layer, seq_len, dim) + float* value_cache; // (layer, seq_len, dim) +} RunState; + +typedef struct { + Config config; // the hyperparameters of the architecture (the blueprint) + TransformerWeights weights; // the weights of the model + RunState state; // buffers for the "wave" of activations in the forward pass + unsigned char* data; // model container kept in memory + size_t file_size; // size of the checkpoint file in bytes +} Transformer; + +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)); + s->hb = calloc(p->hidden_dim, sizeof(float)); + s->hb2 = calloc(p->hidden_dim, sizeof(float)); + s->xq = (QuantizedTensor) { .q = calloc(p->dim, sizeof(int8_t)), .s = calloc(p->dim, sizeof(float)) }; + s->hq = (QuantizedTensor) { .q = calloc(p->hidden_dim, sizeof(int8_t)), .s = calloc(p->hidden_dim, sizeof(float)) }; + s->q = calloc(p->dim, sizeof(float)); + s->k = calloc(kv_dim, sizeof(float)); + s->v = calloc(kv_dim, sizeof(float)); + s->att = calloc(p->n_heads * p->seq_len, sizeof(float)); + s->logits = calloc(p->vocab_size, 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 + || !s->value_cache) { + fprintf(stderr, "malloc failed!\n"); + exit(EXIT_FAILURE); + } +} + +void free_run_state(RunState* s) { + free(s->x); + free(s->xb); + free(s->xb2); + free(s->hb); + free(s->hb2); + free(s->xq.q); + free(s->xq.s); + free(s->hq.q); + free(s->hq.s); + free(s->q); + free(s->k); + free(s->v); + free(s->att); + free(s->logits); + free(s->key_cache); + free(s->value_cache); +} + +// ---------------------------------------------------------------------------- +// Quantization functions + +void dequantize(QuantizedTensor *qx, float* x, int n) { + for (int i = 0; i < n; i++) { + x[i] = qx->q[i] * qx->s[i / GS]; + } +} + +void quantize(QuantizedTensor *qx, float* x, int n) { + int num_groups = n / GS; + float Q_MAX = 127.0f; + + for (int group = 0; group < num_groups; group++) { + + // find the max absolute value in the current group + float wmax = 0.0; + for (int i = 0; i < GS; i++) { + float val = fabs(x[group * GS + i]); + if (val > wmax) { + wmax = val; + } + } + + // calculate and write the scaling factor + float scale = wmax / Q_MAX; + qx->s[group] = scale; + + // calculate and write the quantized values + for (int i = 0; i < GS; i++) { + float quant_value = x[group * GS + i] / scale; // scale + int8_t quantized = (int8_t) round(quant_value); // round and clamp + qx->q[group * GS + i] = quantized; + } + } +} + +/* initialize `n` x quantized tensor (with `size_each` elements), starting from memory pointed at *ptr */ +QuantizedTensor *init_quantized_tensors(void **ptr, int n, int size_each) { + void *p = *ptr; + QuantizedTensor *res = malloc(n * sizeof(QuantizedTensor)); + for(int i=0; idim / p->n_heads; + // first are the parameters that are kept in fp32 (the rmsnorm (1D) weights) + float* fptr = (float*) ptr; // cast our pointer to float* + w->rms_att_weight = fptr; + fptr += p->n_layers * p->dim; + w->rms_ffn_weight = fptr; + fptr += p->n_layers * p->dim; + w->rms_final_weight = fptr; + fptr += p->dim; + + // now read all the quantized weights + ptr = (void*)fptr; // now cast the pointer back to void* + w->q_tokens = init_quantized_tensors(&ptr, 1, p->vocab_size * p->dim); + // dequantize token embedding table + w->token_embedding_table = malloc(p->vocab_size * p->dim * sizeof(float)); + dequantize(w->q_tokens, w->token_embedding_table, p->vocab_size * p->dim); + + w->wq = init_quantized_tensors(&ptr, p->n_layers, p->dim * (p->n_heads * head_size)); + w->wk = init_quantized_tensors(&ptr, p->n_layers, p->dim * (p->n_kv_heads * head_size)); + w->wv = init_quantized_tensors(&ptr, p->n_layers, p->dim * (p->n_kv_heads * head_size)); + w->wo = init_quantized_tensors(&ptr, p->n_layers, (p->n_heads * head_size) * p->dim); + + w->w1 = init_quantized_tensors(&ptr, p->n_layers, p->dim * p->hidden_dim); + w->w2 = init_quantized_tensors(&ptr, p->n_layers, p->hidden_dim * p->dim); + w->w3 = init_quantized_tensors(&ptr, p->n_layers, p->dim * p->hidden_dim); + + w->wcls = shared_classifier ? w->q_tokens : init_quantized_tensors(&ptr, 1, p->dim * p->vocab_size); +} + +void read_checkpoint(char* checkpoint, uint32_t checkpoint_offset, Config* config, + TransformerWeights* weights, unsigned char** data, + size_t* file_size) { + FILE *file = fopen(checkpoint, "rb"); + if (!file) { fprintf(stderr, "Couldn't open file %s\n", checkpoint); exit(EXIT_FAILURE); } + if (fseek(file, (long)checkpoint_offset, SEEK_SET) != 0) { + fprintf(stderr, "Couldn't seek to embedded model\n"); + exit(EXIT_FAILURE); + } + // read in magic number (uint32), has to be 0x616b3432, i.e. "ak42" in ASCII + uint32_t magic_number; + if (fread(&magic_number, sizeof(uint32_t), 1, file) != 1) { exit(EXIT_FAILURE); } + if (magic_number != 0x616b3432) { fprintf(stderr, "Bad magic number\n"); exit(EXIT_FAILURE); } + // read in the version number (uint32), has to be 2 + int version; + if (fread(&version, sizeof(int), 1, file) != 1) { exit(EXIT_FAILURE); } + if (version != 2) { fprintf(stderr, "Bad version %d, need version 2\n", version); exit(EXIT_FAILURE); } + int header_size = 256; // the header size for version 2 in bytes + // read in the Config + if (fread(config, sizeof(Config), 1, file) != 1) { exit(EXIT_FAILURE); } + // read in flags + uint8_t shared_classifier; // a byte to indicate if the classifier is shared + if (fread(&shared_classifier, sizeof(uint8_t), 1, file) != 1) { exit(EXIT_FAILURE); } + int group_size; // the group size used in quantization + if (fread(&group_size, sizeof(int), 1, file) != 1) { exit(EXIT_FAILURE); } + GS = group_size; // set as global, as it will be used in many places + // figure out the file size + fseek(file, 0, SEEK_END); // move file pointer to end of file + long length = ftell(file); + if (length <= 0) { fprintf(stderr, "Couldn't measure checkpoint\n"); exit(EXIT_FAILURE); } + *file_size = (size_t)length; + *data = malloc(*file_size); + if (!*data) { fprintf(stderr, "malloc failed!\n"); exit(EXIT_FAILURE); } + rewind(file); + if (fread(*data, 1, *file_size, file) != *file_size) { + fprintf(stderr, "Couldn't read checkpoint\n"); + exit(EXIT_FAILURE); + } + fclose(file); + void* weights_ptr = ((char*)*data) + checkpoint_offset + header_size; + memory_map_weights(weights, config, weights_ptr, shared_classifier); +} + +void build_transformer(Transformer *t, char* checkpoint_path, uint32_t checkpoint_offset, + int context_override) { + // read in the Config and the Weights from the checkpoint + read_checkpoint(checkpoint_path, checkpoint_offset, &t->config, &t->weights, + &t->data, &t->file_size); + if (context_override > 0) { + t->config.seq_len = context_override; + } + // allocate the RunState buffers + malloc_run_state(&t->state, &t->config); +} + +void free_transformer(Transformer* t) { + // free QuantizedTensors + free(t->weights.q_tokens); + free(t->weights.token_embedding_table); + free(t->weights.wq); + free(t->weights.wk); + free(t->weights.wv); + free(t->weights.wo); + free(t->weights.w1); + free(t->weights.w2); + free(t->weights.w3); + if(t->weights.wcls != t->weights.q_tokens) { free(t->weights.wcls); } + free(t->data); + // free the RunState buffers + free_run_state(&t->state); +} + +// ---------------------------------------------------------------------------- +// neural net blocks; the dynamics of the Transformer + +void rmsnorm(float* o, float* x, float* weight, int size) { + // calculate sum of squares + float ss = 0.0f; + for (int j = 0; j < size; j++) { + ss += x[j] * x[j]; + } + ss /= size; + ss += 1e-5f; + ss = 1.0f / sqrtf(ss); + // normalize and scale + for (int j = 0; j < size; j++) { + o[j] = weight[j] * (ss * x[j]); + } +} + +void softmax(float* x, int size) { + // find max value (for numerical stability) + float max_val = x[0]; + for (int i = 1; i < size; i++) { + if (x[i] > max_val) { + max_val = x[i]; + } + } + // exp and sum + float sum = 0.0f; + for (int i = 0; i < size; i++) { + x[i] = expf(x[i] - max_val); + sum += x[i]; + } + // normalize + for (int i = 0; i < size; i++) { + x[i] /= sum; + } +} + +void matmul(float* xout, QuantizedTensor *x, QuantizedTensor *w, int n, int d) { + // W (d,n) @ x (n,) -> xout (d,) + // by far the most amount of time is spent inside this little function + // inputs to this function are both quantized + + int i; + #pragma omp parallel for private(i) + for (i = 0; i < d; i++) { + + float val = 0.0f; + int32_t ival = 0; + int in = i * n; + + // do the matmul in groups of GS + int j; + for (j = 0; j <= n - GS; j += GS) { + for (int k = 0; k < GS; k++) { + ival += ((int32_t) x->q[j + k]) * ((int32_t) w->q[in + j + k]); + } + val += ((float) ival) * w->s[(in + j) / GS] * x->s[j / GS]; + ival = 0; + } + + xout[i] = val; + } +} + +float* forward(Transformer* transformer, int token, int pos) { + + // a few convenience variables + Config* p = &transformer->config; + TransformerWeights* w = &transformer->weights; + RunState* s = &transformer->state; + 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; + + // copy the token embedding into x + memcpy(x, w->token_embedding_table + token*dim, dim * sizeof(float)); + + // forward all the layers + for(int l = 0; l < p->n_layers; l++) { + + // attention rmsnorm + rmsnorm(s->xb, x, w->rms_att_weight + l*dim, dim); + + // qkv matmuls for this position + quantize(&s->xq, s->xb, dim); + matmul(s->q, &s->xq, w->wq + l, dim, dim); + matmul(s->k, &s->xq, w->wk + l, dim, kv_dim); + matmul(s->v, &s->xq, w->wv + l, dim, kv_dim); + + // RoPE relative positional encoding: complex-valued rotate q and k in each head + for (int i = 0; i < dim; i+=2) { + int head_dim = i % head_size; + float freq = 1.0f / powf(10000.0f, head_dim / (float)head_size); + float val = pos * freq; + float fcr = cosf(val); + float fci = sinf(val); + int rotn = i < kv_dim ? 2 : 1; // how many vectors? 2 = q & k, 1 = q only + for (int v = 0; v < rotn; v++) { + float* vec = v == 0 ? s->q : s->k; // the vector to rotate (query or key) + float v0 = vec[i]; + float v1 = vec[i+1]; + vec[i] = v0 * fcr - v1 * fci; + vec[i+1] = v0 * fci + v1 * fcr; + } + } + + // save key,value at this time step (pos) to our kv cache + 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; + #pragma omp parallel for private(h) + for (h = 0; h < p->n_heads; h++) { + // get the query vector for this head + float* q = s->q + h * head_size; + // attention scores for this head + float* att = s->att + h * p->seq_len; + // 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 * 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++) { + score += q[i] * k[i]; + } + score /= sqrtf(head_size); + // save the score to the attention buffer + att[t] = score; + } + + // softmax the scores to get attention weights, from 0..pos inclusively + softmax(att, pos + 1); + + // weighted sum of the values, store back into xb + float* xb = s->xb + h * head_size; + 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 * kv_dim + (h / kv_mul) * head_size; + // get the attention weight for this timestep + float a = att[t]; + // accumulate the weighted value into xb + for (int i = 0; i < head_size; i++) { + xb[i] += a * v[i]; + } + } + } + + // final matmul to get the output of the attention + quantize(&s->xq, s->xb, dim); + matmul(s->xb2, &s->xq, w->wo + l, dim, dim); + + // residual connection back into x + 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); + + // Now for FFN in PyTorch we have: self.w2(F.silu(self.w1(x)) * self.w3(x)) + // first calculate self.w1(x) and self.w3(x) + quantize(&s->xq, s->xb, dim); + matmul(s->hb, &s->xq, w->w1 + l, dim, hidden_dim); + matmul(s->hb2, &s->xq, w->w3 + l, dim, hidden_dim); + + // SwiGLU non-linearity + for (int i = 0; i < hidden_dim; i++) { + float val = s->hb[i]; + // silu(x)=x*σ(x), where σ(x) is the logistic sigmoid + val *= (1.0f / (1.0f + expf(-val))); + // elementwise multiply with w3(x) + val *= s->hb2[i]; + s->hb[i] = val; + } + + // final matmul to get the output of the ffn + quantize(&s->hq, s->hb, hidden_dim); + matmul(s->xb, &s->hq, w->w2 + l, hidden_dim, dim); + + // residual connection + for (int i = 0; i < dim; i++) { + x[i] += s->xb[i]; + } + } + + // final rmsnorm + rmsnorm(x, x, w->rms_final_weight, dim); + + // classifier into logits + quantize(&s->xq, x, dim); + matmul(s->logits, &s->xq, w->wcls, dim, p->vocab_size); + return s->logits; +} + +// ---------------------------------------------------------------------------- +// The Byte Pair Encoding (BPE) Tokenizer that translates strings <-> tokens + +typedef struct { + char *str; + int id; +} TokenIndex; + +typedef struct { + char** vocab; + float* vocab_scores; + TokenIndex *sorted_vocab; + int vocab_size; + unsigned int max_token_length; + unsigned char byte_pieces[512]; // stores all single-byte strings +} Tokenizer; + +int compare_tokens(const void *a, const void *b) { + return strcmp(((TokenIndex*)a)->str, ((TokenIndex*)b)->str); +} + +void build_tokenizer(Tokenizer* t, char* tokenizer_path, uint32_t tokenizer_offset, + int vocab_size) { + // i should have written the vocab_size into the tokenizer file... sigh + t->vocab_size = vocab_size; + // malloc space to hold the scores and the strings + t->vocab = (char**)malloc(vocab_size * sizeof(char*)); + t->vocab_scores = (float*)malloc(vocab_size * sizeof(float)); + t->sorted_vocab = NULL; // initialized lazily + for (int i = 0; i < 256; i++) { + t->byte_pieces[i * 2] = (unsigned char)i; + t->byte_pieces[i * 2 + 1] = '\0'; + } + // read in the file + FILE *file = fopen(tokenizer_path, "rb"); + if (!file) { fprintf(stderr, "couldn't load %s\n", tokenizer_path); exit(EXIT_FAILURE); } + if (fseek(file, (long)tokenizer_offset, SEEK_SET) != 0) { + fprintf(stderr, "couldn't seek to embedded tokenizer\n"); + exit(EXIT_FAILURE); + } + if (fread(&t->max_token_length, sizeof(int), 1, file) != 1) { fprintf(stderr, "failed read\n"); exit(EXIT_FAILURE); } + int len; + for (int i = 0; i < vocab_size; i++) { + if (fread(t->vocab_scores + i, sizeof(float), 1, file) != 1) { fprintf(stderr, "failed read\n"); exit(EXIT_FAILURE);} + if (fread(&len, sizeof(int), 1, file) != 1) { fprintf(stderr, "failed read\n"); exit(EXIT_FAILURE); } + t->vocab[i] = (char *)malloc(len + 1); + if (fread(t->vocab[i], len, 1, file) != 1) { fprintf(stderr, "failed read\n"); exit(EXIT_FAILURE); } + t->vocab[i][len] = '\0'; // add the string terminating token + } + fclose(file); +} + +void free_tokenizer(Tokenizer* t) { + for (int i = 0; i < t->vocab_size; i++) { free(t->vocab[i]); } + free(t->vocab); + free(t->vocab_scores); + free(t->sorted_vocab); +} + +char* decode(Tokenizer* t, int prev_token, int token) { + char *piece = t->vocab[token]; + // following BOS (1) token, sentencepiece decoder strips any leading whitespace (see PR #89) + if (prev_token == 1 && piece[0] == ' ') { piece++; } + // careful, some tokens designate raw bytes, and look like e.g. '<0x01>' + // parse this and convert and return the actual byte + unsigned char byte_val; + if (sscanf(piece, "<0x%02hhX>", &byte_val) == 1) { + piece = (char*)t->byte_pieces + byte_val * 2; + } + return piece; +} + +void safe_printf(char *piece) { + // piece might be a raw byte token, and we only want to print printable chars or whitespace + // because some of the other bytes can be various control codes, backspace, etc. + if (piece == NULL) { return; } + if (piece[0] == '\0') { return; } + if (piece[1] == '\0') { + unsigned char byte_val = piece[0]; + if (!(isprint(byte_val) || isspace(byte_val))) { + return; // bad byte, don't print it + } + } + printf("%s", piece); +} + +int str_lookup(char *str, TokenIndex *sorted_vocab, int vocab_size) { + // efficiently find the perfect match for str in vocab, return its index or -1 if not found + TokenIndex tok = { .str = str }; // acts as the key to search for + TokenIndex *res = bsearch(&tok, sorted_vocab, vocab_size, sizeof(TokenIndex), compare_tokens); + return res != NULL ? res->id : -1; +} + +void encode(Tokenizer* t, char *text, int8_t bos, int8_t eos, int *tokens, int *n_tokens) { + // encode the string text (input) into an upper-bound preallocated tokens[] array + // bos != 0 means prepend the BOS token (=1), eos != 0 means append the EOS token (=2) + if (text == NULL) { fprintf(stderr, "cannot encode NULL text\n"); exit(EXIT_FAILURE); } + + if (t->sorted_vocab == NULL) { + // lazily malloc and sort the vocabulary + t->sorted_vocab = malloc(t->vocab_size * sizeof(TokenIndex)); + for (int i = 0; i < t->vocab_size; i++) { + t->sorted_vocab[i].str = t->vocab[i]; + t->sorted_vocab[i].id = i; + } + qsort(t->sorted_vocab, t->vocab_size, sizeof(TokenIndex), compare_tokens); + } + + // create a temporary buffer that will store merge candidates of always two consecutive tokens + // *2 for concat, +1 for null terminator +2 for UTF8 (in case max_token_length is 1) + char* str_buffer = malloc((t->max_token_length*2 +1 +2) * sizeof(char)); + size_t str_len = 0; + + // start at 0 tokens + *n_tokens = 0; + + // add optional BOS (=1) token, if desired + if (bos) tokens[(*n_tokens)++] = 1; + + // add_dummy_prefix is true by default + // so prepend a dummy prefix token to the input string, but only if text != "" + // TODO: pretty sure this isn't correct in the general case but I don't have the + // energy to read more of the sentencepiece code to figure out what it's doing + if (text[0] != '\0') { + int dummy_prefix = str_lookup(" ", t->sorted_vocab, t->vocab_size); + tokens[(*n_tokens)++] = dummy_prefix; + } + + // Okay UTF-8 time. This will get messy. Here is the reference from Wikipedia: + // Code point ↔ UTF-8 conversion + // First code point Last code point Byte 1 Byte 2 Byte 3 Byte 4 + // U+0000 U+007F 0xxxxxxx + // U+0080 U+07FF 110xxxxx 10xxxxxx + // U+0800 U+FFFF 1110xxxx 10xxxxxx 10xxxxxx + // U+10000 U+10FFFF 11110xxx 10xxxxxx 10xxxxxx 10xxxxxx + + // process the raw (UTF-8) byte sequence of the input string + for (char *c = text; *c != '\0'; c++) { + + // reset buffer if the current byte is ASCII or a leading byte + // 0xC0 is 11000000, so (*c & 0xC0) keeps the first 2 bits and zeros the rest + // 0x80 is 10000000 + // in UTF-8, all continuation bytes start with "10" in first two bits + // so in English this is: "if this byte is not a continuation byte" + if ((*c & 0xC0) != 0x80) { + // this byte must be either a leading byte (11...) or an ASCII char (0x...) + // => reset our location, as we're starting a new UTF-8 codepoint + str_len = 0; + } + + // append the current byte to the buffer + str_buffer[str_len++] = *c; // ++ is post-increment, incremented after this line + str_buffer[str_len] = '\0'; + + // while the next character is a continuation byte, continue appending + // but if there are too many of them, just stop to avoid overruning str_buffer size. + if ((*(c+1) & 0xC0) == 0x80 && str_len < 4) { + continue; + } + + // ok c+1 is not a continuation byte, so we've read in a full codepoint + int id = str_lookup(str_buffer, t->sorted_vocab, t->vocab_size); + + if (id != -1) { + // we found this codepoint in vocab, add it as a token + tokens[(*n_tokens)++] = id; + } else { + // byte_fallback encoding: just encode each byte as a token + // +3 is here because the first 3 vocab elements are , , + // so the individual bytes only start at index 3 + for (size_t i=0; i < str_len; i++) { + tokens[(*n_tokens)++] = (unsigned char)str_buffer[i] + 3; + } + } + str_len = 0; // protect against a sequence of stray UTF8 continuation bytes + } + + // merge the best consecutive pair each iteration, according the scores in vocab_scores + while (1) { + float best_score = -1e10; + int best_id = -1; + int best_idx = -1; + + for (int i=0; i < (*n_tokens-1); i++) { + // check if we can merge the pair (tokens[i], tokens[i+1]) + sprintf(str_buffer, "%s%s", t->vocab[tokens[i]], t->vocab[tokens[i+1]]); + int id = str_lookup(str_buffer, t->sorted_vocab, t->vocab_size); + if (id != -1 && t->vocab_scores[id] > best_score) { + // this merge pair exists in vocab! record its score and position + best_score = t->vocab_scores[id]; + best_id = id; + best_idx = i; + } + } + + if (best_idx == -1) { + break; // we couldn't find any more pairs to merge, so we're done + } + + // merge the consecutive pair (best_idx, best_idx+1) into new token best_id + tokens[best_idx] = best_id; + // delete token at position best_idx+1, shift the entire sequence back 1 + for (int i = best_idx+1; i < (*n_tokens-1); i++) { + tokens[i] = tokens[i+1]; + } + (*n_tokens)--; // token length decreased + } + + // add optional EOS (=2) token, if desired + if (eos) tokens[(*n_tokens)++] = 2; + + free(str_buffer); +} + +// ---------------------------------------------------------------------------- +// The Sampler, which takes logits and returns a sampled token +// sampling can be done in a few ways: greedy argmax, sampling, top-p sampling + +typedef struct { + float prob; + int index; +} ProbIndex; // struct used when sorting probabilities during top-p sampling + +typedef struct { + int vocab_size; + ProbIndex* probindex; // buffer used in top-p sampling + float temperature; + float topp; + unsigned long long rng_state; +} Sampler; + +int sample_argmax(float* probabilities, int n) { + // return the index that has the highest probability + int max_i = 0; + float max_p = probabilities[0]; + for (int i = 1; i < n; i++) { + if (probabilities[i] > max_p) { + max_i = i; + max_p = probabilities[i]; + } + } + return max_i; +} + +int sample_mult(float* probabilities, int n, float coin) { + // sample index from probabilities (they must sum to 1!) + // coin is a random number in [0, 1), usually from random_f32() + float cdf = 0.0f; + for (int i = 0; i < n; i++) { + cdf += probabilities[i]; + if (coin < cdf) { + return i; + } + } + return n - 1; // in case of rounding errors +} + +int compare(const void* a, const void* b) { + ProbIndex* a_ = (ProbIndex*) a; + ProbIndex* b_ = (ProbIndex*) b; + if (a_->prob > b_->prob) return -1; + if (a_->prob < b_->prob) return 1; + return 0; +} + +int sample_topp(float* probabilities, int n, float topp, ProbIndex* probindex, float coin) { + // top-p sampling (or "nucleus sampling") samples from the smallest set of + // 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". + // coin is a random number in [0, 1), usually from random_f32() + + 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++) { + if (probabilities[i] >= cutoff) { + probindex[n0].index = i; + probindex[n0].prob = probabilities[i]; + n0++; + } + } + qsort(probindex, n0, sizeof(ProbIndex), compare); + + // truncate the list where cumulative probability exceeds topp + float cumulative_prob = 0.0f; + 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; + break; // we've exceeded topp by including last_idx + } + } + + // sample from the truncated list + float r = coin * cumulative_prob; + float cdf = 0.0f; + for (int i = 0; i <= last_idx; i++) { + cdf += probindex[i].prob; + if (r < cdf) { + return probindex[i].index; + } + } + return probindex[last_idx].index; // in case of rounding errors +} + +void build_sampler(Sampler* sampler, int vocab_size, float temperature, float topp, unsigned long long rng_seed) { + sampler->vocab_size = vocab_size; + sampler->temperature = temperature; + sampler->topp = topp; + sampler->rng_state = rng_seed; + // buffer only used with nucleus sampling; may not need but it's ~small + sampler->probindex = malloc(sampler->vocab_size * sizeof(ProbIndex)); +} + +void free_sampler(Sampler* sampler) { + free(sampler->probindex); +} + +unsigned int random_u32(unsigned long long *state) { + // xorshift rng: https://en.wikipedia.org/wiki/Xorshift#xorshift.2A + *state ^= *state >> 12; + *state ^= *state << 25; + *state ^= *state >> 27; + return (*state * 0x2545F4914F6CDD1Dull) >> 32; +} +float random_f32(unsigned long long *state) { // random float32 in [0,1) + return (random_u32(state) >> 8) / 16777216.0f; +} + +int sample(Sampler* sampler, float* logits) { + // sample the token given the logits and some hyperparameters + int next; + if (sampler->temperature == 0.0f) { + // greedy argmax sampling: take the token with the highest probability + next = sample_argmax(logits, sampler->vocab_size); + } else { + // apply the temperature to the logits + for (int q=0; qvocab_size; q++) { logits[q] /= sampler->temperature; } + // apply softmax to the logits to get the probabilities for next token + softmax(logits, sampler->vocab_size); + // flip a (float) coin (this is our source of entropy for sampling) + float coin = random_f32(&sampler->rng_state); + // we sample from this distribution to get the next token + if (sampler->topp <= 0 || sampler->topp >= 1) { + // simply sample from the predicted probability distribution + next = sample_mult(logits, sampler->vocab_size, coin); + } else { + // top-p (nucleus) sampling, clamping the least likely tokens to zero + next = sample_topp(logits, sampler->vocab_size, sampler->topp, sampler->probindex, coin); + } + } + return next; +} + +// ---------------------------------------------------------------------------- +// utilities: time + +long time_in_ms() { + // return time in milliseconds, for benchmarking the model speed +#if defined _WIN32 + return (long)GetTickCount(); +#else + struct timespec time; + clock_gettime(CLOCK_REALTIME, &time); + return time.tv_sec * 1000 + time.tv_nsec / 1000000; +#endif +} + +// ---------------------------------------------------------------------------- +// generation loop + +void generate(Transformer *transformer, Tokenizer *tokenizer, Sampler *sampler, + char *prompt, int max_new_tokens) { + char *empty_prompt = ""; + if (prompt == NULL) { prompt = empty_prompt; } + + // encode the (string) prompt into tokens sequence + int num_prompt_tokens = 0; + int* prompt_tokens = (int*)malloc((strlen(prompt)+3) * sizeof(int)); // +3 for '\0', ?BOS, ?EOS + encode(tokenizer, prompt, 1, 0, prompt_tokens, &num_prompt_tokens); + if (num_prompt_tokens < 1) { + fprintf(stderr, "something is wrong, expected at least 1 prompt token\n"); + exit(EXIT_FAILURE); + } + int steps = num_prompt_tokens + max_new_tokens; + if (max_new_tokens == 0 || steps > transformer->config.seq_len) { + steps = transformer->config.seq_len; + } + if (num_prompt_tokens >= transformer->config.seq_len) { + fprintf(stderr, "prompt is too long for the %d-token context\n", + transformer->config.seq_len); + free(prompt_tokens); + return; + } + + // start the main loop + long start = 0; // used to time our code, only initialized after first iteration + int next = 0; // will store the next token in the sequence + int token = prompt_tokens[0]; // kick off with the first token in the prompt + int pos = 0; // position in the sequence + while (pos < steps) { + + // forward the transformer to get logits for the next token + float* logits = forward(transformer, token, pos); + + // advance the state state machine + if (pos < num_prompt_tokens - 1) { + // if we are still processing the input prompt, force the next prompt token + next = prompt_tokens[pos + 1]; + } else { + // otherwise sample the next token from the logits + next = sample(sampler, logits); + } + pos++; + + // data-dependent terminating condition: the BOS (=1) token delimits sequences + if (next == 1) { break; } + + // print the token as string, decode it with the Tokenizer object + char* piece = decode(tokenizer, token, next); + safe_printf(piece); // same as printf("%s", piece), but skips "unsafe" bytes + fflush(stdout); + token = next; + + // init the timer here because the first iteration can be slower + if (start == 0) { start = time_in_ms(); } + } + printf("\n"); + + // report achieved tok/s (pos-1 because the timer starts after first iteration) + if (show_stats && pos > num_prompt_tokens) { + long end = time_in_ms(); + long elapsed = end - start; + if (elapsed < 1) { elapsed = 1; } + fprintf(stderr, "[%d generated tokens, %.2f tok/s]\n", + pos - num_prompt_tokens, + (pos - num_prompt_tokens) / (double)elapsed * 1000.0); + } + + free(prompt_tokens); +} + +int read_stdin(const char* guide, char* buffer, size_t bufsize) { + // read a line from stdin, up to but not including \n + printf("%s", guide); + if (fgets(buffer, bufsize, stdin) != NULL) { + size_t len = strlen(buffer); + if (len > 0 && buffer[len - 1] == '\n') { + buffer[len - 1] = '\0'; // strip newline + } + return 1; + } + return 0; +} + +// ---------------------------------------------------------------------------- +// chat loop +// I manually inspected the tokens for a few chat conversations compared to +// python reference and that seemed ok, but this was not thoroughly tested and +// is not safely implemented, it's more a proof of concept atm. + +void chat(Transformer *transformer, Tokenizer *tokenizer, Sampler *sampler, + char *cli_user_prompt, char *cli_system_prompt, int steps) { + + // buffers for reading the system prompt and user prompt from stdin + // you'll notice they are soomewhat haphazardly and unsafely set atm + char system_prompt[512]; + char user_prompt[512]; + char rendered_prompt[1152]; + int num_prompt_tokens = 0; + int* prompt_tokens = (int*)malloc(1152 * sizeof(int)); + int user_idx; + + // start the main loop + int8_t user_turn = 1; // user starts + int next = 0; // will store the next token in the sequence + int token; // stores the current token to feed into the transformer + int pos = 0; // position in the sequence + while (pos < steps) { + + // when it is the user's turn to contribute tokens to the dialog... + if (user_turn) { + // get the (optional) system prompt at position 0 + if (pos == 0) { + // at position 0, the user can also contribute a system prompt + if (cli_system_prompt == NULL) { + // system prompt was not passed in, attempt to get it from stdin + read_stdin("Enter system prompt (optional): ", system_prompt, sizeof(system_prompt)); + } else { + // system prompt was passed in, use it + strcpy(system_prompt, cli_system_prompt); + } + } + // get the user prompt + if (pos == 0 && cli_user_prompt != NULL) { + // user prompt for position 0 was passed in, use it + strcpy(user_prompt, cli_user_prompt); + } else { + // otherwise get user prompt from stdin + read_stdin("User: ", user_prompt, sizeof(user_prompt)); + } + // render user/system prompts into the Llama 2 Chat schema + if (pos == 0 && system_prompt[0] != '\0') { + char system_template[] = "[INST] <>\n%s\n<>\n\n%s [/INST]"; + sprintf(rendered_prompt, system_template, system_prompt, user_prompt); + } else { + char user_template[] = "[INST] %s [/INST]"; + sprintf(rendered_prompt, user_template, user_prompt); + } + // encode the rendered prompt into tokens + encode(tokenizer, rendered_prompt, 1, 0, prompt_tokens, &num_prompt_tokens); + user_idx = 0; // reset the user index + user_turn = 0; + printf("Assistant: "); + } + + // determine the token to pass into the transformer next + if (user_idx < num_prompt_tokens) { + // if we are still processing the input prompt, force the next prompt token + token = prompt_tokens[user_idx++]; + } else { + // otherwise use the next token sampled from previous turn + token = next; + } + // EOS (=2) token ends the Assistant turn + if (token == 2) { user_turn = 1; } + + // forward the transformer to get logits for the next token + float* logits = forward(transformer, token, pos); + next = sample(sampler, logits); + pos++; + + if (user_idx >= num_prompt_tokens && next != 2) { + // the Assistant is responding, so print its output + char* piece = decode(tokenizer, token, next); + safe_printf(piece); // same as printf("%s", piece), but skips "unsafe" bytes + fflush(stdout); + } + if (next == 2) { printf("\n"); } + } + printf("\n"); + free(prompt_tokens); +} + + +// ---------------------------------------------------------------------------- +// CLI, include only if not testing +#ifndef TESTING + +void print_usage(FILE* stream) { + fprintf(stream, "FloppyLLM - a complete language model in one floppy-sized executable\n\n"); + fprintf(stream, "Usage: floppy-llm [options]\n\n"); + fprintf(stream, "Without --prompt, FloppyLLM starts an interactive prompt loop.\n\n"); + fprintf(stream, "Options:\n"); + fprintf(stream, " -i, --prompt TEXT complete TEXT and exit\n"); + fprintf(stream, " -n, --tokens N maximum new tokens (default %d; 0 fills context)\n", + DEFAULT_MAX_NEW_TOKENS); + fprintf(stream, " -c, --context N runtime context size (compiled default %d; 0 uses model value)\n", + DEFAULT_CONTEXT_OVERRIDE); + fprintf(stream, " -t, --temperature N sampling temperature (default 0.8; 0 is greedy)\n"); + fprintf(stream, " -p, --top-p N nucleus sampling threshold (default 0.9)\n"); + fprintf(stream, " -s, --seed N random seed (default current time)\n"); + fprintf(stream, " --stats print generation speed\n"); + fprintf(stream, " --model FILE use an external Q8 model (developer option)\n"); + fprintf(stream, " --tokenizer FILE use an external tokenizer (developer option)\n"); + fprintf(stream, " -h, --help show this help\n"); + fprintf(stream, " --version show version information\n"); +} + +char* option_value(int argc, char** argv, int* index) { + if (*index + 1 >= argc) { + fprintf(stderr, "missing value after %s\n", argv[*index]); + print_usage(stderr); + exit(EXIT_FAILURE); + } + *index += 1; + return argv[*index]; +} + +int main(int argc, char *argv[]) { + char *external_model = NULL; + char *external_tokenizer = NULL; + char *prompt = NULL; + float temperature = 0.8f; + float topp = 0.9f; + int max_new_tokens = DEFAULT_MAX_NEW_TOKENS; + int context_override = DEFAULT_CONTEXT_OVERRIDE; + unsigned long long rng_seed = 0; + + for (int i = 1; i < argc; i++) { + char* arg = argv[i]; + if (strcmp(arg, "-h") == 0 || strcmp(arg, "--help") == 0) { + print_usage(stdout); + return 0; + } else if (strcmp(arg, "--version") == 0) { + printf("FloppyLLM 0.1.0 (tiny1m, Q8_0 group 32)\n"); + return 0; + } else if (strcmp(arg, "--stats") == 0) { + show_stats = 1; + } else if (strcmp(arg, "-i") == 0 || strcmp(arg, "--prompt") == 0) { + prompt = option_value(argc, argv, &i); + } else if (strcmp(arg, "-n") == 0 || strcmp(arg, "--tokens") == 0) { + max_new_tokens = atoi(option_value(argc, argv, &i)); + } else if (strcmp(arg, "-c") == 0 || strcmp(arg, "--context") == 0) { + context_override = atoi(option_value(argc, argv, &i)); + } else if (strcmp(arg, "-t") == 0 || strcmp(arg, "--temperature") == 0) { + temperature = (float)atof(option_value(argc, argv, &i)); + } else if (strcmp(arg, "-p") == 0 || strcmp(arg, "--top-p") == 0) { + topp = (float)atof(option_value(argc, argv, &i)); + } else if (strcmp(arg, "-s") == 0 || strcmp(arg, "--seed") == 0) { + rng_seed = (unsigned long long)strtoul(option_value(argc, argv, &i), NULL, 10); + } else if (strcmp(arg, "--model") == 0) { + external_model = option_value(argc, argv, &i); + } else if (strcmp(arg, "--tokenizer") == 0) { + external_tokenizer = option_value(argc, argv, &i); + } else { + fprintf(stderr, "unknown option: %s\n", arg); + print_usage(stderr); + return 1; + } + } + + if (temperature < 0.0f) { temperature = 0.0f; } + if (topp < 0.0f || topp > 1.0f) { topp = 0.9f; } + if (max_new_tokens < 0) { max_new_tokens = 0; } + if (context_override > 0 && context_override < 16) { + fprintf(stderr, "context must be at least 16 tokens\n"); + return 1; + } + if (rng_seed == 0) { rng_seed = (unsigned long long)time(NULL); } + + char self_path[1024]; + OverlayFooter footer; + char *checkpoint_path = external_model; + char *tokenizer_path = external_tokenizer; + uint32_t checkpoint_offset = 0; + uint32_t tokenizer_offset = 0; + + if (checkpoint_path == NULL) { + if (!executable_path(self_path, sizeof(self_path), argv[0]) || + !find_overlay(self_path, &footer)) { + fprintf(stderr, "this is an unbundled development runner; use --model and --tokenizer\n"); + return 1; + } + checkpoint_path = self_path; + tokenizer_path = self_path; + checkpoint_offset = footer.model_offset; + tokenizer_offset = footer.tokenizer_offset; + } else if (tokenizer_path == NULL) { + fprintf(stderr, "--tokenizer is required when using --model\n"); + return 1; + } + + Transformer transformer; + build_transformer(&transformer, checkpoint_path, checkpoint_offset, context_override); + + Tokenizer tokenizer; + build_tokenizer(&tokenizer, tokenizer_path, tokenizer_offset, + transformer.config.vocab_size); + + Sampler sampler; + build_sampler(&sampler, transformer.config.vocab_size, temperature, topp, rng_seed); + + if (prompt != NULL) { + generate(&transformer, &tokenizer, &sampler, prompt, max_new_tokens); + } else { + char input[1024]; + printf("FloppyLLM 0.1.0 | 896K parameters | context %d\n", + transformer.config.seq_len); + printf("Tiny story completion model. Type /help or /exit.\n\n"); + while (read_stdin("PROMPT> ", input, sizeof(input))) { + if (strcmp(input, "/exit") == 0 || strcmp(input, "/quit") == 0) { + break; + } + if (strcmp(input, "/help") == 0) { + printf("Enter the beginning of a simple story. /exit quits.\n"); + continue; + } + if (input[0] == '\0') { continue; } + generate(&transformer, &tokenizer, &sampler, input, max_new_tokens); + printf("\n"); + } + } + + free_sampler(&sampler); + free_tokenizer(&tokenizer); + free_transformer(&transformer); + return 0; +} +#endif diff --git a/tools/check_size.py b/tools/check_size.py new file mode 100644 index 0000000..bf5f2a8 --- /dev/null +++ b/tools/check_size.py @@ -0,0 +1,13 @@ +#!/usr/bin/env python3 +"""Fail when an artifact is larger than the requested byte budget.""" + +import sys +from pathlib import Path + +path = Path(sys.argv[1]) +budget = int(sys.argv[2]) +size = path.stat().st_size +remaining = budget - size +print(f"{path}: {size:,}/{budget:,} bytes ({remaining:,} bytes free)") +raise SystemExit(size > budget) + diff --git a/tools/legacy_to_q8.py b/tools/legacy_to_q8.py new file mode 100644 index 0000000..41b434c --- /dev/null +++ b/tools/legacy_to_q8.py @@ -0,0 +1,78 @@ +import struct +import sys + +import numpy as np + + +src, dst = sys.argv[1:3] +raw = open(src, "rb").read() +dim, hidden, layers, heads, kv_heads, vocab_signed, seq_len = struct.unpack_from( + "<7i", raw, 0 +) +vocab = abs(vocab_signed) +shared = vocab_signed > 0 +head_size = dim // heads +kv_dim = dim * kv_heads // heads + +values = np.frombuffer(raw, dtype="