From 2f51451561ee24a671ba69a617749e14cf275b1a Mon Sep 17 00:00:00 2001 From: saharNooby Date: Thu, 30 Mar 2023 17:55:30 +0400 Subject: [PATCH] Initial commit --- README.md | 378 ++----------------- examples/CMakeLists.txt | 1 + examples/main_rwkv/CMakeLists.txt | 4 + examples/main_rwkv/README.md | 3 + examples/main_rwkv/main_rwkv.cpp | 520 +++++++++++++++++++++++++++ rwkv/convert_pytorch_rwkv_to_ggml.py | 175 +++++++++ 6 files changed, 738 insertions(+), 343 deletions(-) create mode 100644 examples/main_rwkv/CMakeLists.txt create mode 100644 examples/main_rwkv/README.md create mode 100644 examples/main_rwkv/main_rwkv.cpp create mode 100644 rwkv/convert_pytorch_rwkv_to_ggml.py diff --git a/README.md b/README.md index e30452e..1005b55 100644 --- a/README.md +++ b/README.md @@ -1,365 +1,57 @@ -# llama.cpp +# rwkv.cpp -![llama](https://user-images.githubusercontent.com/1991296/227761327-6d83e30e-2200-41a6-bfbb-f575231c54f4.png) +This is a port of [RWKV-LM by @BlinkDL](https://github.com/BlinkDL/RWKV-LM) to [ggml library by @ggerganov](https://github.com/ggerganov/ggml). The end goal is to allow 4-bit quanized inference on CPU. -[![Actions Status](https://github.com/ggerganov/llama.cpp/workflows/CI/badge.svg)](https://github.com/ggerganov/llama.cpp/actions) -[![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](https://opensource.org/licenses/MIT) +**WORK IN PROGRESS: NOTHING WORKS YET!** If you know C/C++/ggml, please help! -Inference of [LLaMA](https://arxiv.org/abs/2302.13971) model in pure C/C++ +Converting and loading the model works, but I'm not sure if the element/dimension order is correct -- more debugging needed. -**Hot topics:** +## Plan -- [Roadmap (short-term)](https://github.com/ggerganov/llama.cpp/discussions/457) -- Support for [GPT4All](https://github.com/ggerganov/llama.cpp#using-gpt4all) +1. Make FP32 inference work +2. Validate states and logits against [reference implementation](https://github.com/BlinkDL/ChatRWKV/blob/main/RWKV_in_150_lines.py) by creating a testing script +3. Heavily refactor code; optimize where possible +4. Make FP16 inference work +5. Create fancy interface with sockets/shared memory/pipes/etc. +6. Create Python wrapper with sampling and simple chat interface +7. Write a good `README.md` and publish links to this repo +8. Make INT4 inference work +9. Create pull request to main `ggml` repo with all improvements made here -## Description +## Structure -The main goal is to run the model using 4-bit quantization on a MacBook +This repo is based on the [llama.cpp repo](https://github.com/ggerganov/llama.cpp). RWKV-related code is in these directories: -- Plain C/C++ implementation without dependencies -- Apple silicon first-class citizen - optimized via ARM NEON and Accelerate framework -- AVX2 support for x86 architectures -- Mixed F16 / F32 precision -- 4-bit quantization support -- Runs on the CPU +- `./rwkv`: directory containing Python scripts +- `./examples/main_rwkw`: directory containing script that loads and infers RWKV model -This was [hacked in an evening](https://github.com/ggerganov/llama.cpp/issues/33#issuecomment-1465108022) - I have no idea if it works correctly. -Please do not make conclusions about the models based on the results from this implementation. -For all I know, it can be completely wrong. This project is for educational purposes. -New features will probably be added mostly through community contributions. +Please do not change files in other directories -- this will make pulling recent changes easier. -Supported platforms: +## How to use -- [X] Mac OS -- [X] Linux -- [X] Windows (via CMake) -- [X] Docker +### Windows -Supported models: +Requirements: [git](https://gitforwindows.org/), [CMake](https://cmake.org/download/), MSVC compiler, Python 3.x with PyTorch. -- [X] LLaMA -- [X] [Alpaca](https://github.com/ggerganov/llama.cpp#instruction-mode-with-alpaca) -- [X] [GPT4All](https://github.com/ggerganov/llama.cpp#using-gpt4all) +Clone the repo and set it up for build: ---- - -Here is a typical run using LLaMA-7B: - -```java -make -j && ./main -m ./models/7B/ggml-model-q4_0.bin -p "Building a website can be done in 10 simple steps:" -n 512 -I llama.cpp build info: -I UNAME_S: Darwin -I UNAME_P: arm -I UNAME_M: arm64 -I CFLAGS: -I. -O3 -DNDEBUG -std=c11 -fPIC -pthread -DGGML_USE_ACCELERATE -I CXXFLAGS: -I. -I./examples -O3 -DNDEBUG -std=c++11 -fPIC -pthread -I LDFLAGS: -framework Accelerate -I CC: Apple clang version 14.0.0 (clang-1400.0.29.202) -I CXX: Apple clang version 14.0.0 (clang-1400.0.29.202) - -make: Nothing to be done for `default'. -main: seed = 1678486056 -llama_model_load: loading model from './models/7B/ggml-model-q4_0.bin' - please wait ... -llama_model_load: n_vocab = 32000 -llama_model_load: n_ctx = 512 -llama_model_load: n_embd = 4096 -llama_model_load: n_mult = 256 -llama_model_load: n_head = 32 -llama_model_load: n_layer = 32 -llama_model_load: n_rot = 128 -llama_model_load: f16 = 2 -llama_model_load: n_ff = 11008 -llama_model_load: ggml ctx size = 4529.34 MB -llama_model_load: memory_size = 512.00 MB, n_mem = 16384 -llama_model_load: .................................... done -llama_model_load: model size = 4017.27 MB / num tensors = 291 - -main: prompt: 'Building a website can be done in 10 simple steps:' -main: number of tokens in prompt = 15 - 1 -> '' - 8893 -> 'Build' - 292 -> 'ing' - 263 -> ' a' - 4700 -> ' website' - 508 -> ' can' - 367 -> ' be' - 2309 -> ' done' - 297 -> ' in' - 29871 -> ' ' - 29896 -> '1' - 29900 -> '0' - 2560 -> ' simple' - 6576 -> ' steps' - 29901 -> ':' - -sampling parameters: temp = 0.800000, top_k = 40, top_p = 0.950000 - - -Building a website can be done in 10 simple steps: -1) Select a domain name and web hosting plan -2) Complete a sitemap -3) List your products -4) Write product descriptions -5) Create a user account -6) Build the template -7) Start building the website -8) Advertise the website -9) Provide email support -10) Submit the website to search engines -A website is a collection of web pages that are formatted with HTML. HTML is the code that defines what the website looks like and how it behaves. -The HTML code is formatted into a template or a format. Once this is done, it is displayed on the user's browser. -The web pages are stored in a web server. The web server is also called a host. When the website is accessed, it is retrieved from the server and displayed on the user's computer. -A website is known as a website when it is hosted. This means that it is displayed on a host. The host is usually a web server. -A website can be displayed on different browsers. The browsers are basically the software that renders the website on the user's screen. -A website can also be viewed on different devices such as desktops, tablets and smartphones. -Hence, to have a website displayed on a browser, the website must be hosted. -A domain name is an address of a website. It is the name of the website. -The website is known as a website when it is hosted. This means that it is displayed on a host. The host is usually a web server. -A website can be displayed on different browsers. The browsers are basically the software that renders the website on the user’s screen. -A website can also be viewed on different devices such as desktops, tablets and smartphones. Hence, to have a website displayed on a browser, the website must be hosted. -A domain name is an address of a website. It is the name of the website. -A website is an address of a website. It is a collection of web pages that are formatted with HTML. HTML is the code that defines what the website looks like and how it behaves. -The HTML code is formatted into a template or a format. Once this is done, it is displayed on the user’s browser. -A website is known as a website when it is hosted - -main: mem per token = 14434244 bytes -main: load time = 1332.48 ms -main: sample time = 1081.40 ms -main: predict time = 31378.77 ms / 61.41 ms per token -main: total time = 34036.74 ms +```commandline +git clone https://github.com/saharNooby/rwkv.cpp.git +cd rwkv.cpp +cmake . ``` -And here is another demo of running both LLaMA-7B and [whisper.cpp](https://github.com/ggerganov/whisper.cpp) on a single M1 Pro MacBook: +Download an RWKV model from [Huggingface](https://huggingface.co/BlinkDL) and convert it into `ggml` format: -https://user-images.githubusercontent.com/1991296/224442907-7693d4be-acaa-4e01-8b4f-add84093ffff.mp4 - -## Usage - -Here are the step for the LLaMA-7B model: - -```bash -# build this repo -git clone https://github.com/ggerganov/llama.cpp -cd llama.cpp -make - -# obtain the original LLaMA model weights and place them in ./models -ls ./models -65B 30B 13B 7B tokenizer_checklist.chk tokenizer.model - -# install Python dependencies -python3 -m pip install torch numpy sentencepiece - -# convert the 7B model to ggml FP16 format -python3 convert-pth-to-ggml.py models/7B/ 1 - -# quantize the model to 4-bits -python3 quantize.py 7B - -# run the inference -./main -m ./models/7B/ggml-model-q4_0.bin -n 128 +```commandline +python convert_pytorch_rwkv_to_ggml.py C:\RWKV-4-Pile-169M-20220807-8023.pth C:\rwkv.cpp-169M.bin float32 ``` -Currently, it's best to use Python 3.9 or Python 3.10, as `sentencepiece` has not yet published a wheel for Python 3.11. +Compile and run the script: -When running the larger models, make sure you have enough disk space to store all the intermediate files. - -### Memory/Disk Requirements - -As the models are currently fully loaded into memory, you will need adequate disk space to save them -and sufficient RAM to load them. At the moment, memory and disk requirements are the same. - -| model | original size | quantized size (4-bit) | -|-------|---------------|------------------------| -| 7B | 13 GB | 3.9 GB | -| 13B | 24 GB | 7.8 GB | -| 30B | 60 GB | 19.5 GB | -| 65B | 120 GB | 38.5 GB | - -### Interactive mode - -If you want a more ChatGPT-like experience, you can run in interactive mode by passing `-i` as a parameter. -In this mode, you can always interrupt generation by pressing Ctrl+C and enter one or more lines of text which will be converted into tokens and appended to the current context. You can also specify a *reverse prompt* with the parameter `-r "reverse prompt string"`. This will result in user input being prompted whenever the exact tokens of the reverse prompt string are encountered in the generation. A typical use is to use a prompt which makes LLaMa emulate a chat between multiple users, say Alice and Bob, and pass `-r "Alice:"`. - -Here is an example few-shot interaction, invoked with the command - -```bash -# default arguments using 7B model -./examples/chat.sh - -# advanced chat with 13B model -./examples/chat-13B.sh - -# custom arguments using 13B model -./main -m ./models/13B/ggml-model-q4_0.bin -n 256 --repeat_penalty 1.0 --color -i -r "User:" -f prompts/chat-with-bob.txt +```commandline +cmake --build . --config Release +bin\Release\main_rwkv.exe "C:\rwkv.cpp-169M.bin" 123 "C:\state_in.bin" "C:\state_out.bin" "C:\logits_out.bin" ``` -Note the use of `--color` to distinguish between user input and generated text. - -![image](https://user-images.githubusercontent.com/1991296/224575029-2af3c7dc-5a65-4f64-a6bb-517a532aea38.png) - -### Instruction mode with Alpaca - -1. First, download the `ggml` Alpaca model into the `./models` folder -2. Run the `main` tool like this: - -``` -./examples/alpaca.sh -``` - -Sample run: - -``` -== Running in interactive mode. == - - Press Ctrl+C to interject at any time. - - Press Return to return control to LLaMa. - - If you want to submit another line, end your input in '\'. - - Below is an instruction that describes a task. Write a response that appropriately completes the request. - -> How many letters are there in the English alphabet? -There 26 letters in the English Alphabet -> What is the most common way of transportation in Amsterdam? -The majority (54%) are using public transit. This includes buses, trams and metros with over 100 lines throughout the city which make it very accessible for tourists to navigate around town as well as locals who commute by tram or metro on a daily basis -> List 5 words that start with "ca". -cadaver, cauliflower, cabbage (vegetable), catalpa (tree) and Cailleach. -> -``` - -### Using [GPT4All](https://github.com/nomic-ai/gpt4all) - -- Obtain the `gpt4all-lora-quantized.bin` model -- It is distributed in the old `ggml` format which is now obsoleted -- You have to convert it to the new format using [./convert-gpt4all-to-ggml.py](./convert-gpt4all-to-ggml.py): - - ```bash - python3 convert-gpt4all-to-ggml.py models/gpt4all-7B/gpt4all-lora-quantized.bin ./models/tokenizer.model - ``` - -- You can now use the newly generated `gpt4all-lora-quantized.bin` model in exactly the same way as all other models -- The original model is saved in the same folder with a suffix `.orig` - -### Obtaining and verifying the Facebook LLaMA original model and Stanford Alpaca model data - -- **Under no circumstances share IPFS, magnet links, or any other links to model downloads anywhere in this respository, including in issues, discussions or pull requests. They will be immediately deleted.** -- The LLaMA models are officially distributed by Facebook and will **never** be provided through this repository. -- Refer to [Facebook's LLaMA repository](https://github.com/facebookresearch/llama/pull/73/files) if you need to request access to the model data. -- Please verify the sha256 checksums of all downloaded model files to confirm that you have the correct model data files before creating an issue relating to your model files. -- The following command will verify if you have all possible latest files in your self-installed `./models` subdirectory: - - `sha256sum --ignore-missing -c SHA256SUMS` on Linux - - or - - `shasum -a 256 --ignore-missing -c SHA256SUMS` on macOS - -- If your issue is with model generation quality then please at least scan the following links and papers to understand the limitations of LLaMA models. This is especially important when choosing an appropriate model size and appreciating both the significant and subtle differences between LLaMA models and ChatGPT: - - LLaMA: - - [Introducing LLaMA: A foundational, 65-billion-parameter large language model](https://ai.facebook.com/blog/large-language-model-llama-meta-ai/) - - [LLaMA: Open and Efficient Foundation Language Models](https://arxiv.org/abs/2302.13971) - - GPT-3 - - [Language Models are Few-Shot Learners](https://arxiv.org/abs/2005.14165) - - GPT-3.5 / InstructGPT / ChatGPT: - - [Aligning language models to follow instructions](https://openai.com/research/instruction-following) - - [Training language models to follow instructions with human feedback](https://arxiv.org/abs/2203.02155) - -### Perplexity (Measuring model quality) - -You can use the `perplexity` example to measure perplexity over the given prompt. For more background, -see https://huggingface.co/docs/transformers/perplexity. However, in general, lower perplexity is better for LLMs. - -#### Latest measurements - -The latest perplexity scores for the various model sizes and quantizations are being tracked in [discussion #406](https://github.com/ggerganov/llama.cpp/discussions/406). `llama.cpp` is measuring very well -compared to the baseline implementations. Quantization has a small negative impact to quality, but, as you can see, running -13B at q4_0 beats the 7B f16 model by a significant amount. - -All measurements are done against wikitext2 test dataset (https://paperswithcode.com/dataset/wikitext-2), with default options (512 length context). -Note that the changing the context length will have a significant impact on perplexity (longer context = better perplexity). -``` -Perplexity - model options -5.5985 - 13B, q4_0 -5.9565 - 7B, f16 -6.3001 - 7B, q4_1 -6.5949 - 7B, q4_0 -6.5995 - 7B, q4_0, --memory_f16 -``` - -#### How to run - -1. Download/extract: https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-raw-v1.zip?ref=salesforce-research -2. Run `./perplexity -m models/7B/ggml-model-q4_0.bin -f wiki.test.raw` -3. Output: -``` -perplexity : calculating perplexity over 655 chunks -24.43 seconds per pass - ETA 4.45 hours -[1]4.5970,[2]5.1807,[3]6.0382,... -``` -And after 4.45 hours, you will have the final perplexity. - -### Android - -You can easily run `llama.cpp` on Android device with [termux](https://play.google.com/store/apps/details?id=com.termux). -First, obtain the [Android NDK](https://developer.android.com/ndk) and then build with CMake: -``` -$ mkdir build-android -$ cd build-android -$ export NDK= -$ cmake -DCMAKE_TOOLCHAIN_FILE=$NDK/build/cmake/android.toolchain.cmake -DANDROID_ABI=arm64-v8a -DANDROID_PLATFORM=android-23 -DCMAKE_C_FLAGS=-march=armv8.4a+dotprod .. -$ make -``` -Install [termux](https://play.google.com/store/apps/details?id=com.termux) on your device and run `termux-setup-storage` to get access to your SD card. -Finally, copy the `llama` binary and the model files to your device storage. Here is a demo of an interactive session running on Pixel 5 phone: - -https://user-images.githubusercontent.com/271616/225014776-1d567049-ad71-4ef2-b050-55b0b3b9274c.mp4 - -### Docker - -#### Prerequisites -* Docker must be installed and running on your system. -* Create a folder to store big models & intermediate files (in ex. im using /llama/models) - -#### Images -We have two Docker images available for this project: - -1. `ghcr.io/ggerganov/llama.cpp:full`: This image includes both the main executable file and the tools to convert LLaMA models into ggml and convert into 4-bit quantization. -2. `ghcr.io/ggerganov/llama.cpp:light`: This image only includes the main executable file. - -#### Usage - -The easiest way to download the models, convert them to ggml and optimize them is with the --all-in-one command which includes the full docker image. - - ```bash -docker run -v /llama/models:/models ghcr.io/ggerganov/llama.cpp:full --all-in-one "/models/" 7B -``` - -On complete, you are ready to play! - -```bash -docker run -v /llama/models:/models ghcr.io/ggerganov/llama.cpp:full --run -m /models/7B/ggml-model-q4_0.bin -p "Building a website can be done in 10 simple steps:" -n 512 -``` - -or with light image: - -```bash -docker run -v /llama/models:/models ghcr.io/ggerganov/llama.cpp:light -m /models/7B/ggml-model-q4_0.bin -p "Building a website can be done in 10 simple steps:" -n 512 -``` - -### Contributing - -- Contributors can open PRs -- Collaborators can push to branches in the `llama.cpp` repo and merge PRs into the `master` branch -- Collaborators will be invited based on contributions -- Any help with managing issues and PRs is very appreciated! -- Make sure to read this: [Inference at the edge](https://github.com/ggerganov/llama.cpp/discussions/205) -- A bit of backstory for those who are interested: [Changelog podcast](https://changelog.com/podcast/532) - -### Coding guidelines - -- Avoid adding third-party dependencies, extra files, extra headers, etc. -- Always consider cross-compatibility with other operating systems and architectures -- Avoid fancy looking modern STL constructs, use basic `for` loops, avoid templates, keep it simple -- There are no strict rules for the code style, but try to follow the patterns in the code (indentation, spaces, etc.). Vertical alignment makes things more readable and easier to batch edit -- Clean-up any trailing whitespaces, use 4 spaces indentation, brackets on same line, `void * ptr`, `int & a` -- See [good first issues](https://github.com/ggerganov/llama.cpp/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22) for tasks suitable for first contributions - +The script will read state from `state_in.bin`, do single inference using the state and token `123` as an input, save new state into `state_out.bin` and logits into `logits_out.bin`. diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index ce3a347..c8f16df 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -30,6 +30,7 @@ include_directories(${CMAKE_CURRENT_SOURCE_DIR}) if (EMSCRIPTEN) else() add_subdirectory(main) + add_subdirectory(main_rwkv) add_subdirectory(quantize) add_subdirectory(perplexity) add_subdirectory(embedding) diff --git a/examples/main_rwkv/CMakeLists.txt b/examples/main_rwkv/CMakeLists.txt new file mode 100644 index 0000000..e59aec1 --- /dev/null +++ b/examples/main_rwkv/CMakeLists.txt @@ -0,0 +1,4 @@ +set(TARGET main_rwkv) +add_executable(${TARGET} main_rwkv.cpp) +target_link_libraries(${TARGET} PRIVATE common ggml ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_11) diff --git a/examples/main_rwkv/README.md b/examples/main_rwkv/README.md new file mode 100644 index 0000000..41464bb --- /dev/null +++ b/examples/main_rwkv/README.md @@ -0,0 +1,3 @@ +# main_rwkv + +Runner for RWKV language model. diff --git a/examples/main_rwkv/main_rwkv.cpp b/examples/main_rwkv/main_rwkv.cpp new file mode 100644 index 0000000..797420a --- /dev/null +++ b/examples/main_rwkv/main_rwkv.cpp @@ -0,0 +1,520 @@ +#include "common.h" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// --- Utilities --- + +#define RWKV_ASSERT(x, ...) \ + do { \ + if (!(x)) { \ + fprintf(stderr, "*** Assertion failed ***\n"); \ + fprintf(stderr, __VA_ARGS__); \ + fprintf(stderr, "\n%s:%d: %s\n", __FILE__, __LINE__, #x); \ + abort(); \ + } \ + } while (0) + +// TODO Move to ggml, if correct +float ggml_get_f32_2d(struct ggml_tensor * tensor, int i, int j) { + RWKV_ASSERT(tensor->n_dims == 2, "Not a 2D tensor"); + RWKV_ASSERT(tensor->type == GGML_TYPE_F32, "Unsupported data type"); + return *(float *) ((char *) tensor->data + j * tensor->nb[1] + i * tensor->nb[0]); +} + +// TODO Move to ggml, if correct +float ggml_get_f32_3d(struct ggml_tensor * tensor, int i, int j, int k) { + RWKV_ASSERT(tensor->n_dims == 3, "Not a 3D tensor"); + RWKV_ASSERT(tensor->type == GGML_TYPE_F32, "Unsupported data type"); + return *(float *) ((char *) tensor->data + k * tensor->nb[2] + j * tensor->nb[1] + i * tensor->nb[0]); +} + +void print_tensor(struct ggml_tensor * tensor, char * name) { + int n_dims = tensor->n_dims; + + if (n_dims == 1) { + int x = tensor->ne[0]; + + RWKV_ASSERT(x >= 6, "Too small tensor"); + + printf( + "1D tensor %s, shape (%d), [%f %f %f ... %f %f %f]\n", + name, + x, + ggml_get_f32_1d(tensor, 0), + ggml_get_f32_1d(tensor, 1), + ggml_get_f32_1d(tensor, 2), + ggml_get_f32_1d(tensor, x - 3), + ggml_get_f32_1d(tensor, x - 2), + ggml_get_f32_1d(tensor, x - 1) + ); + } else if (n_dims == 2) { + int x = tensor->ne[0]; + int y = tensor->ne[1]; + + if (y < 6) { + printf( + "2D tensor %s, shape (%d, %d), [[%f %f %f ... %f %f %f]]\n", + name, + x, + y, + ggml_get_f32_2d(tensor, 0, 0), + ggml_get_f32_2d(tensor, 1, 0), + ggml_get_f32_2d(tensor, 2, 0), + ggml_get_f32_2d(tensor, x - 3, y - 1), + ggml_get_f32_2d(tensor, x - 2, y - 1), + ggml_get_f32_2d(tensor, x - 1, y - 1) + ); + } else { + printf( + "2D tensor %s, shape (%d, %d), [[%f %f %f ... ] ... [ ... %f %f %f]]\n", + name, + x, + y, + ggml_get_f32_2d(tensor, 0, 0), + ggml_get_f32_2d(tensor, 0, 1), + ggml_get_f32_2d(tensor, 0, 2), + ggml_get_f32_2d(tensor, x - 1, y - 3), + ggml_get_f32_2d(tensor, x - 1, y - 2), + ggml_get_f32_2d(tensor, x - 1, y - 1) + ); + } + } else if (n_dims == 3) { + int x = tensor->ne[0]; + int y = tensor->ne[1]; + int z = tensor->ne[2]; + + RWKV_ASSERT(z >= 6, "Too small tensor"); + + printf( + "3D tensor %s, shape (%d, %d, %d), [[[%f %f %f ...] ... [... %f %f %f]]]\n", + name, + x, + y, + z, + ggml_get_f32_3d(tensor, 0, 0, 0), + ggml_get_f32_3d(tensor, 0, 0, 1), + ggml_get_f32_3d(tensor, 0, 0, 2), + ggml_get_f32_3d(tensor, x - 1, y - 1, z - 3), + ggml_get_f32_3d(tensor, x - 1, y - 1, z - 2), + ggml_get_f32_3d(tensor, x - 1, y - 1, z - 1) + ); + } else { + RWKV_ASSERT(false, "Unsupported dimension count %d", n_dims); + } +} + +// Prints tensor name, dimensionality, shape and part of its contents. +#define PRINT_TENSOR(x) print_tensor(x, #x) + +// Computes value of the tensor and all tensors it depends on. +void compute_graph(struct ggml_context * ctx, struct ggml_tensor * tensor) { + struct ggml_cgraph graph = ggml_build_forward(tensor); + + // TODO Move to script arguments + graph.n_threads = std::max(1, (int32_t) std::thread::hardware_concurrency() / 2); + + ggml_graph_compute(ctx, &graph); +} + +// --- Model definition and loading code --- + +struct rwkv_layer { + struct ggml_tensor * ln1_weight; + struct ggml_tensor * ln1_bias; + + // RWKV, also called "attention" by the author. + struct ggml_tensor * att_time_mix_k; + struct ggml_tensor * att_time_mix_v; + struct ggml_tensor * att_time_mix_r; + struct ggml_tensor * att_time_first; + struct ggml_tensor * att_time_decay; + struct ggml_tensor * att_key; + struct ggml_tensor * att_value; + struct ggml_tensor * att_receptance; + struct ggml_tensor * att_output; + + struct ggml_tensor * ln2_weight; + struct ggml_tensor * ln2_bias; + + // FFN. + struct ggml_tensor * ffn_time_mix_k; + struct ggml_tensor * ffn_time_mix_r; + struct ggml_tensor * ffn_key; + struct ggml_tensor * ffn_value; + struct ggml_tensor * ffn_receptance; +}; + +struct rwkv_model { + int32_t n_vocab; + int32_t n_embed; + int32_t n_layer; + // 0 for float32, 1 for float16. + int32_t data_type; + + struct ggml_tensor * emb; + + struct ggml_tensor * ln0_weight; + struct ggml_tensor * ln0_bias; + + std::vector layers; + + struct ggml_tensor * ln_out_weight; + struct ggml_tensor * ln_out_bias; + + struct ggml_tensor * head; +}; + +// Reads single int32 value from a file. +void read_int32(FILE * file, int32_t * dest) { + // TODO Will not read correct values on machine with different endianness + RWKV_ASSERT(fread(dest, 4, 1, file) == 1, "Failed to read an int32 value from a file"); +} + +// Finds model parameter by key and sets it into dest. +// If the parameter was not found, aborts the execution. +void set_parameter(std::unordered_map * parameters, char * key, struct ggml_tensor ** dest) { + struct ggml_tensor * parameter = (*parameters)[key]; + RWKV_ASSERT(parameter != NULL, "Parameter %s not found in model file", key); + *dest = parameter; +} + +// Finds block parameter by block index and key and sets it into dest. +// If the parameter was not found, aborts the execution. +void set_block_parameter(std::unordered_map * parameters, int32_t block_index, char * key, struct ggml_tensor ** dest) { + char full_key[128]; + sprintf(full_key, "blocks.%d.%s", block_index, key); + set_parameter(parameters, full_key, dest); +} + +// Loads RWKV model metadata and parameters from a file. +void load_rwkv_model(ggml_context * ctx, char * file_path, struct rwkv_model * model) { + printf("Loading model from %s\n", file_path); + FILE * file = fopen(file_path, "rb"); + RWKV_ASSERT(file != NULL, "Failed to open file %s", file_path); + + int32_t magic; + read_int32(file, &magic); + RWKV_ASSERT(magic == 0x67676d66, "Unexpected magic value %d", magic); + + int32_t version; + read_int32(file, &version); + RWKV_ASSERT(version == 100, "Unsupported file version %d", version); + + read_int32(file, &(model->n_vocab)); + RWKV_ASSERT(model->n_vocab > 0, "Non-positive n_vocab %d", model->n_vocab); + + read_int32(file, &(model->n_embed)); + RWKV_ASSERT(model->n_embed > 0, "Non-positive n_embed %d", model->n_embed); + + read_int32(file, &(model->n_layer)); + RWKV_ASSERT(model->n_layer > 0, "Non-positive n_layer %d", model->n_layer); + + read_int32(file, &(model->data_type)); + RWKV_ASSERT(model->data_type == 0 || model->data_type == 1, "Unsupported model data type %d", model->data_type); + + RWKV_ASSERT(model->data_type == 0, "Data types other than float32 are not yet supported"); // TODO + + printf("n_vocab = %d\n", model->n_vocab); + printf("n_embed = %d\n", model->n_embed); + printf("n_layer = %d\n", model->n_layer); + + std::unordered_map parameters; + + while (true) { + int32_t dim_count; + fread(&dim_count, 4, 1, file); + + if (feof(file)) { + break; + } + + RWKV_ASSERT(dim_count == 1 || dim_count == 2 || dim_count == 3, "Unsupported dimension count %d", dim_count); + + int32_t key_length; + read_int32(file, &key_length); + RWKV_ASSERT(key_length > 0, "Non-positive key length %d", key_length); + + int32_t data_type; + read_int32(file, &data_type); + RWKV_ASSERT(data_type == 0 || data_type == 1, "Unsupported parameter data type %d", data_type); + + RWKV_ASSERT(data_type == 0, "Data types other than float32 are not yet supported"); // TODO + + struct ggml_tensor * tensor; + + int32_t x = -1; + int32_t y = -1; + int32_t z = -1; + int32_t element_count; + + if (dim_count == 1) { + read_int32(file, &x); + element_count = x; + tensor = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, x); + } else if (dim_count == 2) { + read_int32(file, &x); + read_int32(file, &y); + element_count = x * y; + // Not a typo, dimensions should be reversed here + tensor = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, y, x); + } else if (dim_count == 3) { + read_int32(file, &x); + read_int32(file, &y); + read_int32(file, &z); + element_count = x * y * z; + // Not a typo, dimensions should be reversed here + tensor = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, z, y, x); + } else { + abort(); + } + + std::string key(key_length, 0); + RWKV_ASSERT(fread(&key[0], 1, key_length, file) == key_length, "Failed to read parameter key"); + + // TODO Use ggml_type_size + size_t element_size = data_type == 0 ? 4 : 2; + size_t byte_count = element_count * element_size; + + RWKV_ASSERT(fread(tensor->data, 1, byte_count, file) == byte_count, "Failed to read parameter data"); + + parameters[key] = tensor; + } + + fclose(file); + + printf("Initializing model parameters\n"); + + model->layers.resize(model->n_layer); + + set_parameter(¶meters, "emb.weight", &(model->emb)); + + set_parameter(¶meters, "blocks.0.ln0.weight", &(model->ln0_weight)); + set_parameter(¶meters, "blocks.0.ln0.bias", &(model->ln0_bias)); + + for (int i = 0; i < model->n_layer; i++) { + rwkv_layer layer = model->layers[i]; + + set_block_parameter(¶meters, i, "ln1.weight", &(layer.ln1_weight)); + set_block_parameter(¶meters, i, "ln1.bias", &(layer.ln1_bias)); + + set_block_parameter(¶meters, i, "att.time_mix_k", &(layer.att_time_mix_k)); + set_block_parameter(¶meters, i, "att.time_mix_v", &(layer.att_time_mix_v)); + set_block_parameter(¶meters, i, "att.time_mix_r", &(layer.att_time_mix_r)); + set_block_parameter(¶meters, i, "att.time_first", &(layer.att_time_first)); + set_block_parameter(¶meters, i, "att.time_decay", &(layer.att_time_decay)); + set_block_parameter(¶meters, i, "att.key.weight", &(layer.att_key)); + set_block_parameter(¶meters, i, "att.value.weight", &(layer.att_value)); + set_block_parameter(¶meters, i, "att.receptance.weight", &(layer.att_receptance)); + set_block_parameter(¶meters, i, "att.output.weight", &(layer.att_output)); + + set_block_parameter(¶meters, i, "ln2.weight", &(layer.ln2_weight)); + set_block_parameter(¶meters, i, "ln2.bias", &(layer.ln2_bias)); + + set_block_parameter(¶meters, i, "ffn.time_mix_k", &(layer.ffn_time_mix_k)); + set_block_parameter(¶meters, i, "ffn.time_mix_r", &(layer.ffn_time_mix_r)); + set_block_parameter(¶meters, i, "ffn.key.weight", &(layer.ffn_key)); + set_block_parameter(¶meters, i, "ffn.value.weight", &(layer.ffn_value)); + set_block_parameter(¶meters, i, "ffn.receptance.weight", &(layer.ffn_receptance)); + + model->layers[i] = layer; + } + + set_parameter(¶meters, "ln_out.weight", &(model->ln_out_weight)); + set_parameter(¶meters, "ln_out.bias", &(model->ln_out_bias)); + + set_parameter(¶meters, "head.weight", &(model->head)); + + // Verify order of dimensions + struct ggml_tensor * emb = model->emb; + RWKV_ASSERT(emb->n_dims == 2, "Unexpected dimension count of embedding matrix %d", emb->n_dims); + RWKV_ASSERT(emb->ne[0] == model->n_vocab, "Unexpected dimension of embedding matrix %d", emb->ne[0]); + RWKV_ASSERT(emb->ne[1] == model->n_embed, "Unexpected dimension of embedding matrix %d", emb->ne[1]); +} + +// --- Operators --- + +// TODO Fuse and benchmark +struct ggml_tensor * ggml_layer_norm(ggml_context * ctx, struct ggml_tensor * x, struct ggml_tensor * weight, struct ggml_tensor * bias) { + // LayerNorm in RWKV is: + // x = (x - mean(x)) / sqrt(variance(x) + 1e-5) * weight + bias + // Looks like ggml_norm does the first part, we only need to apply weight & bias + x = ggml_norm(ctx, x); + x = ggml_mul(ctx, x, weight); + x = ggml_add(ctx, x, bias); + return x; +} + +// TODO Fuse and benchmark +struct ggml_tensor * ggml_sigmoid(ggml_context * ctx, struct ggml_tensor * x) { + // ggml has no native sigmoid, but silu(x) / x can be an approximation + x = ggml_silu(ctx, x); + x = ggml_div(ctx, x, x); + return x; +} + +// --- Script --- + +// Usage: main_rwkv.exe "C:\model.bin" "C:\state_in.bin" "C:\state_out.bin" "C:\logits_out.bin" +// Token index is 0-based. +// To start from new state, pass empty string instead of input state file path. +int main(int argc, char ** argv) { + RWKV_ASSERT(argc - 1 == 5, "Expected 5 arguments, got %d", argc - 1); + char * model_path = argv[1]; + char * token_s = argv[2]; + char * state_in_path = argv[3]; + char * state_out_path = argv[4]; + char * logits_out_path = argv[5]; + + int32_t token = strtol(token_s, (char **) NULL, 10); + printf("Token index is %d\n", token); + + bool create_new_state = strcmp(state_in_path, "") == 0; + + // Initialize ggml + struct ggml_init_params params; + // TODO Calculate required memory (automatically or manually) + params.mem_size = 1024 * 1024 * 1024; + params.mem_buffer = NULL; + + struct ggml_context * ctx = ggml_init(params); + + // Load model + struct rwkv_model model; + load_rwkv_model(ctx, model_path, &model); + + PRINT_TENSOR(model.emb); + + int32_t n_vocab = model.n_vocab; + int32_t n_embed = model.n_embed; + int32_t n_layer = model.n_layer; + + // Load input state + int32_t state_element_count = n_layer * 5 * n_embed; + struct ggml_tensor * state = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, state_element_count); + + if (create_new_state) { + printf("Creating new state\n"); + ggml_set_f32(state, 0.0F); + + // TODO Verify correctness + for (int i = 0; i < n_layer; i++) { + // state[5 * i + 4] = -1e30 + int32_t offset_in_bytes = (5 * i + 4) * n_embed * 4; + struct ggml_tensor * state_part = ggml_view_1d(ctx, state, n_embed, offset_in_bytes); + ggml_set_f32(state_part, -1e30F); + } + } else { + printf("Loading state from %s\n", state_in_path); + int32_t state_file_size = state_element_count * 4; + + FILE * state_in_file = fopen(state_in_path, "rb"); + RWKV_ASSERT(state_in_file != NULL, "Failed to open file %s", state_in_path); + + // TODO Saving/loading raw data makes state cache machine-dependent + RWKV_ASSERT(fread(state->data, 1, state_file_size, state_in_file) == state_file_size, "Failed to read tensor data from a file"); + + fclose(state_in_file); + } + + // --- Evaluate model --- + + struct ggml_tensor * ones = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_embed); + ggml_set_f32(ones, 1.0F); + + // x = self.w.emb.weight[token] + struct ggml_tensor * token_index = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_set_i32_1d(token_index, 0, token); + // TODO Is transpose copying the tensor? + struct ggml_tensor * x = ggml_get_rows(ctx, ggml_transpose(ctx, model.emb), token_index); + + compute_graph(ctx, x); + // For token 123, should be [-0.1836, 0.4434, 0.3848 ... -0.4102, -0.3164, -0.1826] + // TODO NOT CORRECT + PRINT_TENSOR(x); + + // x = self.layer_norm(x, self.w.blocks[0].ln0) + x = ggml_layer_norm(ctx, x, model.ln0_weight, model.ln0_bias); + + compute_graph(ctx, x); + // For token 123, should be [-0.4194, 1.1698, 0.7798 ... -1.1838, -0.8716, -0.2765] + // TODO NOT CORRECT + PRINT_TENSOR(x); + + for (int i = 0; i < n_layer; i++) { + auto layer = model.layers[i]; + + // RWKV/time mixing + { + struct ggml_tensor * x0 = ggml_layer_norm(ctx, x, layer.ln1_weight, layer.ln1_bias); + // TODO Implement + x = ggml_add(ctx, x, x0); + } + + // FFN/channel mixing + { + // self.layer_norm(x, self.w.blocks[i].ln2) + struct ggml_tensor * x0 = ggml_layer_norm(ctx, x, layer.ln2_weight, layer.ln2_bias); + // state[5 * i + 0] + int32_t offset_in_bytes = (5 * i + 0) * n_embed * 4; + struct ggml_tensor * x_prev = ggml_view_1d(ctx, state, n_embed, offset_in_bytes); + // xk = x * time_mix_k + state[5 * i + 0] * (1 - time_mix_k) + // xr = x * time_mix_r + state[5 * i + 0] * (1 - time_mix_r) + struct ggml_tensor * xk = ggml_add( + ctx, + ggml_mul(ctx, x0, layer.ffn_time_mix_k), + ggml_mul(ctx, x_prev, ggml_sub(ctx, ones, layer.ffn_time_mix_k)) + ); + struct ggml_tensor * xr = ggml_add( + ctx, + ggml_mul(ctx, x0, layer.ffn_time_mix_r), + ggml_mul(ctx, x_prev, ggml_sub(ctx, ones, layer.ffn_time_mix_r)) + ); + // state[5 * i + 0] = x + ggml_cpy(ctx, x0, x_prev); + + // r = torch.sigmoid(rw @ xr) + struct ggml_tensor * r = ggml_sigmoid( + ctx, + ggml_mul_mat(ctx, layer.ffn_receptance, xr) + ); + // k = torch.square(torch.relu(kw @ xk)) + // TODO Does not work; shape mismatch + //struct ggml_tensor * k = ggml_sqr(ctx, ggml_relu( + // ctx, + // ggml_mul_mat(ctx, layer.ffn_key, xk) + //)); + // r * (vw @ k) + // TODO Does not work; shape mismatch + // x0 = ggml_mul(ctx, r, ggml_mul_mat(ctx, layer.ffn_value, k)); + // x = x + self.channel_mixing(...) + x = ggml_add(ctx, x, x0); + } + } + + // x = self.layer_norm(x, self.w.ln_out) + x = ggml_layer_norm(ctx, x, model.ln_out_weight, model.ln_out_bias); + + // x = (self.w.head.weight @ x).float() + // TODO Is transpose copying the tensor? + struct ggml_tensor * logits = ggml_mul_mat(ctx, x, ggml_transpose(ctx, model.head)); + + compute_graph(ctx, logits); + + PRINT_TENSOR(logits); + + ggml_free(ctx); + + printf("OK\n"); + + return 0; +} diff --git a/rwkv/convert_pytorch_rwkv_to_ggml.py b/rwkv/convert_pytorch_rwkv_to_ggml.py new file mode 100644 index 0000000..68c740f --- /dev/null +++ b/rwkv/convert_pytorch_rwkv_to_ggml.py @@ -0,0 +1,175 @@ +# Converts an RWKV model checkpoint to an rwkv.cpp compatible file. +# Usage: python convert_pytorch_rwkv_to_ggml.py C:\RWKV-4-Pile-169M-20220807-8023.pth C:\rwkv.cpp-169M.bin float32 +# Get model checkpoints from https://huggingface.co/BlinkDL + +# File format: +# +# RWKVModelFile { +# // All ints and floats are in machine byte order. +# // Magic is "ggml" string bytes. +# int32 magic = 0x67676d66; +# int32 version; +# int32 n_vocab; +# int32 n_embed; +# int32 n_layer; +# // 0 if float32, 1 if float16. +# int32 data_type; +# // Read until EOF. +# Parameter[] parameters; +# } +# +# Parameter { +# int32 dim_count; +# int32 key_length; +# // 0 if float32, 1 if float16. +# int32 data_type; +# int32[dim_count] shape; +# // Keys are like "emb.weight", "block.0.ln1.weight". +# uint8[key_length] key_utf8; +# // Can be either float32 or float16. +# float[product(shape)] data; +# } + +import os +import argparse +import struct +import torch +from typing import Dict + +def parse_args(): + parser = argparse.ArgumentParser(description='Convert an RWKV model checkpoint to an rwkv.cpp compatible file') + parser.add_argument('src_path', help='Path to PyTorch checkpoint file') + parser.add_argument('dest_path', help='Path to rwkv.cpp checkpoint file, will be overwritten') + parser.add_argument('data_type', help='Data type, float16 or float32', type=str, choices=['float16', 'float32'], default='float32') + return parser.parse_args() + +def get_layer_count(state_dict: Dict[str, torch.Tensor]) -> int: + n_layer = 0 + + while f'blocks.{n_layer}.ln1.weight' in state_dict: + n_layer += 1 + + assert n_layer > 0 + + return n_layer + +def write_state_dict(state_dict: Dict[str, torch.Tensor], dest_path: str, data_type: str) -> None: + emb_weight: torch.Tensor = state_dict['emb.weight'] + + n_layer = get_layer_count(state_dict) + n_vocab = emb_weight.shape[0] + n_embed = emb_weight.shape[1] + + with open(dest_path, 'wb') as out_file: + out_file.write(struct.pack( + # Disable padding with '=' + '=iiiiii', + # Magic: 'ggmf' in hex + 0x67676d66, + # llama.cpp uses file versions 1+, let's use 100+ for rwkv.cpp + 100, + n_vocab, + n_embed, + n_layer, + 1 if data_type == 'float16' else 0 + )) + + for k in state_dict.keys(): + tensor = state_dict[k].float() + + # Same processing as in "RWKV_in_150_lines.py" + if '.time_' in k: + # (1, 1, n_embed) -> (n_embed) + tensor = tensor.squeeze() + + if '.time_decay' in k: + tensor = -torch.exp(tensor) + + # Keep 1-dim vectors in fp32 + if data_type == 'float16' and len(tensor.shape) > 1: + tensor = tensor.half() + + # ggml stores tensor values in other way than PyTorch, need to flip dimension order here + tensor = torch.permute(tensor, dims=[i for i in reversed(range(len(tensor.shape)))]).contiguous() + + shape = tensor.shape + + print(f'Writing {k}, shape {shape}, type {tensor.dtype}') + + k_encoded: bytes = k.encode('utf-8') + + out_file.write(struct.pack( + '=iii', + len(shape), + len(k_encoded), + 1 if tensor.dtype == torch.float16 else 0 + )) + + for dim in tensor.shape: + out_file.write(struct.pack('=i', dim)) + + out_file.write(k_encoded) + + tensor.numpy().tofile(out_file) + +def main() -> None: + args = parse_args() + + print(f'Reading {args.src_path}') + + state_dict: Dict[str, torch.Tensor] = torch.load(args.src_path, map_location='cpu') + + write_state_dict(state_dict, args.dest_path, args.data_type) + + print('Done') + +# --- Tests --- + +def test() -> None: + test_file_path = 'convert_pytorch_rwkv_to_ggml_test.tmp' + + try: + state_dict: Dict[str, torch.Tensor] = { + 'emb.weight': torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.float32), + 'blocks.0.ln1.weight': torch.tensor([1], dtype=torch.float32) + } + + write_state_dict(state_dict, dest_path=test_file_path, data_type='float32') + + with open(test_file_path, 'rb') as input: + actual_bytes: bytes = input.read() + + expected_bytes: bytes = struct.pack( + '=iiiiii' + 'iiiii10sffffff' + 'iiii19sf', + 0x67676d66, + 100, + 3, + 2, + 1, + 0, + # emb.weight + 2, + 10, + 0, + 2, 3, + 'emb.weight'.encode('utf-8'), + 1.0, 3.0, 5.0, + 2.0, 4.0, 6.0, + # blocks.0.ln1.weight + 1, + 19, + 0, + 1, + 'blocks.0.ln1.weight'.encode('utf-8'), + 1.0 + ) + + assert list(actual_bytes) == list(expected_bytes), f'\nActual: {list(actual_bytes)}\nExpected: {list(expected_bytes)}' + + print('All tests pass') + finally: + if os.path.isfile(test_file_path): + os.remove(test_file_path) + +if __name__ == "__main__": + main()