From 2a01ff5fb140eca5c6159982a75df5a9123f0e6d Mon Sep 17 00:00:00 2001 From: Lizonghang <870644199@qq.com> Date: Wed, 23 Oct 2024 09:42:32 +0400 Subject: [PATCH] init --- Makefile | 8 + README.md | 480 +------ common/arg.cpp | 35 + common/common.cpp | 65 +- common/common.h | 8 +- examples/main/main.cpp | 233 ++-- include/llama.h | 23 +- include/zmq.hpp | 2796 ++++++++++++++++++++++++++++++++++++++++ include/zmq_addon.hpp | 848 ++++++++++++ src/llama.cpp | 1255 ++++++++++++------ 10 files changed, 4725 insertions(+), 1026 deletions(-) create mode 100644 include/zmq.hpp create mode 100644 include/zmq_addon.hpp diff --git a/Makefile b/Makefile index 2793978c..4bd59eb5 100644 --- a/Makefile +++ b/Makefile @@ -262,6 +262,14 @@ MK_CFLAGS = -std=c11 -fPIC MK_CXXFLAGS = -std=c++11 -fPIC MK_NVCCFLAGS = -std=c++11 +ifeq ($(UNAME_S),Darwin) + MK_CPPFLAGS += -I/opt/homebrew/include + MK_LDFLAGS += -L/opt/homebrew/lib -lzmq +else ifeq ($(UNAME_S),Linux) + MK_CPPFLAGS += -I/usr/local/include + MK_LDFLAGS += -L/usr/local/lib -lzmq +endif + ifdef LLAMA_NO_CCACHE GGML_NO_CCACHE := 1 DEPRECATE_WARNING := 1 diff --git a/README.md b/README.md index 41e5e544..cbbd7d55 100644 --- a/README.md +++ b/README.md @@ -1,479 +1,3 @@ -# llama.cpp +# prima.cpp -![llama](https://user-images.githubusercontent.com/1991296/230134379-7181e485-c521-4d23-a0d6-f7b3b61ba524.png) - -[![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](https://opensource.org/licenses/MIT) -[![Server](https://github.com/ggerganov/llama.cpp/actions/workflows/server.yml/badge.svg)](https://github.com/ggerganov/llama.cpp/actions/workflows/server.yml) -[![Conan Center](https://shields.io/conan/v/llama-cpp)](https://conan.io/center/llama-cpp) - -[Roadmap](https://github.com/users/ggerganov/projects/7) / [Project status](https://github.com/ggerganov/llama.cpp/discussions/3471) / [Manifesto](https://github.com/ggerganov/llama.cpp/discussions/205) / [ggml](https://github.com/ggerganov/ggml) - -Inference of Meta's [LLaMA](https://arxiv.org/abs/2302.13971) model (and others) in pure C/C++ - -## Recent API changes - -- [Changelog for `libllama` API](https://github.com/ggerganov/llama.cpp/issues/9289) -- [Changelog for `llama-server` REST API](https://github.com/ggerganov/llama.cpp/issues/9291) - -## Hot topics - -- **Hugging Face Inference Endpoints now support GGUF out of the box! https://github.com/ggerganov/llama.cpp/discussions/9669** -- Hugging Face GGUF editor: [discussion](https://github.com/ggerganov/llama.cpp/discussions/9268) | [tool](https://huggingface.co/spaces/CISCai/gguf-editor) - ----- - -## Description - -The main goal of `llama.cpp` is to enable LLM inference with minimal setup and state-of-the-art performance on a wide -variety of hardware - locally and in the cloud. - -- Plain C/C++ implementation without any dependencies -- Apple silicon is a first-class citizen - optimized via ARM NEON, Accelerate and Metal frameworks -- AVX, AVX2 and AVX512 support for x86 architectures -- 1.5-bit, 2-bit, 3-bit, 4-bit, 5-bit, 6-bit, and 8-bit integer quantization for faster inference and reduced memory use -- Custom CUDA kernels for running LLMs on NVIDIA GPUs (support for AMD GPUs via HIP) -- Vulkan and SYCL backend support -- CPU+GPU hybrid inference to partially accelerate models larger than the total VRAM capacity - -Since its [inception](https://github.com/ggerganov/llama.cpp/issues/33#issuecomment-1465108022), the project has -improved significantly thanks to many contributions. It is the main playground for developing new features for the -[ggml](https://github.com/ggerganov/ggml) library. - -**Supported models:** - -Typically finetunes of the base models below are supported as well. - -- [X] LLaMA πŸ¦™ -- [x] LLaMA 2 πŸ¦™πŸ¦™ -- [x] LLaMA 3 πŸ¦™πŸ¦™πŸ¦™ -- [X] [Mistral 7B](https://huggingface.co/mistralai/Mistral-7B-v0.1) -- [x] [Mixtral MoE](https://huggingface.co/models?search=mistral-ai/Mixtral) -- [x] [DBRX](https://huggingface.co/databricks/dbrx-instruct) -- [X] [Falcon](https://huggingface.co/models?search=tiiuae/falcon) -- [X] [Chinese LLaMA / Alpaca](https://github.com/ymcui/Chinese-LLaMA-Alpaca) and [Chinese LLaMA-2 / Alpaca-2](https://github.com/ymcui/Chinese-LLaMA-Alpaca-2) -- [X] [Vigogne (French)](https://github.com/bofenghuang/vigogne) -- [X] [BERT](https://github.com/ggerganov/llama.cpp/pull/5423) -- [X] [Koala](https://bair.berkeley.edu/blog/2023/04/03/koala/) -- [X] [Baichuan 1 & 2](https://huggingface.co/models?search=baichuan-inc/Baichuan) + [derivations](https://huggingface.co/hiyouga/baichuan-7b-sft) -- [X] [Aquila 1 & 2](https://huggingface.co/models?search=BAAI/Aquila) -- [X] [Starcoder models](https://github.com/ggerganov/llama.cpp/pull/3187) -- [X] [Refact](https://huggingface.co/smallcloudai/Refact-1_6B-fim) -- [X] [MPT](https://github.com/ggerganov/llama.cpp/pull/3417) -- [X] [Bloom](https://github.com/ggerganov/llama.cpp/pull/3553) -- [x] [Yi models](https://huggingface.co/models?search=01-ai/Yi) -- [X] [StableLM models](https://huggingface.co/stabilityai) -- [x] [Deepseek models](https://huggingface.co/models?search=deepseek-ai/deepseek) -- [x] [Qwen models](https://huggingface.co/models?search=Qwen/Qwen) -- [x] [PLaMo-13B](https://github.com/ggerganov/llama.cpp/pull/3557) -- [x] [Phi models](https://huggingface.co/models?search=microsoft/phi) -- [x] [GPT-2](https://huggingface.co/gpt2) -- [x] [Orion 14B](https://github.com/ggerganov/llama.cpp/pull/5118) -- [x] [InternLM2](https://huggingface.co/models?search=internlm2) -- [x] [CodeShell](https://github.com/WisdomShell/codeshell) -- [x] [Gemma](https://ai.google.dev/gemma) -- [x] [Mamba](https://github.com/state-spaces/mamba) -- [x] [Grok-1](https://huggingface.co/keyfan/grok-1-hf) -- [x] [Xverse](https://huggingface.co/models?search=xverse) -- [x] [Command-R models](https://huggingface.co/models?search=CohereForAI/c4ai-command-r) -- [x] [SEA-LION](https://huggingface.co/models?search=sea-lion) -- [x] [GritLM-7B](https://huggingface.co/GritLM/GritLM-7B) + [GritLM-8x7B](https://huggingface.co/GritLM/GritLM-8x7B) -- [x] [OLMo](https://allenai.org/olmo) -- [x] [OLMoE](https://huggingface.co/allenai/OLMoE-1B-7B-0924) -- [x] [Granite models](https://huggingface.co/collections/ibm-granite/granite-code-models-6624c5cec322e4c148c8b330) -- [x] [GPT-NeoX](https://github.com/EleutherAI/gpt-neox) + [Pythia](https://github.com/EleutherAI/pythia) -- [x] [Snowflake-Arctic MoE](https://huggingface.co/collections/Snowflake/arctic-66290090abe542894a5ac520) -- [x] [Smaug](https://huggingface.co/models?search=Smaug) -- [x] [Poro 34B](https://huggingface.co/LumiOpen/Poro-34B) -- [x] [Bitnet b1.58 models](https://huggingface.co/1bitLLM) -- [x] [Flan T5](https://huggingface.co/models?search=flan-t5) -- [x] [Open Elm models](https://huggingface.co/collections/apple/openelm-instruct-models-6619ad295d7ae9f868b759ca) -- [x] [ChatGLM3-6b](https://huggingface.co/THUDM/chatglm3-6b) + [ChatGLM4-9b](https://huggingface.co/THUDM/glm-4-9b) -- [x] [SmolLM](https://huggingface.co/collections/HuggingFaceTB/smollm-6695016cad7167254ce15966) -- [x] [EXAONE-3.0-7.8B-Instruct](https://huggingface.co/LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct) -- [x] [FalconMamba Models](https://huggingface.co/collections/tiiuae/falconmamba-7b-66b9a580324dd1598b0f6d4a) -- [x] [Jais](https://huggingface.co/inceptionai/jais-13b-chat) -- [x] [Bielik-11B-v2.3](https://huggingface.co/collections/speakleash/bielik-11b-v23-66ee813238d9b526a072408a) - -(instructions for supporting more models: [HOWTO-add-model.md](./docs/development/HOWTO-add-model.md)) - -**Multimodal models:** - -- [x] [LLaVA 1.5 models](https://huggingface.co/collections/liuhaotian/llava-15-653aac15d994e992e2677a7e), [LLaVA 1.6 models](https://huggingface.co/collections/liuhaotian/llava-16-65b9e40155f60fd046a5ccf2) -- [x] [BakLLaVA](https://huggingface.co/models?search=SkunkworksAI/Bakllava) -- [x] [Obsidian](https://huggingface.co/NousResearch/Obsidian-3B-V0.5) -- [x] [ShareGPT4V](https://huggingface.co/models?search=Lin-Chen/ShareGPT4V) -- [x] [MobileVLM 1.7B/3B models](https://huggingface.co/models?search=mobileVLM) -- [x] [Yi-VL](https://huggingface.co/models?search=Yi-VL) -- [x] [Mini CPM](https://huggingface.co/models?search=MiniCPM) -- [x] [Moondream](https://huggingface.co/vikhyatk/moondream2) -- [x] [Bunny](https://github.com/BAAI-DCAI/Bunny) - -**Bindings:** - -- Python: [abetlen/llama-cpp-python](https://github.com/abetlen/llama-cpp-python) -- Go: [go-skynet/go-llama.cpp](https://github.com/go-skynet/go-llama.cpp) -- Node.js: [withcatai/node-llama-cpp](https://github.com/withcatai/node-llama-cpp) -- JS/TS (llama.cpp server client): [lgrammel/modelfusion](https://modelfusion.dev/integration/model-provider/llamacpp) -- JS/TS (Programmable Prompt Engine CLI): [offline-ai/cli](https://github.com/offline-ai/cli) -- JavaScript/Wasm (works in browser): [tangledgroup/llama-cpp-wasm](https://github.com/tangledgroup/llama-cpp-wasm) -- Typescript/Wasm (nicer API, available on npm): [ngxson/wllama](https://github.com/ngxson/wllama) -- Ruby: [yoshoku/llama_cpp.rb](https://github.com/yoshoku/llama_cpp.rb) -- Rust (more features): [edgenai/llama_cpp-rs](https://github.com/edgenai/llama_cpp-rs) -- Rust (nicer API): [mdrokz/rust-llama.cpp](https://github.com/mdrokz/rust-llama.cpp) -- Rust (more direct bindings): [utilityai/llama-cpp-rs](https://github.com/utilityai/llama-cpp-rs) -- C#/.NET: [SciSharp/LLamaSharp](https://github.com/SciSharp/LLamaSharp) -- Scala 3: [donderom/llm4s](https://github.com/donderom/llm4s) -- Clojure: [phronmophobic/llama.clj](https://github.com/phronmophobic/llama.clj) -- React Native: [mybigday/llama.rn](https://github.com/mybigday/llama.rn) -- Java: [kherud/java-llama.cpp](https://github.com/kherud/java-llama.cpp) -- Zig: [deins/llama.cpp.zig](https://github.com/Deins/llama.cpp.zig) -- Flutter/Dart: [netdur/llama_cpp_dart](https://github.com/netdur/llama_cpp_dart) -- PHP (API bindings and features built on top of llama.cpp): [distantmagic/resonance](https://github.com/distantmagic/resonance) [(more info)](https://github.com/ggerganov/llama.cpp/pull/6326) -- Guile Scheme: [guile_llama_cpp](https://savannah.nongnu.org/projects/guile-llama-cpp) - -**UI:** - -Unless otherwise noted these projects are open-source with permissive licensing: - -- [MindWorkAI/AI-Studio](https://github.com/MindWorkAI/AI-Studio) (FSL-1.1-MIT) -- [iohub/collama](https://github.com/iohub/coLLaMA) -- [janhq/jan](https://github.com/janhq/jan) (AGPL) -- [nat/openplayground](https://github.com/nat/openplayground) -- [Faraday](https://faraday.dev/) (proprietary) -- [LMStudio](https://lmstudio.ai/) (proprietary) -- [Layla](https://play.google.com/store/apps/details?id=com.laylalite) (proprietary) -- [ramalama](https://github.com/containers/ramalama) (MIT) -- [LocalAI](https://github.com/mudler/LocalAI) (MIT) -- [LostRuins/koboldcpp](https://github.com/LostRuins/koboldcpp) (AGPL) -- [Mozilla-Ocho/llamafile](https://github.com/Mozilla-Ocho/llamafile) -- [nomic-ai/gpt4all](https://github.com/nomic-ai/gpt4all) -- [ollama/ollama](https://github.com/ollama/ollama) -- [oobabooga/text-generation-webui](https://github.com/oobabooga/text-generation-webui) (AGPL) -- [psugihara/FreeChat](https://github.com/psugihara/FreeChat) -- [cztomsik/ava](https://github.com/cztomsik/ava) (MIT) -- [ptsochantaris/emeltal](https://github.com/ptsochantaris/emeltal) -- [pythops/tenere](https://github.com/pythops/tenere) (AGPL) -- [RAGNA Desktop](https://ragna.app/) (proprietary) -- [RecurseChat](https://recurse.chat/) (proprietary) -- [semperai/amica](https://github.com/semperai/amica) -- [withcatai/catai](https://github.com/withcatai/catai) -- [Mobile-Artificial-Intelligence/maid](https://github.com/Mobile-Artificial-Intelligence/maid) (MIT) -- [Msty](https://msty.app) (proprietary) -- [LLMFarm](https://github.com/guinmoon/LLMFarm?tab=readme-ov-file) (MIT) -- [KanTV](https://github.com/zhouwg/kantv?tab=readme-ov-file)(Apachev2.0 or later) -- [Dot](https://github.com/alexpinel/Dot) (GPL) -- [MindMac](https://mindmac.app) (proprietary) -- [KodiBot](https://github.com/firatkiral/kodibot) (GPL) -- [eva](https://github.com/ylsdamxssjxxdd/eva) (MIT) -- [AI Sublime Text plugin](https://github.com/yaroslavyaroslav/OpenAI-sublime-text) (MIT) -- [AIKit](https://github.com/sozercan/aikit) (MIT) -- [LARS - The LLM & Advanced Referencing Solution](https://github.com/abgulati/LARS) (AGPL) -- [LLMUnity](https://github.com/undreamai/LLMUnity) (MIT) -- [Llama Assistant](https://github.com/vietanhdev/llama-assistant) (GPL) - -*(to have a project listed here, it should clearly state that it depends on `llama.cpp`)* - -**Tools:** - -- [akx/ggify](https://github.com/akx/ggify) – download PyTorch models from HuggingFace Hub and convert them to GGML -- [akx/ollama-dl](https://github.com/akx/ollama-dl) – download models from the Ollama library to be used directly with llama.cpp -- [crashr/gppm](https://github.com/crashr/gppm) – launch llama.cpp instances utilizing NVIDIA Tesla P40 or P100 GPUs with reduced idle power consumption -- [gpustack/gguf-parser](https://github.com/gpustack/gguf-parser-go/tree/main/cmd/gguf-parser) - review/check the GGUF file and estimate the memory usage -- [Styled Lines](https://marketplace.unity.com/packages/tools/generative-ai/styled-lines-llama-cpp-model-292902) (proprietary licensed, async wrapper of inference part for game development in Unity3d with prebuild Mobile and Web platform wrappers and a model example) - -**Infrastructure:** - -- [Paddler](https://github.com/distantmagic/paddler) - Stateful load balancer custom-tailored for llama.cpp -- [GPUStack](https://github.com/gpustack/gpustack) - Manage GPU clusters for running LLMs - -**Games:** -- [Lucy's Labyrinth](https://github.com/MorganRO8/Lucys_Labyrinth) - A simple maze game where agents controlled by an AI model will try to trick you. - -## Demo - -
-Typical run using LLaMA v2 13B on M2 Ultra - -``` -$ make -j && ./llama-cli -m models/llama-13b-v2/ggml-model-q4_0.gguf -p "Building a website can be done in 10 simple steps:\nStep 1:" -n 400 -e -I llama.cpp build info: -I UNAME_S: Darwin -I UNAME_P: arm -I UNAME_M: arm64 -I CFLAGS: -I. -O3 -std=c11 -fPIC -DNDEBUG -Wall -Wextra -Wpedantic -Wcast-qual -Wdouble-promotion -Wshadow -Wstrict-prototypes -Wpointer-arith -Wmissing-prototypes -pthread -DGGML_USE_K_QUANTS -DGGML_USE_ACCELERATE -I CXXFLAGS: -I. -I./common -O3 -std=c++11 -fPIC -DNDEBUG -Wall -Wextra -Wpedantic -Wcast-qual -Wno-unused-function -Wno-multichar -pthread -DGGML_USE_K_QUANTS -I LDFLAGS: -framework Accelerate -I CC: Apple clang version 14.0.3 (clang-1403.0.22.14.1) -I CXX: Apple clang version 14.0.3 (clang-1403.0.22.14.1) - -make: Nothing to be done for `default'. -main: build = 1041 (cf658ad) -main: seed = 1692823051 -llama_model_loader: loaded meta data with 16 key-value pairs and 363 tensors from models/llama-13b-v2/ggml-model-q4_0.gguf (version GGUF V1 (latest)) -llama_model_loader: - type f32: 81 tensors -llama_model_loader: - type q4_0: 281 tensors -llama_model_loader: - type q6_K: 1 tensors -llm_load_print_meta: format = GGUF V1 (latest) -llm_load_print_meta: arch = llama -llm_load_print_meta: vocab type = SPM -llm_load_print_meta: n_vocab = 32000 -llm_load_print_meta: n_merges = 0 -llm_load_print_meta: n_ctx_train = 4096 -llm_load_print_meta: n_ctx = 512 -llm_load_print_meta: n_embd = 5120 -llm_load_print_meta: n_head = 40 -llm_load_print_meta: n_head_kv = 40 -llm_load_print_meta: n_layer = 40 -llm_load_print_meta: n_rot = 128 -llm_load_print_meta: n_gqa = 1 -llm_load_print_meta: f_norm_eps = 1.0e-05 -llm_load_print_meta: f_norm_rms_eps = 1.0e-05 -llm_load_print_meta: n_ff = 13824 -llm_load_print_meta: freq_base = 10000.0 -llm_load_print_meta: freq_scale = 1 -llm_load_print_meta: model type = 13B -llm_load_print_meta: model ftype = mostly Q4_0 -llm_load_print_meta: model size = 13.02 B -llm_load_print_meta: general.name = LLaMA v2 -llm_load_print_meta: BOS token = 1 '' -llm_load_print_meta: EOS token = 2 '' -llm_load_print_meta: UNK token = 0 '' -llm_load_print_meta: LF token = 13 '<0x0A>' -llm_load_tensors: ggml ctx size = 0.11 MB -llm_load_tensors: mem required = 7024.01 MB (+ 400.00 MB per state) -................................................................................................... -llama_new_context_with_model: kv self size = 400.00 MB -llama_new_context_with_model: compute buffer total size = 75.41 MB - -system_info: n_threads = 16 / 24 | AVX = 0 | AVX2 = 0 | AVX512 = 0 | AVX512_VBMI = 0 | AVX512_VNNI = 0 | FMA = 0 | NEON = 1 | ARM_FMA = 1 | F16C = 0 | FP16_VA = 1 | WASM_SIMD = 0 | BLAS = 1 | SSE3 = 0 | VSX = 0 | -sampling: repeat_last_n = 64, repeat_penalty = 1.100000, presence_penalty = 0.000000, frequency_penalty = 0.000000, top_k = 40, tfs_z = 1.000000, top_p = 0.950000, typical_p = 1.000000, temp = 0.800000, mirostat = 0, mirostat_lr = 0.100000, mirostat_ent = 5.000000 -generate: n_ctx = 512, n_batch = 512, n_predict = 400, n_keep = 0 - - - Building a website can be done in 10 simple steps: -Step 1: Find the right website platform. -Step 2: Choose your domain name and hosting plan. -Step 3: Design your website layout. -Step 4: Write your website content and add images. -Step 5: Install security features to protect your site from hackers or spammers -Step 6: Test your website on multiple browsers, mobile devices, operating systems etc… -Step 7: Test it again with people who are not related to you personally – friends or family members will work just fine! -Step 8: Start marketing and promoting the website via social media channels or paid ads -Step 9: Analyze how many visitors have come to your site so far, what type of people visit more often than others (e.g., men vs women) etc… -Step 10: Continue to improve upon all aspects mentioned above by following trends in web design and staying up-to-date on new technologies that can enhance user experience even further! -How does a Website Work? -A website works by having pages, which are made of HTML code. This code tells your computer how to display the content on each page you visit – whether it’s an image or text file (like PDFs). In order for someone else’s browser not only be able but also want those same results when accessing any given URL; some additional steps need taken by way of programming scripts that will add functionality such as making links clickable! -The most common type is called static HTML pages because they remain unchanged over time unless modified manually (either through editing files directly or using an interface such as WordPress). They are usually served up via HTTP protocols – this means anyone can access them without having any special privileges like being part of a group who is allowed into restricted areas online; however, there may still exist some limitations depending upon where one lives geographically speaking. -How to -llama_print_timings: load time = 576.45 ms -llama_print_timings: sample time = 283.10 ms / 400 runs ( 0.71 ms per token, 1412.91 tokens per second) -llama_print_timings: prompt eval time = 599.83 ms / 19 tokens ( 31.57 ms per token, 31.68 tokens per second) -llama_print_timings: eval time = 24513.59 ms / 399 runs ( 61.44 ms per token, 16.28 tokens per second) -llama_print_timings: total time = 25431.49 ms -``` - -
- -
-Demo of running both LLaMA-7B and whisper.cpp on a single M1 Pro MacBook - -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: - -https://user-images.githubusercontent.com/1991296/224442907-7693d4be-acaa-4e01-8b4f-add84093ffff.mp4 - -
- -## Usage - -Here are the end-to-end binary build and model conversion steps for most supported models. - -### Basic usage - -Firstly, you need to get the binary. There are different methods that you can follow: -- Method 1: Clone this repository and build locally, see [how to build](./docs/build.md) -- Method 2: If you are using MacOS or Linux, you can install llama.cpp via [brew, flox or nix](./docs/install.md) -- Method 3: Use a Docker image, see [documentation for Docker](./docs/docker.md) -- Method 4: Download pre-built binary from [releases](https://github.com/ggerganov/llama.cpp/releases) - -You can run a basic completion using this command: - -```bash -llama-cli -m your_model.gguf -p "I believe the meaning of life is" -n 128 - -# Output: -# I believe the meaning of life is to find your own truth and to live in accordance with it. For me, this means being true to myself and following my passions, even if they don't align with societal expectations. I think that's what I love about yoga – it's not just a physical practice, but a spiritual one too. It's about connecting with yourself, listening to your inner voice, and honoring your own unique journey. -``` - -See [this page](./examples/main/README.md) for a full list of parameters. - -### Conversation mode - -If you want a more ChatGPT-like experience, you can run in conversation mode by passing `-cnv` as a parameter: - -```bash -llama-cli -m your_model.gguf -p "You are a helpful assistant" -cnv - -# Output: -# > hi, who are you? -# Hi there! I'm your helpful assistant! I'm an AI-powered chatbot designed to assist and provide information to users like you. I'm here to help answer your questions, provide guidance, and offer support on a wide range of topics. I'm a friendly and knowledgeable AI, and I'm always happy to help with anything you need. What's on your mind, and how can I assist you today? -# -# > what is 1+1? -# Easy peasy! The answer to 1+1 is... 2! -``` - -By default, the chat template will be taken from the input model. If you want to use another chat template, pass `--chat-template NAME` as a parameter. See the list of [supported templates](https://github.com/ggerganov/llama.cpp/wiki/Templates-supported-by-llama_chat_apply_template) - -```bash -./llama-cli -m your_model.gguf -p "You are a helpful assistant" -cnv --chat-template chatml -``` - -You can also use your own template via in-prefix, in-suffix and reverse-prompt parameters: - -```bash -./llama-cli -m your_model.gguf -p "You are a helpful assistant" -cnv --in-prefix 'User: ' --reverse-prompt 'User:' -``` - -### Web server - -[llama.cpp web server](./examples/server/README.md) is a lightweight [OpenAI API](https://github.com/openai/openai-openapi) compatible HTTP server that can be used to serve local models and easily connect them to existing clients. - -Example usage: - -```bash -./llama-server -m your_model.gguf --port 8080 - -# Basic web UI can be accessed via browser: http://localhost:8080 -# Chat completion endpoint: http://localhost:8080/v1/chat/completions -``` - -### Interactive mode - -> [!NOTE] -> If you prefer basic usage, please consider using conversation mode instead of interactive mode - -In this mode, you can always interrupt generation by pressing Ctrl+C and entering 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 that makes LLaMA emulate a chat between multiple users, say Alice and Bob, and pass `-r "Alice:"`. - -Here is an example of a few-shot interaction, invoked with the command - -```bash -# default arguments using a 7B model -./examples/chat.sh - -# advanced chat with a 13B model -./examples/chat-13B.sh - -# custom arguments using a 13B model -./llama-cli -m ./models/13B/ggml-model-q4_0.gguf -n 256 --repeat_penalty 1.0 --color -i -r "User:" -f prompts/chat-with-bob.txt -``` - -Note the use of `--color` to distinguish between user input and generated text. Other parameters are explained in more detail in the [README](examples/main/README.md) for the `llama-cli` example program. - -![image](https://user-images.githubusercontent.com/1991296/224575029-2af3c7dc-5a65-4f64-a6bb-517a532aea38.png) - -### Persistent Interaction - -The prompt, user inputs, and model generations can be saved and resumed across calls to `./llama-cli` by leveraging `--prompt-cache` and `--prompt-cache-all`. The `./examples/chat-persistent.sh` script demonstrates this with support for long-running, resumable chat sessions. To use this example, you must provide a file to cache the initial chat prompt and a directory to save the chat session, and may optionally provide the same variables as `chat-13B.sh`. The same prompt cache can be reused for new chat sessions. Note that both prompt cache and chat directory are tied to the initial prompt (`PROMPT_TEMPLATE`) and the model file. - -```bash -# Start a new chat -PROMPT_CACHE_FILE=chat.prompt.bin CHAT_SAVE_DIR=./chat/default ./examples/chat-persistent.sh - -# Resume that chat -PROMPT_CACHE_FILE=chat.prompt.bin CHAT_SAVE_DIR=./chat/default ./examples/chat-persistent.sh - -# Start a different chat with the same prompt/model -PROMPT_CACHE_FILE=chat.prompt.bin CHAT_SAVE_DIR=./chat/another ./examples/chat-persistent.sh - -# Different prompt cache for different prompt/model -PROMPT_TEMPLATE=./prompts/chat-with-bob.txt PROMPT_CACHE_FILE=bob.prompt.bin \ - CHAT_SAVE_DIR=./chat/bob ./examples/chat-persistent.sh -``` - -### Constrained output with grammars - -`llama.cpp` supports grammars to constrain model output. For example, you can force the model to output JSON only: - -```bash -./llama-cli -m ./models/13B/ggml-model-q4_0.gguf -n 256 --grammar-file grammars/json.gbnf -p 'Request: schedule a call at 8pm; Command:' -``` - -The `grammars/` folder contains a handful of sample grammars. To write your own, check out the [GBNF Guide](./grammars/README.md). - -For authoring more complex JSON grammars, you can also check out https://grammar.intrinsiclabs.ai/, a browser app that lets you write TypeScript interfaces which it compiles to GBNF grammars that you can save for local use. Note that the app is built and maintained by members of the community, please file any issues or FRs on [its repo](http://github.com/intrinsiclabsai/gbnfgen) and not this one. - -## Build - -Please refer to [Build llama.cpp locally](./docs/build.md) - -## Supported backends - -| Backend | Target devices | -| --- | --- | -| [Metal](./docs/build.md#metal-build) | Apple Silicon | -| [BLAS](./docs/build.md#blas-build) | All | -| [BLIS](./docs/backend/BLIS.md) | All | -| [SYCL](./docs/backend/SYCL.md) | Intel and Nvidia GPU | -| [MUSA](./docs/build.md#musa) | Moore Threads GPU | -| [CUDA](./docs/build.md#cuda) | Nvidia GPU | -| [hipBLAS](./docs/build.md#hipblas) | AMD GPU | -| [Vulkan](./docs/build.md#vulkan) | GPU | -| [CANN](./docs/build.md#cann) | Ascend NPU | - -## Tools - -### Prepare and Quantize - -> [!NOTE] -> You can use the [GGUF-my-repo](https://huggingface.co/spaces/ggml-org/gguf-my-repo) space on Hugging Face to quantise your model weights without any setup too. It is synced from `llama.cpp` main every 6 hours. - -To obtain the official LLaMA 2 weights please see the Obtaining and using the Facebook LLaMA 2 model section. There is also a large selection of pre-quantized `gguf` models available on Hugging Face. - -Note: `convert.py` has been moved to `examples/convert_legacy_llama.py` and shouldn't be used for anything other than `Llama/Llama2/Mistral` models and their derivatives. -It does not support LLaMA 3, you can use `convert_hf_to_gguf.py` with LLaMA 3 downloaded from Hugging Face. - -To learn more about quantizing model, [read this documentation](./examples/quantize/README.md) - -### Perplexity (measuring model quality) - -You can use the `perplexity` example to measure perplexity over a given prompt (lower perplexity is better). -For more information, see [https://huggingface.co/docs/transformers/perplexity](https://huggingface.co/docs/transformers/perplexity). - -To learn more how to measure perplexity using llama.cpp, [read this documentation](./examples/perplexity/README.md) - -## 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, PRs and projects is very appreciated! -- 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 -- Read the [CONTRIBUTING.md](CONTRIBUTING.md) for more information -- 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) - -## Other documentations - -- [main (cli)](./examples/main/README.md) -- [server](./examples/server/README.md) -- [jeopardy](./examples/jeopardy/README.md) -- [GBNF grammars](./grammars/README.md) - -**Development documentations** - -- [How to build](./docs/build.md) -- [Running on Docker](./docs/docker.md) -- [Build on Android](./docs/android.md) -- [Performance troubleshooting](./docs/development/token_generation_performance_tips.md) -- [GGML tips & tricks](https://github.com/ggerganov/llama.cpp/wiki/GGML-Tips-&-Tricks) - -**Seminal papers and background on the models** - -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) +This is a distributed implementation of [llama.cpp](https://github.com/ggerganov/llama.cpp), coming soon. \ No newline at end of file diff --git a/common/arg.cpp b/common/arg.cpp index 2a85ad84..0d3ddccf 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -661,6 +661,41 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, params.n_ctx = value; } ).set_env("LLAMA_ARG_CTX_SIZE")); + add_opt(llama_arg( + {"-w", "--world", "--world-size"}, "N", + format("number of devices to use (default: %d)", params.n_world), + [](gpt_params & params, int value) { + params.n_world = value; + } + ).set_env("LLAMA_ARG_N_WORLD")); + add_opt(llama_arg( + {"-r", "--rank", "--my-rank"}, "N", + format("my rank for distributed inference (default: %d)", params.rank), + [](gpt_params & params, int value) { + params.rank = value; + } + ).set_env("LLAMA_ARG_RANK")); + add_opt(llama_arg( + {"-lw", "--layer-window", "--n-layer-window"}, "N", + format("number of layers to process in each compute (default: %d)", params.n_layer_window), + [](gpt_params & params, int value) { + params.n_layer_window = value; + } + ).set_env("LLAMA_ARG_N_LAYER_WINDOW")); + add_opt(llama_arg( + {"-mip", "--master", "--master-ip"}, "IPAddress", + format("ip address of the master node (default: %s)", params.master_ip.c_str()), + [](gpt_params & params, const std::string & value) { + params.master_ip = value; + } + ).set_env("LLAMA_ARG_MASTER_IP")); + add_opt(llama_arg( + {"-nip", "--next", "--next-node", "--next-ip", "--next-node-ip"}, "IPAddress", + format("ip address of the next node (default: %s)", params.next_node_ip.c_str()), + [](gpt_params & params, const std::string & value) { + params.next_node_ip = value; + } + ).set_env("LLAMA_ARG_NEXT_NODE_IP")); add_opt(llama_arg( {"-n", "--predict", "--n-predict"}, "N", format("number of tokens to predict (default: %d, -1 = infinity, -2 = until context filled)", params.n_predict), diff --git a/common/common.cpp b/common/common.cpp index 29df16c9..a4a68899 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -872,6 +872,8 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) { return iparams; } + llama_init_sockets(lctx, cparams.n_world, cparams.rank); + if (!params.control_vectors.empty()) { if (params.control_vector_layer_start <= 0) params.control_vector_layer_start = 1; if (params.control_vector_layer_end <= 0) params.control_vector_layer_end = llama_n_layer(model); @@ -924,28 +926,34 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) { if (params.warmup) { LOG_WRN("%s: warming up the model with an empty run - please wait ... (--no-warmup to disable)\n", __func__); + const uint32_t my_rank = cparams.rank; std::vector tmp; - llama_token bos = llama_token_bos(model); - llama_token eos = llama_token_eos(model); - // some models (e.g. T5) don't have a BOS token - if (bos != LLAMA_TOKEN_NULL) { - tmp.push_back(bos); - } - if (eos != LLAMA_TOKEN_NULL) { - tmp.push_back(eos); - } - if (tmp.empty()) { - tmp.push_back(0); - } - if (llama_model_has_encoder(model)) { - llama_encode(lctx, llama_batch_get_one(tmp.data(), tmp.size(), 0, 0)); - llama_token decoder_start_token_id = llama_model_decoder_start_token(model); - if (decoder_start_token_id == -1) { - decoder_start_token_id = bos; + if (my_rank == 0) { + llama_token bos = llama_token_bos(model); + llama_token eos = llama_token_eos(model); + // some models (e.g. T5) don't have a BOS token + if (bos != LLAMA_TOKEN_NULL) { + tmp.push_back(bos); + } + if (eos != LLAMA_TOKEN_NULL) { + tmp.push_back(eos); + } + if (tmp.empty()) { + tmp.push_back(0); + } + + if (llama_model_has_encoder(model)) { + throw std::runtime_error("this model is currently not supported"); + + llama_encode(lctx, llama_batch_get_one(tmp.data(), tmp.size(), 0, 0)); + llama_token decoder_start_token_id = llama_model_decoder_start_token(model); + if (decoder_start_token_id == -1) { + decoder_start_token_id = bos; + } + tmp.clear(); + tmp.push_back(decoder_start_token_id); } - tmp.clear(); - tmp.push_back(decoder_start_token_id); } if (llama_model_has_decoder(model)) { llama_decode(lctx, llama_batch_get_one(tmp.data(), std::min(tmp.size(), (size_t) params.n_batch), 0, 0)); @@ -976,6 +984,9 @@ struct llama_model_params llama_model_params_from_gpt_params(const gpt_params & if (params.n_gpu_layers != -1) { mparams.n_gpu_layers = params.n_gpu_layers; } + mparams.n_world = params.n_world; + mparams.rank = params.rank; + mparams.n_layer_window = params.n_layer_window; mparams.rpc_servers = params.rpc_servers.c_str(); mparams.main_gpu = params.main_gpu; mparams.split_mode = params.split_mode; @@ -1025,6 +1036,22 @@ static ggml_type kv_cache_type_from_str(const std::string & s) { struct llama_context_params llama_context_params_from_gpt_params(const gpt_params & params) { auto cparams = llama_context_default_params(); + cparams.n_world = params.n_world; + cparams.rank = params.rank; + cparams.n_layer_window = params.n_layer_window; + + if (cparams.master_ip != nullptr) { + delete[] cparams.master_ip; + } + cparams.master_ip = new char[params.master_ip.length() + 1]; + std::strcpy(cparams.master_ip, params.master_ip.c_str()); + + if (cparams.next_node_ip != nullptr) { + delete[] cparams.next_node_ip; + } + cparams.next_node_ip = new char[params.next_node_ip.length() + 1]; + std::strcpy(cparams.next_node_ip, params.next_node_ip.c_str()); + cparams.n_ctx = params.n_ctx; cparams.n_seq_max = params.n_parallel; cparams.n_batch = params.n_batch; diff --git a/common/common.h b/common/common.h index 8b84cf9a..21540dac 100644 --- a/common/common.h +++ b/common/common.h @@ -142,6 +142,11 @@ struct gpt_sampler_params { }; struct gpt_params { + int32_t n_world = 1; // number of devices to use + int32_t rank = 0; // my rank for distributed inference + int32_t n_layer_window = 32; // number of layers to process in each compute + std::string master_ip = "localhost"; // ip address of the master node + std::string next_node_ip = "localhost"; // ip address of my next node int32_t n_predict = -1; // new tokens to predict int32_t n_ctx = 0; // context size int32_t n_batch = 2048; // logical batch size for prompt processing (must be >=32 to use BLAS) @@ -229,8 +234,7 @@ struct gpt_params { bool multiple_choice = false; // compute TruthfulQA score over random tasks from datafile supplied in prompt size_t multiple_choice_tasks = 0; // number of tasks to use when computing the TruthfulQA score. If 0, all tasks will be computed - bool kl_divergence = false; // compute KL divergence - + bool kl_divergence = false; // compute KL divergence bool usage = false; // print usage bool use_color = false; // use color to distinguish generations and inputs bool special = false; // enable special token output diff --git a/examples/main/main.cpp b/examples/main/main.cpp index 6bbb1e13..d85a7bb0 100644 --- a/examples/main/main.cpp +++ b/examples/main/main.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) #include @@ -141,6 +142,9 @@ int main(int argc, char ** argv) { if (!gpt_params_parse(argc, argv, params, LLAMA_EXAMPLE_MAIN, print_usage)) { return 1; } + const uint32_t n_world = params.n_world; + const uint32_t my_rank = params.rank; + GGML_ASSERT(!(n_world == 1 && my_rank > 0)); gpt_init(); @@ -151,22 +155,6 @@ int main(int argc, char ** argv) { console::init(params.simple_io, params.use_color); atexit([]() { console::cleanup(); }); - if (params.logits_all) { - LOG_ERR("************\n"); - LOG_ERR("%s: please use the 'perplexity' tool for perplexity calculations\n", __func__); - LOG_ERR("************\n\n"); - - return 0; - } - - if (params.embedding) { - LOG_ERR("************\n"); - LOG_ERR("%s: please use the 'embedding' tool for embedding calculations\n", __func__); - LOG_ERR("************\n\n"); - - return 0; - } - if (params.n_ctx != 0 && params.n_ctx < 8) { LOG_WRN("%s: warning: minimum context size is 8, using minimum size.\n", __func__); params.n_ctx = 8; @@ -290,7 +278,7 @@ int main(int argc, char ** argv) { std::vector embd_inp; - { + if (my_rank == 0) { auto prompt = (params.conversation && params.enable_chat_template && !params.prompt.empty()) ? chat_add_and_format(model, chat_msgs, "system", params.prompt) // format the system prompt in conversation mode : params.prompt; @@ -304,23 +292,23 @@ int main(int argc, char ** argv) { LOG_DBG("prompt: \"%s\"\n", prompt.c_str()); LOG_DBG("tokens: %s\n", string_from(ctx, embd_inp).c_str()); - } - // Should not run without any tokens - if (embd_inp.empty()) { - if (add_bos) { - embd_inp.push_back(llama_token_bos(model)); - LOG_WRN("embd_inp was considered empty and bos was added: %s\n", string_from(ctx, embd_inp).c_str()); - } else { - LOG_ERR("input is empty\n"); - return -1; + // should not run without any tokens + if (embd_inp.empty()) { + if (add_bos) { + embd_inp.push_back(llama_token_bos(model)); + LOG_WRN("embd_inp was considered empty and bos was added: %s\n", string_from(ctx, embd_inp).c_str()); + } else { + LOG_ERR("input is empty\n"); + return -1; + } } - } - // Tokenize negative prompt - if ((int) embd_inp.size() > n_ctx - 4) { - LOG_ERR("%s: prompt is too long (%d tokens, max %d)\n", __func__, (int) embd_inp.size(), n_ctx - 4); - return 1; + // tokenize negative prompt + if ((int) embd_inp.size() > n_ctx - 4) { + LOG_ERR("%s: prompt is too long (%d tokens, max %d)\n", __func__, (int) embd_inp.size(), n_ctx - 4); + return 1; + } } // debug message about similarity of saved session, if applicable @@ -448,18 +436,18 @@ int main(int argc, char ** argv) { } } - smpl = gpt_sampler_init(model, sparams); - if (!smpl) { - LOG_ERR("%s: failed to initialize sampling subsystem\n", __func__); - return 1; + if (my_rank == 0) { + smpl = gpt_sampler_init(model, sparams); + if (!smpl) { + LOG_ERR("%s: failed to initialize sampling subsystem\n", __func__); + return 1; + } + LOG_INF("sampler seed: %u\n", gpt_sampler_get_seed(smpl)); + LOG_INF("sampler params: \n%s\n", sparams.print().c_str()); + LOG_INF("sampler chain: %s\n", gpt_sampler_print(smpl).c_str()); + LOG_INF("generate: n_ctx = %d, n_batch = %d, n_predict = %d, n_keep = %d\n", n_ctx, params.n_batch, params.n_predict, params.n_keep); } - LOG_INF("sampler seed: %u\n", gpt_sampler_get_seed(smpl)); - LOG_INF("sampler params: \n%s\n", sparams.print().c_str()); - LOG_INF("sampler chain: %s\n", gpt_sampler_print(smpl).c_str()); - - LOG_INF("generate: n_ctx = %d, n_batch = %d, n_predict = %d, n_keep = %d\n", n_ctx, params.n_batch, params.n_predict, params.n_keep); - // group-attention state // number of grouped KV tokens so far (used only if params.grp_attn_n > 1) int ga_i = 0; @@ -487,9 +475,7 @@ int main(int argc, char ** argv) { " - If you want to submit another line, end your input with '\\'.\n"; } LOG_INF("== Running in interactive mode. ==\n"); -#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) || defined (_WIN32) - LOG_INF( " - Press Ctrl+C to interject at any time.\n"); -#endif + LOG_INF( " - Enter quit or exit to quit chat.\n"); LOG_INF( "%s\n", control_message); is_interacting = params.interactive_first; @@ -525,6 +511,8 @@ int main(int argc, char ** argv) { } if (llama_model_has_encoder(model)) { + throw std::runtime_error("this model is currently not supported"); + int enc_input_size = embd_inp.size(); llama_token * enc_input_buf = embd_inp.data(); @@ -542,9 +530,16 @@ int main(int argc, char ** argv) { embd_inp.push_back(decoder_start_token_id); } + char * stop_signal = nullptr; + std::thread signal_thread; + + if (my_rank != 0) { + signal_thread = std::thread(llama_free_sockets, ctx, &stop_signal); + } + while ((n_remain != 0 && !is_antiprompt) || params.interactive) { // predict - if (!embd.empty()) { + if (!embd.empty() || my_rank != 0) { // Note: (n_ctx - 4) here is to match the logic for commandline prompt handling via // --prompt or --file which uses the same value. int max_embd_size = n_ctx - 4; @@ -640,25 +635,22 @@ int main(int argc, char ** argv) { } } - for (int i = 0; i < (int) embd.size(); i += params.n_batch) { - int n_eval = (int) embd.size() - i; - if (n_eval > params.n_batch) { - n_eval = params.n_batch; + if (my_rank == 0) { + for (int i = 0; i < (int) embd.size(); i += params.n_batch) { + int n_eval = (int) embd.size() - i; + if (n_eval > params.n_batch) { + n_eval = params.n_batch; + } + if (llama_decode(ctx, llama_batch_get_one(&embd[i], n_eval, n_past, 0)) != 0) { + LOG_ERR("%s : failed to eval\n", __func__); + return 1; + } + n_past += n_eval; } - - LOG_DBG("eval: %s\n", string_from(ctx, embd).c_str()); - - if (llama_decode(ctx, llama_batch_get_one(&embd[i], n_eval, n_past, 0))) { - LOG_ERR("%s : failed to eval\n", __func__); - return 1; - } - - n_past += n_eval; - - LOG_DBG("n_past = %d\n", n_past); - // Display total tokens alongside total time - if (params.n_print > 0 && n_past % params.n_print == 0) { - LOG_DBG("\n\033[31mTokens consumed so far = %d / %d \033[0m\n", n_past, n_ctx); + } else { + llama_decode(ctx, llama_batch_get_one(embd.data(), 0, 0, 0)); + if (stop_signal != nullptr && std::strcmp(stop_signal, "STOP") == 0) { + break; } } @@ -670,70 +662,70 @@ int main(int argc, char ** argv) { embd.clear(); - if ((int) embd_inp.size() <= n_consumed && !is_interacting) { - // optionally save the session on first sample (for faster prompt loading next time) - if (!path_session.empty() && need_to_save_session && !params.prompt_cache_ro) { - need_to_save_session = false; - llama_state_save_file(ctx, path_session.c_str(), session_tokens.data(), session_tokens.size()); + if (my_rank == 0) { + if ((int) embd_inp.size() <= n_consumed && !is_interacting) { + // optionally save the session on first sample (for faster prompt loading next time) + if (!path_session.empty() && need_to_save_session && !params.prompt_cache_ro) { + need_to_save_session = false; + llama_state_save_file(ctx, path_session.c_str(), session_tokens.data(), session_tokens.size()); + LOG_DBG("saved session to %s\n", path_session.c_str()); + } - LOG_DBG("saved session to %s\n", path_session.c_str()); - } + const llama_token id = gpt_sampler_sample(smpl, ctx, -1); - const llama_token id = gpt_sampler_sample(smpl, ctx, -1); + gpt_sampler_accept(smpl, id, /* accept_grammar= */ true); - gpt_sampler_accept(smpl, id, /* accept_grammar= */ true); + embd.push_back(id); - // LOG_DBG("last: %s\n", string_from(ctx, smpl->prev.to_vector()).c_str()); + // echo this to console + input_echo = true; - embd.push_back(id); + // decrement remaining sampling budget + --n_remain; - // echo this to console - input_echo = true; + LOG_DBG("n_remain: %d\n", n_remain); + } else { + // some user input remains from prompt or interaction, forward it to processing + while ((int) embd_inp.size() > n_consumed) { + embd.push_back(embd_inp[n_consumed]); - // decrement remaining sampling budget - --n_remain; + // push the prompt in the sampling context in order to apply repetition penalties later + // for the prompt, we don't apply grammar rules + gpt_sampler_accept(smpl, embd_inp[n_consumed], /* accept_grammar= */ false); - LOG_DBG("n_remain: %d\n", n_remain); - } else { - // some user input remains from prompt or interaction, forward it to processing - LOG_DBG("embd_inp.size(): %d, n_consumed: %d\n", (int) embd_inp.size(), n_consumed); - while ((int) embd_inp.size() > n_consumed) { - embd.push_back(embd_inp[n_consumed]); - - // push the prompt in the sampling context in order to apply repetition penalties later - // for the prompt, we don't apply grammar rules - gpt_sampler_accept(smpl, embd_inp[n_consumed], /* accept_grammar= */ false); - - ++n_consumed; - if ((int) embd.size() >= params.n_batch) { - break; + ++n_consumed; + if ((int) embd.size() >= params.n_batch) { + break; + } } } } // display text - if (input_echo && display) { - for (auto id : embd) { - const std::string token_str = llama_token_to_piece(ctx, id, params.special); + if (my_rank == 0) { + if (input_echo && display) { + for (auto id : embd) { + const std::string token_str = llama_token_to_piece(ctx, id, params.special); - // Console/Stream Output - LOG("%s", token_str.c_str()); + // Console/Stream Output + LOG("%s", token_str.c_str()); - // Record Displayed Tokens To Log - // Note: Generated tokens are created one by one hence this check - if (embd.size() > 1) { - // Incoming Requested Tokens - input_tokens.push_back(id); - } else { - // Outgoing Generated Tokens - output_tokens.push_back(id); - output_ss << token_str; + // Record Displayed Tokens To Log + // Note: Generated tokens are created one by one hence this check + if (embd.size() > 1) { + // Incoming Requested Tokens + input_tokens.push_back(id); + } else { + // Outgoing Generated Tokens + output_tokens.push_back(id); + output_ss << token_str; + } } } } // reset color to default if there is no pending user input - if (input_echo && (int) embd_inp.size() == n_consumed) { + if (my_rank == 0 && input_echo && (int) embd_inp.size() == n_consumed) { console::set_display(console::reset); display = true; } @@ -782,7 +774,7 @@ int main(int argc, char ** argv) { } // deal with end of generation tokens in interactive mode - if (llama_token_is_eog(model, gpt_sampler_last(smpl))) { + if (my_rank == 0 && llama_token_is_eog(model, gpt_sampler_last(smpl))) { LOG_DBG("found an EOG token\n"); if (params.interactive) { @@ -840,6 +832,10 @@ int main(int argc, char ** argv) { console::set_display(console::reset); display = true; + if (buffer == "quit\n" || buffer == "exit\n") { + break; + } + // Add tokens to embd only if the input buffer is non-empty // Entering a empty line lets the user pass control back if (buffer.length() > 1) { @@ -924,19 +920,20 @@ int main(int argc, char ** argv) { llama_state_save_file(ctx, path_session.c_str(), session_tokens.data(), session_tokens.size()); } - LOG("\n\n"); - gpt_perf_print(ctx, smpl); - write_logfile(ctx, params, model, input_tokens, output_ss.str(), output_tokens); - - gpt_sampler_free(smpl); - + if (my_rank == 0) { + LOG("\n\n"); + gpt_perf_print(ctx, smpl); + write_logfile(ctx, params, model, input_tokens, output_ss.str(), output_tokens); + gpt_sampler_free(smpl); + llama_free_sockets(ctx, &stop_signal); + } + if (my_rank != 0 && signal_thread.joinable()) { + signal_thread.join(); + } llama_free(ctx); llama_free_model(model); - llama_backend_free(); - ggml_threadpool_free(threadpool); ggml_threadpool_free(threadpool_batch); - return 0; -} +} \ No newline at end of file diff --git a/include/llama.h b/include/llama.h index 7cae1bbe..b0965123 100644 --- a/include/llama.h +++ b/include/llama.h @@ -276,6 +276,9 @@ extern "C" { }; struct llama_model_params { + uint32_t n_world; // number of nodes + uint32_t rank; // my node rank + uint32_t n_layer_window; // number of layers to kept each time int32_t n_gpu_layers; // number of layers to store in VRAM enum llama_split_mode split_mode; // how to split the model across multiple GPUs @@ -312,12 +315,17 @@ extern "C" { // NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations // https://github.com/ggerganov/llama.cpp/pull/7544 struct llama_context_params { - uint32_t n_ctx; // text context, 0 = from model - uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode - uint32_t n_ubatch; // physical maximum batch size - uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models) - int32_t n_threads; // number of threads to use for generation - int32_t n_threads_batch; // number of threads to use for batch processing + uint32_t n_world; // world size + uint32_t rank; // my rank + uint32_t n_layer_window; // number of layers to process in each compute + char * master_ip; // ip address of the master node + char * next_node_ip; // ip address of the next node + uint32_t n_ctx; // text context, 0 = from model + uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode + uint32_t n_ubatch; // physical maximum batch size + uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models) + int32_t n_threads; // number of threads to use for generation + int32_t n_threads_batch; // number of threads to use for batch processing enum llama_rope_scaling_type rope_scaling_type; // RoPE scaling type, from `enum llama_rope_scaling_type` enum llama_pooling_type pooling_type; // whether to pool (sum) embedding results by sequence id @@ -418,6 +426,9 @@ extern "C" { LLAMA_API void llama_free_model(struct llama_model * model); + LLAMA_API void llama_init_sockets(struct llama_context * ctx, uint32_t n_world, uint32_t my_rank); + LLAMA_API void llama_free_sockets(struct llama_context * ctx, char ** msg); + // TODO: rename to llama_init_from_model LLAMA_API struct llama_context * llama_new_context_with_model( struct llama_model * model, diff --git a/include/zmq.hpp b/include/zmq.hpp new file mode 100644 index 00000000..a434eb61 --- /dev/null +++ b/include/zmq.hpp @@ -0,0 +1,2796 @@ +/* + Copyright (c) 2016-2017 ZeroMQ community + Copyright (c) 2009-2011 250bpm s.r.o. + Copyright (c) 2011 Botond Ballo + Copyright (c) 2007-2009 iMatix Corporation + + 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. +*/ + +#ifndef __ZMQ_HPP_INCLUDED__ +#define __ZMQ_HPP_INCLUDED__ + +#ifdef _WIN32 +#ifndef NOMINMAX +#define NOMINMAX +#endif +#endif + +// included here for _HAS_CXX* macros +#include + +#if defined(_MSVC_LANG) +#define CPPZMQ_LANG _MSVC_LANG +#else +#define CPPZMQ_LANG __cplusplus +#endif +// overwrite if specific language macros indicate higher version +#if defined(_HAS_CXX14) && _HAS_CXX14 && CPPZMQ_LANG < 201402L +#undef CPPZMQ_LANG +#define CPPZMQ_LANG 201402L +#endif +#if defined(_HAS_CXX17) && _HAS_CXX17 && CPPZMQ_LANG < 201703L +#undef CPPZMQ_LANG +#define CPPZMQ_LANG 201703L +#endif + +// macros defined if has a specific standard or greater +#if CPPZMQ_LANG >= 201103L || (defined(_MSC_VER) && _MSC_VER >= 1900) +#define ZMQ_CPP11 +#endif +#if CPPZMQ_LANG >= 201402L +#define ZMQ_CPP14 +#endif +#if CPPZMQ_LANG >= 201703L +#define ZMQ_CPP17 +#endif + +#if defined(ZMQ_CPP14) && !defined(_MSC_VER) +#define ZMQ_DEPRECATED(msg) [[deprecated(msg)]] +#elif defined(_MSC_VER) +#define ZMQ_DEPRECATED(msg) __declspec(deprecated(msg)) +#elif defined(__GNUC__) +#define ZMQ_DEPRECATED(msg) __attribute__((deprecated(msg))) +#else +#define ZMQ_DEPRECATED(msg) +#endif + +#if defined(ZMQ_CPP17) +#define ZMQ_NODISCARD [[nodiscard]] +#else +#define ZMQ_NODISCARD +#endif + +#if defined(ZMQ_CPP11) +#define ZMQ_NOTHROW noexcept +#define ZMQ_EXPLICIT explicit +#define ZMQ_OVERRIDE override +#define ZMQ_NULLPTR nullptr +#define ZMQ_CONSTEXPR_FN constexpr +#define ZMQ_CONSTEXPR_VAR constexpr +#define ZMQ_CPP11_DEPRECATED(msg) ZMQ_DEPRECATED(msg) +#else +#define ZMQ_NOTHROW throw() +#define ZMQ_EXPLICIT +#define ZMQ_OVERRIDE +#define ZMQ_NULLPTR 0 +#define ZMQ_CONSTEXPR_FN +#define ZMQ_CONSTEXPR_VAR const +#define ZMQ_CPP11_DEPRECATED(msg) +#endif +#if defined(ZMQ_CPP14) && (!defined(_MSC_VER) || _MSC_VER > 1900) && (!defined(__GNUC__) || __GNUC__ > 5 || (__GNUC__ == 5 && __GNUC_MINOR__ > 3)) +#define ZMQ_EXTENDED_CONSTEXPR +#endif +#if defined(ZMQ_CPP17) +#define ZMQ_INLINE_VAR inline +#define ZMQ_CONSTEXPR_IF constexpr +#else +#define ZMQ_INLINE_VAR +#define ZMQ_CONSTEXPR_IF +#endif + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#ifdef ZMQ_CPP11 +#include +#include +#include +#include +#endif + +#if defined(__has_include) && defined(ZMQ_CPP17) +#define CPPZMQ_HAS_INCLUDE_CPP17(X) __has_include(X) +#else +#define CPPZMQ_HAS_INCLUDE_CPP17(X) 0 +#endif + +#if CPPZMQ_HAS_INCLUDE_CPP17() && !defined(CPPZMQ_HAS_OPTIONAL) +#define CPPZMQ_HAS_OPTIONAL 1 +#endif +#ifndef CPPZMQ_HAS_OPTIONAL +#define CPPZMQ_HAS_OPTIONAL 0 +#elif CPPZMQ_HAS_OPTIONAL +#include +#endif + +#if CPPZMQ_HAS_INCLUDE_CPP17() && !defined(CPPZMQ_HAS_STRING_VIEW) +#define CPPZMQ_HAS_STRING_VIEW 1 +#endif +#ifndef CPPZMQ_HAS_STRING_VIEW +#define CPPZMQ_HAS_STRING_VIEW 0 +#elif CPPZMQ_HAS_STRING_VIEW +#include +#endif + +/* Version macros for compile-time API version detection */ +#define CPPZMQ_VERSION_MAJOR 4 +#define CPPZMQ_VERSION_MINOR 10 +#define CPPZMQ_VERSION_PATCH 0 + +#define CPPZMQ_VERSION \ + ZMQ_MAKE_VERSION(CPPZMQ_VERSION_MAJOR, CPPZMQ_VERSION_MINOR, \ + CPPZMQ_VERSION_PATCH) + +// Detect whether the compiler supports C++11 rvalue references. +#if (defined(__GNUC__) && (__GNUC__ > 4 || (__GNUC__ == 4 && __GNUC_MINOR__ > 2)) \ + && defined(__GXX_EXPERIMENTAL_CXX0X__)) +#define ZMQ_HAS_RVALUE_REFS +#define ZMQ_DELETED_FUNCTION = delete +#elif defined(__clang__) +#if __has_feature(cxx_rvalue_references) +#define ZMQ_HAS_RVALUE_REFS +#endif + +#if __has_feature(cxx_deleted_functions) +#define ZMQ_DELETED_FUNCTION = delete +#else +#define ZMQ_DELETED_FUNCTION +#endif +#elif defined(_MSC_VER) && (_MSC_VER >= 1900) +#define ZMQ_HAS_RVALUE_REFS +#define ZMQ_DELETED_FUNCTION = delete +#elif defined(_MSC_VER) && (_MSC_VER >= 1600) +#define ZMQ_HAS_RVALUE_REFS +#define ZMQ_DELETED_FUNCTION +#else +#define ZMQ_DELETED_FUNCTION +#endif + +#if defined(ZMQ_CPP11) && !defined(__llvm__) && !defined(__INTEL_COMPILER) \ + && defined(__GNUC__) && __GNUC__ < 5 +#define ZMQ_CPP11_PARTIAL +#elif defined(__GLIBCXX__) && __GLIBCXX__ < 20160805 +//the date here is the last date of gcc 4.9.4, which +// effectively means libstdc++ from gcc 5.5 and higher won't trigger this branch +#define ZMQ_CPP11_PARTIAL +#endif + +#ifdef ZMQ_CPP11 +#ifdef ZMQ_CPP11_PARTIAL +#define ZMQ_IS_TRIVIALLY_COPYABLE(T) __has_trivial_copy(T) +#else +#include +#define ZMQ_IS_TRIVIALLY_COPYABLE(T) std::is_trivially_copyable::value +#endif +#endif + +#if ZMQ_VERSION >= ZMQ_MAKE_VERSION(3, 3, 0) +#define ZMQ_NEW_MONITOR_EVENT_LAYOUT +#endif + +#if ZMQ_VERSION >= ZMQ_MAKE_VERSION(4, 1, 0) +#define ZMQ_HAS_PROXY_STEERABLE +/* Socket event data */ +typedef struct +{ + uint16_t event; // id of the event as bitfield + int32_t value; // value is either error code, fd or reconnect interval +} zmq_event_t; +#endif + +// Avoid using deprecated message receive function when possible +#if ZMQ_VERSION < ZMQ_MAKE_VERSION(3, 2, 0) +#define zmq_msg_recv(msg, socket, flags) zmq_recvmsg(socket, msg, flags) +#endif + + +// In order to prevent unused variable warnings when building in non-debug +// mode use this macro to make assertions. +#ifndef NDEBUG +#define ZMQ_ASSERT(expression) assert(expression) +#else +#define ZMQ_ASSERT(expression) (void) (expression) +#endif + +namespace zmq +{ +#ifdef ZMQ_CPP11 +namespace detail +{ +namespace ranges +{ +using std::begin; +using std::end; +template auto begin(T &&r) -> decltype(begin(std::forward(r))) +{ + return begin(std::forward(r)); +} +template auto end(T &&r) -> decltype(end(std::forward(r))) +{ + return end(std::forward(r)); +} +} // namespace ranges + +template using void_t = void; + +template +using iter_value_t = typename std::iterator_traits::value_type; + +template +using range_iter_t = decltype( + ranges::begin(std::declval::type &>())); + +template using range_value_t = iter_value_t>; + +template struct is_range : std::false_type +{ +}; + +template +struct is_range< + T, + void_t::type &>()) + == ranges::end(std::declval::type &>()))>> + : std::true_type +{ +}; + +} // namespace detail +#endif + +typedef zmq_free_fn free_fn; +typedef zmq_pollitem_t pollitem_t; + +// duplicate definition from libzmq 4.3.3 +#if defined _WIN32 +#if defined _WIN64 +typedef unsigned __int64 fd_t; +#else +typedef unsigned int fd_t; +#endif +#else +typedef int fd_t; +#endif + +class error_t : public std::exception +{ + public: + error_t() ZMQ_NOTHROW : errnum(zmq_errno()) {} + explicit error_t(int err) ZMQ_NOTHROW : errnum(err) {} + virtual const char *what() const ZMQ_NOTHROW ZMQ_OVERRIDE + { + return zmq_strerror(errnum); + } + int num() const ZMQ_NOTHROW { return errnum; } + + private: + int errnum; +}; + +namespace detail { +inline int poll(zmq_pollitem_t *items_, size_t nitems_, long timeout_) +{ + int rc = zmq_poll(items_, static_cast(nitems_), timeout_); + if (rc < 0) + throw error_t(); + return rc; +} +} + +#ifdef ZMQ_CPP11 +ZMQ_DEPRECATED("from 4.8.0, use poll taking std::chrono::duration instead of long") +inline int poll(zmq_pollitem_t *items_, size_t nitems_, long timeout_) +#else +inline int poll(zmq_pollitem_t *items_, size_t nitems_, long timeout_ = -1) +#endif +{ + return detail::poll(items_, nitems_, timeout_); +} + +ZMQ_DEPRECATED("from 4.3.1, use poll taking non-const items") +inline int poll(zmq_pollitem_t const *items_, size_t nitems_, long timeout_ = -1) +{ + return detail::poll(const_cast(items_), nitems_, timeout_); +} + +#ifdef ZMQ_CPP11 +ZMQ_DEPRECATED("from 4.3.1, use poll taking non-const items") +inline int +poll(zmq_pollitem_t const *items, size_t nitems, std::chrono::milliseconds timeout) +{ + return detail::poll(const_cast(items), nitems, + static_cast(timeout.count())); +} + +ZMQ_DEPRECATED("from 4.3.1, use poll taking non-const items") +inline int poll(std::vector const &items, + std::chrono::milliseconds timeout) +{ + return detail::poll(const_cast(items.data()), items.size(), + static_cast(timeout.count())); +} + +ZMQ_DEPRECATED("from 4.3.1, use poll taking non-const items") +inline int poll(std::vector const &items, long timeout_ = -1) +{ + return detail::poll(const_cast(items.data()), items.size(), timeout_); +} + +inline int +poll(zmq_pollitem_t *items, size_t nitems, std::chrono::milliseconds timeout = std::chrono::milliseconds{-1}) +{ + return detail::poll(items, nitems, static_cast(timeout.count())); +} + +inline int poll(std::vector &items, + std::chrono::milliseconds timeout = std::chrono::milliseconds{-1}) +{ + return detail::poll(items.data(), items.size(), static_cast(timeout.count())); +} + +ZMQ_DEPRECATED("from 4.3.1, use poll taking std::chrono::duration instead of long") +inline int poll(std::vector &items, long timeout_) +{ + return detail::poll(items.data(), items.size(), timeout_); +} + +template +inline int poll(std::array &items, + std::chrono::milliseconds timeout = std::chrono::milliseconds{-1}) +{ + return detail::poll(items.data(), items.size(), static_cast(timeout.count())); +} +#endif + + +inline void version(int *major_, int *minor_, int *patch_) +{ + zmq_version(major_, minor_, patch_); +} + +#ifdef ZMQ_CPP11 +inline std::tuple version() +{ + std::tuple v; + zmq_version(&std::get<0>(v), &std::get<1>(v), &std::get<2>(v)); + return v; +} + +#if !defined(ZMQ_CPP11_PARTIAL) +namespace detail +{ +template struct is_char_type +{ + // true if character type for string literals in C++11 + static constexpr bool value = + std::is_same::value || std::is_same::value + || std::is_same::value || std::is_same::value; +}; +} +#endif + +#endif + +class message_t +{ + public: + message_t() ZMQ_NOTHROW + { + int rc = zmq_msg_init(&msg); + ZMQ_ASSERT(rc == 0); + } + + explicit message_t(size_t size_) + { + int rc = zmq_msg_init_size(&msg, size_); + if (rc != 0) + throw error_t(); + } + + template message_t(ForwardIter first, ForwardIter last) + { + typedef typename std::iterator_traits::value_type value_t; + + assert(std::distance(first, last) >= 0); + size_t const size_ = + static_cast(std::distance(first, last)) * sizeof(value_t); + int const rc = zmq_msg_init_size(&msg, size_); + if (rc != 0) + throw error_t(); + std::copy(first, last, data()); + } + + message_t(const void *data_, size_t size_) + { + int rc = zmq_msg_init_size(&msg, size_); + if (rc != 0) + throw error_t(); + if (size_) { + // this constructor allows (nullptr, 0), + // memcpy with a null pointer is UB + memcpy(data(), data_, size_); + } + } + + message_t(void *data_, size_t size_, free_fn *ffn_, void *hint_ = ZMQ_NULLPTR) + { + int rc = zmq_msg_init_data(&msg, data_, size_, ffn_, hint_); + if (rc != 0) + throw error_t(); + } + + // overload set of string-like types and generic containers +#if defined(ZMQ_CPP11) && !defined(ZMQ_CPP11_PARTIAL) + // NOTE this constructor will include the null terminator + // when called with a string literal. + // An overload taking const char* can not be added because + // it would be preferred over this function and break compatiblity. + template< + class Char, + size_t N, + typename = typename std::enable_if::value>::type> + ZMQ_DEPRECATED("from 4.7.0, use constructors taking iterators, (pointer, size) " + "or strings instead") + explicit message_t(const Char (&data)[N]) : + message_t(detail::ranges::begin(data), detail::ranges::end(data)) + { + } + + template::value + && ZMQ_IS_TRIVIALLY_COPYABLE(detail::range_value_t) + && !detail::is_char_type>::value + && !std::is_same::value>::type> + explicit message_t(const Range &rng) : + message_t(detail::ranges::begin(rng), detail::ranges::end(rng)) + { + } + + explicit message_t(const std::string &str) : message_t(str.data(), str.size()) {} + +#if CPPZMQ_HAS_STRING_VIEW + explicit message_t(std::string_view str) : message_t(str.data(), str.size()) {} +#endif + +#endif + +#ifdef ZMQ_HAS_RVALUE_REFS + message_t(message_t &&rhs) ZMQ_NOTHROW : msg(rhs.msg) + { + int rc = zmq_msg_init(&rhs.msg); + ZMQ_ASSERT(rc == 0); + } + + message_t &operator=(message_t &&rhs) ZMQ_NOTHROW + { + std::swap(msg, rhs.msg); + return *this; + } +#endif + + ~message_t() ZMQ_NOTHROW + { + int rc = zmq_msg_close(&msg); + ZMQ_ASSERT(rc == 0); + } + + void rebuild() + { + int rc = zmq_msg_close(&msg); + if (rc != 0) + throw error_t(); + rc = zmq_msg_init(&msg); + ZMQ_ASSERT(rc == 0); + } + + void rebuild(size_t size_) + { + int rc = zmq_msg_close(&msg); + if (rc != 0) + throw error_t(); + rc = zmq_msg_init_size(&msg, size_); + if (rc != 0) + throw error_t(); + } + + void rebuild(const void *data_, size_t size_) + { + int rc = zmq_msg_close(&msg); + if (rc != 0) + throw error_t(); + rc = zmq_msg_init_size(&msg, size_); + if (rc != 0) + throw error_t(); + memcpy(data(), data_, size_); + } + + void rebuild(const std::string &str) + { + rebuild(str.data(), str.size()); + } + + void rebuild(void *data_, size_t size_, free_fn *ffn_, void *hint_ = ZMQ_NULLPTR) + { + int rc = zmq_msg_close(&msg); + if (rc != 0) + throw error_t(); + rc = zmq_msg_init_data(&msg, data_, size_, ffn_, hint_); + if (rc != 0) + throw error_t(); + } + + ZMQ_DEPRECATED("from 4.3.1, use move taking non-const reference instead") + void move(message_t const *msg_) + { + int rc = zmq_msg_move(&msg, const_cast(msg_->handle())); + if (rc != 0) + throw error_t(); + } + + void move(message_t &msg_) + { + int rc = zmq_msg_move(&msg, msg_.handle()); + if (rc != 0) + throw error_t(); + } + + ZMQ_DEPRECATED("from 4.3.1, use copy taking non-const reference instead") + void copy(message_t const *msg_) + { + int rc = zmq_msg_copy(&msg, const_cast(msg_->handle())); + if (rc != 0) + throw error_t(); + } + + void copy(message_t &msg_) + { + int rc = zmq_msg_copy(&msg, msg_.handle()); + if (rc != 0) + throw error_t(); + } + + bool more() const ZMQ_NOTHROW + { + int rc = zmq_msg_more(const_cast(&msg)); + return rc != 0; + } + + void *data() ZMQ_NOTHROW { return zmq_msg_data(&msg); } + + const void *data() const ZMQ_NOTHROW + { + return zmq_msg_data(const_cast(&msg)); + } + + size_t size() const ZMQ_NOTHROW + { + return zmq_msg_size(const_cast(&msg)); + } + + ZMQ_NODISCARD bool empty() const ZMQ_NOTHROW { return size() == 0u; } + + template T *data() ZMQ_NOTHROW { return static_cast(data()); } + + template T const *data() const ZMQ_NOTHROW + { + return static_cast(data()); + } + + ZMQ_DEPRECATED("from 4.3.0, use operator== instead") + bool equal(const message_t *other) const ZMQ_NOTHROW { return *this == *other; } + + bool operator==(const message_t &other) const ZMQ_NOTHROW + { + const size_t my_size = size(); + return my_size == other.size() && 0 == memcmp(data(), other.data(), my_size); + } + + bool operator!=(const message_t &other) const ZMQ_NOTHROW + { + return !(*this == other); + } + +#if ZMQ_VERSION >= ZMQ_MAKE_VERSION(3, 2, 0) + int get(int property_) + { + int value = zmq_msg_get(&msg, property_); + if (value == -1) + throw error_t(); + return value; + } +#endif + +#if ZMQ_VERSION >= ZMQ_MAKE_VERSION(4, 1, 0) + const char *gets(const char *property_) + { + const char *value = zmq_msg_gets(&msg, property_); + if (value == ZMQ_NULLPTR) + throw error_t(); + return value; + } +#endif + +#if defined(ZMQ_BUILD_DRAFT_API) && ZMQ_VERSION >= ZMQ_MAKE_VERSION(4, 2, 0) + uint32_t routing_id() const + { + return zmq_msg_routing_id(const_cast(&msg)); + } + + void set_routing_id(uint32_t routing_id) + { + int rc = zmq_msg_set_routing_id(&msg, routing_id); + if (rc != 0) + throw error_t(); + } + + const char *group() const + { + return zmq_msg_group(const_cast(&msg)); + } + + void set_group(const char *group) + { + int rc = zmq_msg_set_group(&msg, group); + if (rc != 0) + throw error_t(); + } +#endif + + // interpret message content as a string + std::string to_string() const + { + return std::string(static_cast(data()), size()); + } +#if CPPZMQ_HAS_STRING_VIEW + // interpret message content as a string + std::string_view to_string_view() const noexcept + { + return std::string_view(static_cast(data()), size()); + } +#endif + + /** Dump content to string for debugging. + * Ascii chars are readable, the rest is printed as hex. + * Probably ridiculously slow. + * Use to_string() or to_string_view() for + * interpreting the message as a string. + */ + std::string str() const + { + // Partly mutuated from the same method in zmq::multipart_t + std::stringstream os; + + const unsigned char *msg_data = this->data(); + unsigned char byte; + size_t size = this->size(); + int is_ascii[2] = {0, 0}; + + os << "zmq::message_t [size " << std::dec << std::setw(3) + << std::setfill('0') << size << "] ("; + // Totally arbitrary + if (size >= 1000) { + os << "... too big to print)"; + } else { + while (size--) { + byte = *msg_data++; + + is_ascii[1] = (byte >= 32 && byte < 127); + if (is_ascii[1] != is_ascii[0]) + os << " "; // Separate text/non text + + if (is_ascii[1]) { + os << byte; + } else { + os << std::hex << std::uppercase << std::setw(2) + << std::setfill('0') << static_cast(byte); + } + is_ascii[0] = is_ascii[1]; + } + os << ")"; + } + return os.str(); + } + + void swap(message_t &other) ZMQ_NOTHROW + { + // this assumes zmq::msg_t from libzmq is trivially relocatable + std::swap(msg, other.msg); + } + + ZMQ_NODISCARD zmq_msg_t *handle() ZMQ_NOTHROW { return &msg; } + ZMQ_NODISCARD const zmq_msg_t *handle() const ZMQ_NOTHROW { return &msg; } + + private: + // The underlying message + zmq_msg_t msg; + + // Disable implicit message copying, so that users won't use shared + // messages (less efficient) without being aware of the fact. + message_t(const message_t &) ZMQ_DELETED_FUNCTION; + void operator=(const message_t &) ZMQ_DELETED_FUNCTION; +}; + +inline void swap(message_t &a, message_t &b) ZMQ_NOTHROW +{ + a.swap(b); +} + +#ifdef ZMQ_CPP11 +enum class ctxopt +{ +#ifdef ZMQ_BLOCKY + blocky = ZMQ_BLOCKY, +#endif +#ifdef ZMQ_IO_THREADS + io_threads = ZMQ_IO_THREADS, +#endif +#ifdef ZMQ_THREAD_SCHED_POLICY + thread_sched_policy = ZMQ_THREAD_SCHED_POLICY, +#endif +#ifdef ZMQ_THREAD_PRIORITY + thread_priority = ZMQ_THREAD_PRIORITY, +#endif +#ifdef ZMQ_THREAD_AFFINITY_CPU_ADD + thread_affinity_cpu_add = ZMQ_THREAD_AFFINITY_CPU_ADD, +#endif +#ifdef ZMQ_THREAD_AFFINITY_CPU_REMOVE + thread_affinity_cpu_remove = ZMQ_THREAD_AFFINITY_CPU_REMOVE, +#endif +#ifdef ZMQ_THREAD_NAME_PREFIX + thread_name_prefix = ZMQ_THREAD_NAME_PREFIX, +#endif +#ifdef ZMQ_MAX_MSGSZ + max_msgsz = ZMQ_MAX_MSGSZ, +#endif +#ifdef ZMQ_ZERO_COPY_RECV + zero_copy_recv = ZMQ_ZERO_COPY_RECV, +#endif +#ifdef ZMQ_MAX_SOCKETS + max_sockets = ZMQ_MAX_SOCKETS, +#endif +#ifdef ZMQ_SOCKET_LIMIT + socket_limit = ZMQ_SOCKET_LIMIT, +#endif +#ifdef ZMQ_IPV6 + ipv6 = ZMQ_IPV6, +#endif +#ifdef ZMQ_MSG_T_SIZE + msg_t_size = ZMQ_MSG_T_SIZE +#endif +}; +#endif + +class context_t +{ + public: + context_t() + { + ptr = zmq_ctx_new(); + if (ptr == ZMQ_NULLPTR) + throw error_t(); + } + + + explicit context_t(int io_threads_, int max_sockets_ = ZMQ_MAX_SOCKETS_DFLT) + { + ptr = zmq_ctx_new(); + if (ptr == ZMQ_NULLPTR) + throw error_t(); + + int rc = zmq_ctx_set(ptr, ZMQ_IO_THREADS, io_threads_); + ZMQ_ASSERT(rc == 0); + + rc = zmq_ctx_set(ptr, ZMQ_MAX_SOCKETS, max_sockets_); + ZMQ_ASSERT(rc == 0); + } + +#ifdef ZMQ_HAS_RVALUE_REFS + context_t(context_t &&rhs) ZMQ_NOTHROW : ptr(rhs.ptr) { rhs.ptr = ZMQ_NULLPTR; } + context_t &operator=(context_t &&rhs) ZMQ_NOTHROW + { + close(); + std::swap(ptr, rhs.ptr); + return *this; + } +#endif + + ~context_t() ZMQ_NOTHROW { close(); } + + ZMQ_CPP11_DEPRECATED("from 4.7.0, use set taking zmq::ctxopt instead") + int setctxopt(int option_, int optval_) + { + int rc = zmq_ctx_set(ptr, option_, optval_); + ZMQ_ASSERT(rc == 0); + return rc; + } + + ZMQ_CPP11_DEPRECATED("from 4.7.0, use get taking zmq::ctxopt instead") + int getctxopt(int option_) { return zmq_ctx_get(ptr, option_); } + +#ifdef ZMQ_CPP11 + void set(ctxopt option, int optval) + { + int rc = zmq_ctx_set(ptr, static_cast(option), optval); + if (rc == -1) + throw error_t(); + } + + ZMQ_NODISCARD int get(ctxopt option) + { + int rc = zmq_ctx_get(ptr, static_cast(option)); + // some options have a default value of -1 + // which is unfortunate, and may result in errors + // that don't make sense + if (rc == -1) + throw error_t(); + return rc; + } +#endif + + // Terminates context (see also shutdown()). + void close() ZMQ_NOTHROW + { + if (ptr == ZMQ_NULLPTR) + return; + + int rc; + do { + rc = zmq_ctx_term(ptr); + } while (rc == -1 && errno == EINTR); + + ZMQ_ASSERT(rc == 0); + ptr = ZMQ_NULLPTR; + } + + // Shutdown context in preparation for termination (close()). + // Causes all blocking socket operations and any further + // socket operations to return with ETERM. + void shutdown() ZMQ_NOTHROW + { + if (ptr == ZMQ_NULLPTR) + return; + int rc = zmq_ctx_shutdown(ptr); + ZMQ_ASSERT(rc == 0); + } + + // Be careful with this, it's probably only useful for + // using the C api together with an existing C++ api. + // Normally you should never need to use this. + ZMQ_EXPLICIT operator void *() ZMQ_NOTHROW { return ptr; } + + ZMQ_EXPLICIT operator void const *() const ZMQ_NOTHROW { return ptr; } + + ZMQ_NODISCARD void *handle() ZMQ_NOTHROW { return ptr; } + + ZMQ_DEPRECATED("from 4.7.0, use handle() != nullptr instead") + operator bool() const ZMQ_NOTHROW { return ptr != ZMQ_NULLPTR; } + + void swap(context_t &other) ZMQ_NOTHROW { std::swap(ptr, other.ptr); } + + private: + void *ptr; + + context_t(const context_t &) ZMQ_DELETED_FUNCTION; + void operator=(const context_t &) ZMQ_DELETED_FUNCTION; +}; + +inline void swap(context_t &a, context_t &b) ZMQ_NOTHROW +{ + a.swap(b); +} + +#ifdef ZMQ_CPP11 + +struct recv_buffer_size +{ + size_t size; // number of bytes written to buffer + size_t untruncated_size; // untruncated message size in bytes + + ZMQ_NODISCARD bool truncated() const noexcept + { + return size != untruncated_size; + } +}; + +#if CPPZMQ_HAS_OPTIONAL + +using send_result_t = std::optional; +using recv_result_t = std::optional; +using recv_buffer_result_t = std::optional; + +#else + +namespace detail +{ +// A C++11 type emulating the most basic +// operations of std::optional for trivial types +template class trivial_optional +{ + public: + static_assert(std::is_trivial::value, "T must be trivial"); + using value_type = T; + + trivial_optional() = default; + trivial_optional(T value) noexcept : _value(value), _has_value(true) {} + + const T *operator->() const noexcept + { + assert(_has_value); + return &_value; + } + T *operator->() noexcept + { + assert(_has_value); + return &_value; + } + + const T &operator*() const noexcept + { + assert(_has_value); + return _value; + } + T &operator*() noexcept + { + assert(_has_value); + return _value; + } + + T &value() + { + if (!_has_value) + throw std::exception(); + return _value; + } + const T &value() const + { + if (!_has_value) + throw std::exception(); + return _value; + } + + explicit operator bool() const noexcept { return _has_value; } + bool has_value() const noexcept { return _has_value; } + + private: + T _value{}; + bool _has_value{false}; +}; +} // namespace detail + +using send_result_t = detail::trivial_optional; +using recv_result_t = detail::trivial_optional; +using recv_buffer_result_t = detail::trivial_optional; + +#endif + +namespace detail +{ +template constexpr T enum_bit_or(T a, T b) noexcept +{ + static_assert(std::is_enum::value, "must be enum"); + using U = typename std::underlying_type::type; + return static_cast(static_cast(a) | static_cast(b)); +} +template constexpr T enum_bit_and(T a, T b) noexcept +{ + static_assert(std::is_enum::value, "must be enum"); + using U = typename std::underlying_type::type; + return static_cast(static_cast(a) & static_cast(b)); +} +template constexpr T enum_bit_xor(T a, T b) noexcept +{ + static_assert(std::is_enum::value, "must be enum"); + using U = typename std::underlying_type::type; + return static_cast(static_cast(a) ^ static_cast(b)); +} +template constexpr T enum_bit_not(T a) noexcept +{ + static_assert(std::is_enum::value, "must be enum"); + using U = typename std::underlying_type::type; + return static_cast(~static_cast(a)); +} +} // namespace detail + +// partially satisfies named requirement BitmaskType +enum class send_flags : int +{ + none = 0, + dontwait = ZMQ_DONTWAIT, + sndmore = ZMQ_SNDMORE +}; + +constexpr send_flags operator|(send_flags a, send_flags b) noexcept +{ + return detail::enum_bit_or(a, b); +} +constexpr send_flags operator&(send_flags a, send_flags b) noexcept +{ + return detail::enum_bit_and(a, b); +} +constexpr send_flags operator^(send_flags a, send_flags b) noexcept +{ + return detail::enum_bit_xor(a, b); +} +constexpr send_flags operator~(send_flags a) noexcept +{ + return detail::enum_bit_not(a); +} + +// partially satisfies named requirement BitmaskType +enum class recv_flags : int +{ + none = 0, + dontwait = ZMQ_DONTWAIT +}; + +constexpr recv_flags operator|(recv_flags a, recv_flags b) noexcept +{ + return detail::enum_bit_or(a, b); +} +constexpr recv_flags operator&(recv_flags a, recv_flags b) noexcept +{ + return detail::enum_bit_and(a, b); +} +constexpr recv_flags operator^(recv_flags a, recv_flags b) noexcept +{ + return detail::enum_bit_xor(a, b); +} +constexpr recv_flags operator~(recv_flags a) noexcept +{ + return detail::enum_bit_not(a); +} + + +// mutable_buffer, const_buffer and buffer are based on +// the Networking TS specification, draft: +// http://www.open-std.org/jtc1/sc22/wg21/docs/papers/2018/n4771.pdf + +class mutable_buffer +{ + public: + constexpr mutable_buffer() noexcept : _data(nullptr), _size(0) {} + constexpr mutable_buffer(void *p, size_t n) noexcept : _data(p), _size(n) + { +#ifdef ZMQ_EXTENDED_CONSTEXPR + assert(p != nullptr || n == 0); +#endif + } + + constexpr void *data() const noexcept { return _data; } + constexpr size_t size() const noexcept { return _size; } + mutable_buffer &operator+=(size_t n) noexcept + { + // (std::min) is a workaround for when a min macro is defined + const auto shift = (std::min)(n, _size); + _data = static_cast(_data) + shift; + _size -= shift; + return *this; + } + + private: + void *_data; + size_t _size; +}; + +inline mutable_buffer operator+(const mutable_buffer &mb, size_t n) noexcept +{ + return mutable_buffer(static_cast(mb.data()) + (std::min)(n, mb.size()), + mb.size() - (std::min)(n, mb.size())); +} +inline mutable_buffer operator+(size_t n, const mutable_buffer &mb) noexcept +{ + return mb + n; +} + +class const_buffer +{ + public: + constexpr const_buffer() noexcept : _data(nullptr), _size(0) {} + constexpr const_buffer(const void *p, size_t n) noexcept : _data(p), _size(n) + { +#ifdef ZMQ_EXTENDED_CONSTEXPR + assert(p != nullptr || n == 0); +#endif + } + constexpr const_buffer(const mutable_buffer &mb) noexcept : + _data(mb.data()), _size(mb.size()) + { + } + + constexpr const void *data() const noexcept { return _data; } + constexpr size_t size() const noexcept { return _size; } + const_buffer &operator+=(size_t n) noexcept + { + const auto shift = (std::min)(n, _size); + _data = static_cast(_data) + shift; + _size -= shift; + return *this; + } + + private: + const void *_data; + size_t _size; +}; + +inline const_buffer operator+(const const_buffer &cb, size_t n) noexcept +{ + return const_buffer(static_cast(cb.data()) + + (std::min)(n, cb.size()), + cb.size() - (std::min)(n, cb.size())); +} +inline const_buffer operator+(size_t n, const const_buffer &cb) noexcept +{ + return cb + n; +} + +// buffer creation + +constexpr mutable_buffer buffer(void *p, size_t n) noexcept +{ + return mutable_buffer(p, n); +} +constexpr const_buffer buffer(const void *p, size_t n) noexcept +{ + return const_buffer(p, n); +} +constexpr mutable_buffer buffer(const mutable_buffer &mb) noexcept +{ + return mb; +} +inline mutable_buffer buffer(const mutable_buffer &mb, size_t n) noexcept +{ + return mutable_buffer(mb.data(), (std::min)(mb.size(), n)); +} +constexpr const_buffer buffer(const const_buffer &cb) noexcept +{ + return cb; +} +inline const_buffer buffer(const const_buffer &cb, size_t n) noexcept +{ + return const_buffer(cb.data(), (std::min)(cb.size(), n)); +} + +namespace detail +{ +template struct is_buffer +{ + static constexpr bool value = + std::is_same::value || std::is_same::value; +}; + +template struct is_pod_like +{ + // NOTE: The networking draft N4771 section 16.11 requires + // T in the buffer functions below to be + // trivially copyable OR standard layout. + // Here we decide to be conservative and require both. + static constexpr bool value = + ZMQ_IS_TRIVIALLY_COPYABLE(T) && std::is_standard_layout::value; +}; + +template constexpr auto seq_size(const C &c) noexcept -> decltype(c.size()) +{ + return c.size(); +} +template +constexpr size_t seq_size(const T (&/*array*/)[N]) noexcept +{ + return N; +} + +template +auto buffer_contiguous_sequence(Seq &&seq) noexcept + -> decltype(buffer(std::addressof(*std::begin(seq)), size_t{})) +{ + using T = typename std::remove_cv< + typename std::remove_reference::type>::type; + static_assert(detail::is_pod_like::value, "T must be POD"); + + const auto size = seq_size(seq); + return buffer(size != 0u ? std::addressof(*std::begin(seq)) : nullptr, + size * sizeof(T)); +} +template +auto buffer_contiguous_sequence(Seq &&seq, size_t n_bytes) noexcept + -> decltype(buffer_contiguous_sequence(seq)) +{ + using T = typename std::remove_cv< + typename std::remove_reference::type>::type; + static_assert(detail::is_pod_like::value, "T must be POD"); + + const auto size = seq_size(seq); + return buffer(size != 0u ? std::addressof(*std::begin(seq)) : nullptr, + (std::min)(size * sizeof(T), n_bytes)); +} + +} // namespace detail + +// C array +template mutable_buffer buffer(T (&data)[N]) noexcept +{ + return detail::buffer_contiguous_sequence(data); +} +template +mutable_buffer buffer(T (&data)[N], size_t n_bytes) noexcept +{ + return detail::buffer_contiguous_sequence(data, n_bytes); +} +template const_buffer buffer(const T (&data)[N]) noexcept +{ + return detail::buffer_contiguous_sequence(data); +} +template +const_buffer buffer(const T (&data)[N], size_t n_bytes) noexcept +{ + return detail::buffer_contiguous_sequence(data, n_bytes); +} +// std::array +template mutable_buffer buffer(std::array &data) noexcept +{ + return detail::buffer_contiguous_sequence(data); +} +template +mutable_buffer buffer(std::array &data, size_t n_bytes) noexcept +{ + return detail::buffer_contiguous_sequence(data, n_bytes); +} +template +const_buffer buffer(std::array &data) noexcept +{ + return detail::buffer_contiguous_sequence(data); +} +template +const_buffer buffer(std::array &data, size_t n_bytes) noexcept +{ + return detail::buffer_contiguous_sequence(data, n_bytes); +} +template +const_buffer buffer(const std::array &data) noexcept +{ + return detail::buffer_contiguous_sequence(data); +} +template +const_buffer buffer(const std::array &data, size_t n_bytes) noexcept +{ + return detail::buffer_contiguous_sequence(data, n_bytes); +} +// std::vector +template +mutable_buffer buffer(std::vector &data) noexcept +{ + return detail::buffer_contiguous_sequence(data); +} +template +mutable_buffer buffer(std::vector &data, size_t n_bytes) noexcept +{ + return detail::buffer_contiguous_sequence(data, n_bytes); +} +template +const_buffer buffer(const std::vector &data) noexcept +{ + return detail::buffer_contiguous_sequence(data); +} +template +const_buffer buffer(const std::vector &data, size_t n_bytes) noexcept +{ + return detail::buffer_contiguous_sequence(data, n_bytes); +} +// std::basic_string +template +mutable_buffer buffer(std::basic_string &data) noexcept +{ + return detail::buffer_contiguous_sequence(data); +} +template +mutable_buffer buffer(std::basic_string &data, + size_t n_bytes) noexcept +{ + return detail::buffer_contiguous_sequence(data, n_bytes); +} +template +const_buffer buffer(const std::basic_string &data) noexcept +{ + return detail::buffer_contiguous_sequence(data); +} +template +const_buffer buffer(const std::basic_string &data, + size_t n_bytes) noexcept +{ + return detail::buffer_contiguous_sequence(data, n_bytes); +} + +#if CPPZMQ_HAS_STRING_VIEW +// std::basic_string_view +template +const_buffer buffer(std::basic_string_view data) noexcept +{ + return detail::buffer_contiguous_sequence(data); +} +template +const_buffer buffer(std::basic_string_view data, size_t n_bytes) noexcept +{ + return detail::buffer_contiguous_sequence(data, n_bytes); +} +#endif + +// Buffer for a string literal (null terminated) +// where the buffer size excludes the terminating character. +// Equivalent to zmq::buffer(std::string_view("...")). +template +constexpr const_buffer str_buffer(const Char (&data)[N]) noexcept +{ + static_assert(detail::is_pod_like::value, "Char must be POD"); +#ifdef ZMQ_EXTENDED_CONSTEXPR + assert(data[N - 1] == Char{0}); +#endif + return const_buffer(static_cast(data), (N - 1) * sizeof(Char)); +} + +namespace literals +{ +constexpr const_buffer operator"" _zbuf(const char *str, size_t len) noexcept +{ + return const_buffer(str, len * sizeof(char)); +} +constexpr const_buffer operator"" _zbuf(const wchar_t *str, size_t len) noexcept +{ + return const_buffer(str, len * sizeof(wchar_t)); +} +constexpr const_buffer operator"" _zbuf(const char16_t *str, size_t len) noexcept +{ + return const_buffer(str, len * sizeof(char16_t)); +} +constexpr const_buffer operator"" _zbuf(const char32_t *str, size_t len) noexcept +{ + return const_buffer(str, len * sizeof(char32_t)); +} +} + +#ifdef ZMQ_CPP11 +enum class socket_type : int +{ + req = ZMQ_REQ, + rep = ZMQ_REP, + dealer = ZMQ_DEALER, + router = ZMQ_ROUTER, + pub = ZMQ_PUB, + sub = ZMQ_SUB, + xpub = ZMQ_XPUB, + xsub = ZMQ_XSUB, + push = ZMQ_PUSH, + pull = ZMQ_PULL, +#if defined(ZMQ_BUILD_DRAFT_API) && ZMQ_VERSION >= ZMQ_MAKE_VERSION(4, 2, 0) + server = ZMQ_SERVER, + client = ZMQ_CLIENT, + radio = ZMQ_RADIO, + dish = ZMQ_DISH, + gather = ZMQ_GATHER, + scatter = ZMQ_SCATTER, + dgram = ZMQ_DGRAM, +#endif +#if defined(ZMQ_BUILD_DRAFT_API) && ZMQ_VERSION >= ZMQ_MAKE_VERSION(4, 3, 3) + peer = ZMQ_PEER, + channel = ZMQ_CHANNEL, +#endif +#if ZMQ_VERSION_MAJOR >= 4 + stream = ZMQ_STREAM, +#endif + pair = ZMQ_PAIR +}; +#endif + +namespace sockopt +{ +// There are two types of options, +// integral type with known compiler time size (int, bool, int64_t, uint64_t) +// and arrays with dynamic size (strings, binary data). + +// BoolUnit: if true accepts values of type bool (but passed as T into libzmq) +template struct integral_option +{ +}; + +// NullTerm: +// 0: binary data +// 1: null-terminated string (`getsockopt` size includes null) +// 2: binary (size 32) or Z85 encoder string of size 41 (null included) +template struct array_option +{ +}; + +#define ZMQ_DEFINE_INTEGRAL_OPT(OPT, NAME, TYPE) \ + using NAME##_t = integral_option; \ + ZMQ_INLINE_VAR ZMQ_CONSTEXPR_VAR NAME##_t NAME {} +#define ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(OPT, NAME, TYPE) \ + using NAME##_t = integral_option; \ + ZMQ_INLINE_VAR ZMQ_CONSTEXPR_VAR NAME##_t NAME {} +#define ZMQ_DEFINE_ARRAY_OPT(OPT, NAME) \ + using NAME##_t = array_option; \ + ZMQ_INLINE_VAR ZMQ_CONSTEXPR_VAR NAME##_t NAME {} +#define ZMQ_DEFINE_ARRAY_OPT_BINARY(OPT, NAME) \ + using NAME##_t = array_option; \ + ZMQ_INLINE_VAR ZMQ_CONSTEXPR_VAR NAME##_t NAME {} +#define ZMQ_DEFINE_ARRAY_OPT_BIN_OR_Z85(OPT, NAME) \ + using NAME##_t = array_option; \ + ZMQ_INLINE_VAR ZMQ_CONSTEXPR_VAR NAME##_t NAME {} + +// deprecated, use zmq::fd_t +using cppzmq_fd_t = ::zmq::fd_t; + +#ifdef ZMQ_AFFINITY +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_AFFINITY, affinity, uint64_t); +#endif +#ifdef ZMQ_BACKLOG +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_BACKLOG, backlog, int); +#endif +#ifdef ZMQ_BINDTODEVICE +ZMQ_DEFINE_ARRAY_OPT_BINARY(ZMQ_BINDTODEVICE, bindtodevice); +#endif +#ifdef ZMQ_CONFLATE +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_CONFLATE, conflate, int); +#endif +#ifdef ZMQ_CONNECT_ROUTING_ID +ZMQ_DEFINE_ARRAY_OPT(ZMQ_CONNECT_ROUTING_ID, connect_routing_id); +#endif +#ifdef ZMQ_CONNECT_TIMEOUT +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_CONNECT_TIMEOUT, connect_timeout, int); +#endif +#ifdef ZMQ_CURVE_PUBLICKEY +ZMQ_DEFINE_ARRAY_OPT_BIN_OR_Z85(ZMQ_CURVE_PUBLICKEY, curve_publickey); +#endif +#ifdef ZMQ_CURVE_SECRETKEY +ZMQ_DEFINE_ARRAY_OPT_BIN_OR_Z85(ZMQ_CURVE_SECRETKEY, curve_secretkey); +#endif +#ifdef ZMQ_CURVE_SERVER +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_CURVE_SERVER, curve_server, int); +#endif +#ifdef ZMQ_CURVE_SERVERKEY +ZMQ_DEFINE_ARRAY_OPT_BIN_OR_Z85(ZMQ_CURVE_SERVERKEY, curve_serverkey); +#endif +#ifdef ZMQ_DISCONNECT_MSG +ZMQ_DEFINE_ARRAY_OPT_BINARY(ZMQ_DISCONNECT_MSG, disconnect_msg); +#endif +#ifdef ZMQ_EVENTS +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_EVENTS, events, int); +#endif +#ifdef ZMQ_FD +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_FD, fd, ::zmq::fd_t); +#endif +#ifdef ZMQ_GSSAPI_PLAINTEXT +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_GSSAPI_PLAINTEXT, gssapi_plaintext, int); +#endif +#ifdef ZMQ_GSSAPI_SERVER +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_GSSAPI_SERVER, gssapi_server, int); +#endif +#ifdef ZMQ_GSSAPI_SERVICE_PRINCIPAL +ZMQ_DEFINE_ARRAY_OPT(ZMQ_GSSAPI_SERVICE_PRINCIPAL, gssapi_service_principal); +#endif +#ifdef ZMQ_GSSAPI_SERVICE_PRINCIPAL_NAMETYPE +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_GSSAPI_SERVICE_PRINCIPAL_NAMETYPE, + gssapi_service_principal_nametype, + int); +#endif +#ifdef ZMQ_GSSAPI_PRINCIPAL +ZMQ_DEFINE_ARRAY_OPT(ZMQ_GSSAPI_PRINCIPAL, gssapi_principal); +#endif +#ifdef ZMQ_GSSAPI_PRINCIPAL_NAMETYPE +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_GSSAPI_PRINCIPAL_NAMETYPE, + gssapi_principal_nametype, + int); +#endif +#ifdef ZMQ_HANDSHAKE_IVL +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_HANDSHAKE_IVL, handshake_ivl, int); +#endif +#ifdef ZMQ_HEARTBEAT_IVL +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_HEARTBEAT_IVL, heartbeat_ivl, int); +#endif +#ifdef ZMQ_HEARTBEAT_TIMEOUT +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_HEARTBEAT_TIMEOUT, heartbeat_timeout, int); +#endif +#ifdef ZMQ_HEARTBEAT_TTL +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_HEARTBEAT_TTL, heartbeat_ttl, int); +#endif +#ifdef ZMQ_HELLO_MSG +ZMQ_DEFINE_ARRAY_OPT_BINARY(ZMQ_HELLO_MSG, hello_msg); +#endif +#ifdef ZMQ_IMMEDIATE +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_IMMEDIATE, immediate, int); +#endif +#ifdef ZMQ_INVERT_MATCHING +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_INVERT_MATCHING, invert_matching, int); +#endif +#ifdef ZMQ_IPV6 +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_IPV6, ipv6, int); +#endif +#ifdef ZMQ_LAST_ENDPOINT +ZMQ_DEFINE_ARRAY_OPT(ZMQ_LAST_ENDPOINT, last_endpoint); +#endif +#ifdef ZMQ_LINGER +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_LINGER, linger, int); +#endif +#ifdef ZMQ_MAXMSGSIZE +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_MAXMSGSIZE, maxmsgsize, int64_t); +#endif +#ifdef ZMQ_MECHANISM +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_MECHANISM, mechanism, int); +#endif +#ifdef ZMQ_METADATA +ZMQ_DEFINE_ARRAY_OPT(ZMQ_METADATA, metadata); +#endif +#ifdef ZMQ_MULTICAST_HOPS +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_MULTICAST_HOPS, multicast_hops, int); +#endif +#ifdef ZMQ_MULTICAST_LOOP +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_MULTICAST_LOOP, multicast_loop, int); +#endif +#ifdef ZMQ_MULTICAST_MAXTPDU +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_MULTICAST_MAXTPDU, multicast_maxtpdu, int); +#endif +#ifdef ZMQ_ONLY_FIRST_SUBSCRIBE +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_ONLY_FIRST_SUBSCRIBE, only_first_subscribe, int); +#endif +#ifdef ZMQ_PLAIN_SERVER +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_PLAIN_SERVER, plain_server, int); +#endif +#ifdef ZMQ_PLAIN_PASSWORD +ZMQ_DEFINE_ARRAY_OPT(ZMQ_PLAIN_PASSWORD, plain_password); +#endif +#ifdef ZMQ_PLAIN_USERNAME +ZMQ_DEFINE_ARRAY_OPT(ZMQ_PLAIN_USERNAME, plain_username); +#endif +#ifdef ZMQ_PRIORITY +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_PRIORITY, priority, int); +#endif +#ifdef ZMQ_USE_FD +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_USE_FD, use_fd, int); +#endif +#ifdef ZMQ_PROBE_ROUTER +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_PROBE_ROUTER, probe_router, int); +#endif +#ifdef ZMQ_RATE +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_RATE, rate, int); +#endif +#ifdef ZMQ_RCVBUF +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_RCVBUF, rcvbuf, int); +#endif +#ifdef ZMQ_RCVHWM +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_RCVHWM, rcvhwm, int); +#endif +#ifdef ZMQ_RCVMORE +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_RCVMORE, rcvmore, int); +#endif +#ifdef ZMQ_RCVTIMEO +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_RCVTIMEO, rcvtimeo, int); +#endif +#ifdef ZMQ_RECONNECT_IVL +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_RECONNECT_IVL, reconnect_ivl, int); +#endif +#ifdef ZMQ_RECONNECT_IVL_MAX +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_RECONNECT_IVL_MAX, reconnect_ivl_max, int); +#endif +#ifdef ZMQ_RECONNECT_STOP +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_RECONNECT_STOP, reconnect_stop, int); +#endif +#ifdef ZMQ_RECOVERY_IVL +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_RECOVERY_IVL, recovery_ivl, int); +#endif +#ifdef ZMQ_REQ_CORRELATE +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_REQ_CORRELATE, req_correlate, int); +#endif +#ifdef ZMQ_REQ_RELAXED +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_REQ_RELAXED, req_relaxed, int); +#endif +#ifdef ZMQ_ROUTER_HANDOVER +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_ROUTER_HANDOVER, router_handover, int); +#endif +#ifdef ZMQ_ROUTER_MANDATORY +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_ROUTER_MANDATORY, router_mandatory, int); +#endif +#ifdef ZMQ_ROUTER_NOTIFY +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_ROUTER_NOTIFY, router_notify, int); +#endif +#ifdef ZMQ_ROUTER_RAW +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_ROUTER_RAW, router_raw, int); +#endif +#ifdef ZMQ_ROUTING_ID +ZMQ_DEFINE_ARRAY_OPT_BINARY(ZMQ_ROUTING_ID, routing_id); +#endif +#ifdef ZMQ_SNDBUF +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_SNDBUF, sndbuf, int); +#endif +#ifdef ZMQ_SNDHWM +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_SNDHWM, sndhwm, int); +#endif +#ifdef ZMQ_SNDTIMEO +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_SNDTIMEO, sndtimeo, int); +#endif +#ifdef ZMQ_SOCKS_PASSWORD +ZMQ_DEFINE_ARRAY_OPT(ZMQ_SOCKS_PASSWORD, socks_password); +#endif +#ifdef ZMQ_SOCKS_PROXY +ZMQ_DEFINE_ARRAY_OPT(ZMQ_SOCKS_PROXY, socks_proxy); +#endif +#ifdef ZMQ_SOCKS_USERNAME +ZMQ_DEFINE_ARRAY_OPT(ZMQ_SOCKS_USERNAME, socks_username); +#endif +#ifdef ZMQ_STREAM_NOTIFY +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_STREAM_NOTIFY, stream_notify, int); +#endif +#ifdef ZMQ_SUBSCRIBE +ZMQ_DEFINE_ARRAY_OPT(ZMQ_SUBSCRIBE, subscribe); +#endif +#ifdef ZMQ_TCP_KEEPALIVE +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_TCP_KEEPALIVE, tcp_keepalive, int); +#endif +#ifdef ZMQ_TCP_KEEPALIVE_CNT +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_TCP_KEEPALIVE_CNT, tcp_keepalive_cnt, int); +#endif +#ifdef ZMQ_TCP_KEEPALIVE_IDLE +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_TCP_KEEPALIVE_IDLE, tcp_keepalive_idle, int); +#endif +#ifdef ZMQ_TCP_KEEPALIVE_INTVL +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_TCP_KEEPALIVE_INTVL, tcp_keepalive_intvl, int); +#endif +#ifdef ZMQ_TCP_MAXRT +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_TCP_MAXRT, tcp_maxrt, int); +#endif +#ifdef ZMQ_THREAD_SAFE +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_THREAD_SAFE, thread_safe, int); +#endif +#ifdef ZMQ_TOS +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_TOS, tos, int); +#endif +#ifdef ZMQ_TYPE +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_TYPE, type, int); +#ifdef ZMQ_CPP11 +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_TYPE, socket_type, socket_type); +#endif // ZMQ_CPP11 +#endif // ZMQ_TYPE +#ifdef ZMQ_UNSUBSCRIBE +ZMQ_DEFINE_ARRAY_OPT(ZMQ_UNSUBSCRIBE, unsubscribe); +#endif +#ifdef ZMQ_VMCI_BUFFER_SIZE +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_VMCI_BUFFER_SIZE, vmci_buffer_size, uint64_t); +#endif +#ifdef ZMQ_VMCI_BUFFER_MIN_SIZE +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_VMCI_BUFFER_MIN_SIZE, vmci_buffer_min_size, uint64_t); +#endif +#ifdef ZMQ_VMCI_BUFFER_MAX_SIZE +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_VMCI_BUFFER_MAX_SIZE, vmci_buffer_max_size, uint64_t); +#endif +#ifdef ZMQ_VMCI_CONNECT_TIMEOUT +ZMQ_DEFINE_INTEGRAL_OPT(ZMQ_VMCI_CONNECT_TIMEOUT, vmci_connect_timeout, int); +#endif +#ifdef ZMQ_XPUB_VERBOSE +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_XPUB_VERBOSE, xpub_verbose, int); +#endif +#ifdef ZMQ_XPUB_VERBOSER +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_XPUB_VERBOSER, xpub_verboser, int); +#endif +#ifdef ZMQ_XPUB_MANUAL +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_XPUB_MANUAL, xpub_manual, int); +#endif +#ifdef ZMQ_XPUB_MANUAL_LAST_VALUE +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_XPUB_MANUAL_LAST_VALUE, xpub_manual_last_value, int); +#endif +#ifdef ZMQ_XPUB_NODROP +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_XPUB_NODROP, xpub_nodrop, int); +#endif +#ifdef ZMQ_XPUB_WELCOME_MSG +ZMQ_DEFINE_ARRAY_OPT(ZMQ_XPUB_WELCOME_MSG, xpub_welcome_msg); +#endif +#ifdef ZMQ_ZAP_ENFORCE_DOMAIN +ZMQ_DEFINE_INTEGRAL_BOOL_UNIT_OPT(ZMQ_ZAP_ENFORCE_DOMAIN, zap_enforce_domain, int); +#endif +#ifdef ZMQ_ZAP_DOMAIN +ZMQ_DEFINE_ARRAY_OPT(ZMQ_ZAP_DOMAIN, zap_domain); +#endif + +} // namespace sockopt +#endif // ZMQ_CPP11 + + +namespace detail +{ +class socket_base +{ + public: + socket_base() ZMQ_NOTHROW : _handle(ZMQ_NULLPTR) {} + ZMQ_EXPLICIT socket_base(void *handle) ZMQ_NOTHROW : _handle(handle) {} + + template + ZMQ_CPP11_DEPRECATED("from 4.7.0, use `set` taking option from zmq::sockopt") + void setsockopt(int option_, T const &optval) + { + setsockopt(option_, &optval, sizeof(T)); + } + + ZMQ_CPP11_DEPRECATED("from 4.7.0, use `set` taking option from zmq::sockopt") + void setsockopt(int option_, const void *optval_, size_t optvallen_) + { + int rc = zmq_setsockopt(_handle, option_, optval_, optvallen_); + if (rc != 0) + throw error_t(); + } + + ZMQ_CPP11_DEPRECATED("from 4.7.0, use `get` taking option from zmq::sockopt") + void getsockopt(int option_, void *optval_, size_t *optvallen_) const + { + int rc = zmq_getsockopt(_handle, option_, optval_, optvallen_); + if (rc != 0) + throw error_t(); + } + + template + ZMQ_CPP11_DEPRECATED("from 4.7.0, use `get` taking option from zmq::sockopt") + T getsockopt(int option_) const + { + T optval; + size_t optlen = sizeof(T); + getsockopt(option_, &optval, &optlen); + return optval; + } + +#ifdef ZMQ_CPP11 + // Set integral socket option, e.g. + // `socket.set(zmq::sockopt::linger, 0)` + template + void set(sockopt::integral_option, const T &val) + { + static_assert(std::is_integral::value, "T must be integral"); + set_option(Opt, &val, sizeof val); + } + + // Set integral socket option from boolean, e.g. + // `socket.set(zmq::sockopt::immediate, false)` + template + void set(sockopt::integral_option, bool val) + { + static_assert(std::is_integral::value, "T must be integral"); + T rep_val = val; + set_option(Opt, &rep_val, sizeof rep_val); + } + + // Set array socket option, e.g. + // `socket.set(zmq::sockopt::plain_username, "foo123")` + template + void set(sockopt::array_option, const char *buf) + { + set_option(Opt, buf, std::strlen(buf)); + } + + // Set array socket option, e.g. + // `socket.set(zmq::sockopt::routing_id, zmq::buffer(id))` + template + void set(sockopt::array_option, const_buffer buf) + { + set_option(Opt, buf.data(), buf.size()); + } + + // Set array socket option, e.g. + // `socket.set(zmq::sockopt::routing_id, id_str)` + template + void set(sockopt::array_option, const std::string &buf) + { + set_option(Opt, buf.data(), buf.size()); + } + +#if CPPZMQ_HAS_STRING_VIEW + // Set array socket option, e.g. + // `socket.set(zmq::sockopt::routing_id, id_str)` + template + void set(sockopt::array_option, std::string_view buf) + { + set_option(Opt, buf.data(), buf.size()); + } +#endif + + // Get scalar socket option, e.g. + // `auto opt = socket.get(zmq::sockopt::linger)` + template + ZMQ_NODISCARD T get(sockopt::integral_option) const + { + static_assert(std::is_scalar::value, "T must be scalar"); + T val; + size_t size = sizeof val; + get_option(Opt, &val, &size); + assert(size == sizeof val); + return val; + } + + // Get array socket option, writes to buf, returns option size in bytes, e.g. + // `size_t optsize = socket.get(zmq::sockopt::routing_id, zmq::buffer(id))` + template + ZMQ_NODISCARD size_t get(sockopt::array_option, + mutable_buffer buf) const + { + size_t size = buf.size(); + get_option(Opt, buf.data(), &size); + return size; + } + + // Get array socket option as string (initializes the string buffer size to init_size) e.g. + // `auto s = socket.get(zmq::sockopt::routing_id)` + // Note: removes the null character from null-terminated string options, + // i.e. the string size excludes the null character. + template + ZMQ_NODISCARD std::string get(sockopt::array_option, + size_t init_size = 1024) const + { + if ZMQ_CONSTEXPR_IF (NullTerm == 2) { + if (init_size == 1024) { + init_size = 41; // get as Z85 string + } + } + std::string str(init_size, '\0'); + size_t size = get(sockopt::array_option{}, buffer(str)); + if ZMQ_CONSTEXPR_IF (NullTerm == 1) { + if (size > 0) { + assert(str[size - 1] == '\0'); + --size; + } + } else if ZMQ_CONSTEXPR_IF (NullTerm == 2) { + assert(size == 32 || size == 41); + if (size == 41) { + assert(str[size - 1] == '\0'); + --size; + } + } + str.resize(size); + return str; + } +#endif + + void bind(std::string const &addr) { bind(addr.c_str()); } + + void bind(const char *addr_) + { + int rc = zmq_bind(_handle, addr_); + if (rc != 0) + throw error_t(); + } + + void unbind(std::string const &addr) { unbind(addr.c_str()); } + + void unbind(const char *addr_) + { + int rc = zmq_unbind(_handle, addr_); + if (rc != 0) + throw error_t(); + } + + void connect(std::string const &addr) { connect(addr.c_str()); } + + void connect(const char *addr_) + { + int rc = zmq_connect(_handle, addr_); + if (rc != 0) + throw error_t(); + } + + void disconnect(std::string const &addr) { disconnect(addr.c_str()); } + + void disconnect(const char *addr_) + { + int rc = zmq_disconnect(_handle, addr_); + if (rc != 0) + throw error_t(); + } + + ZMQ_DEPRECATED("from 4.7.1, use handle() != nullptr or operator bool") + bool connected() const ZMQ_NOTHROW { return (_handle != ZMQ_NULLPTR); } + + ZMQ_CPP11_DEPRECATED("from 4.3.1, use send taking a const_buffer and send_flags") + size_t send(const void *buf_, size_t len_, int flags_ = 0) + { + int nbytes = zmq_send(_handle, buf_, len_, flags_); + if (nbytes >= 0) + return static_cast(nbytes); + if (zmq_errno() == EAGAIN) + return 0; + throw error_t(); + } + + ZMQ_CPP11_DEPRECATED("from 4.3.1, use send taking message_t and send_flags") + bool send(message_t &msg_, + int flags_ = 0) // default until removed + { + int nbytes = zmq_msg_send(msg_.handle(), _handle, flags_); + if (nbytes >= 0) + return true; + if (zmq_errno() == EAGAIN) + return false; + throw error_t(); + } + + template + ZMQ_CPP11_DEPRECATED( + "from 4.4.1, use send taking message_t or buffer (for contiguous " + "ranges), and send_flags") + bool send(T first, T last, int flags_ = 0) + { + zmq::message_t msg(first, last); + int nbytes = zmq_msg_send(msg.handle(), _handle, flags_); + if (nbytes >= 0) + return true; + if (zmq_errno() == EAGAIN) + return false; + throw error_t(); + } + +#ifdef ZMQ_HAS_RVALUE_REFS + ZMQ_CPP11_DEPRECATED("from 4.3.1, use send taking message_t and send_flags") + bool send(message_t &&msg_, + int flags_ = 0) // default until removed + { +#ifdef ZMQ_CPP11 + return send(msg_, static_cast(flags_)).has_value(); +#else + return send(msg_, flags_); +#endif + } +#endif + +#ifdef ZMQ_CPP11 + send_result_t send(const_buffer buf, send_flags flags = send_flags::none) + { + const int nbytes = + zmq_send(_handle, buf.data(), buf.size(), static_cast(flags)); + if (nbytes >= 0) + return static_cast(nbytes); + if (zmq_errno() == EAGAIN) + return {}; + throw error_t(); + } + + send_result_t send(message_t &msg, send_flags flags) + { + int nbytes = zmq_msg_send(msg.handle(), _handle, static_cast(flags)); + if (nbytes >= 0) + return static_cast(nbytes); + if (zmq_errno() == EAGAIN) + return {}; + throw error_t(); + } + + send_result_t send(message_t &&msg, send_flags flags) + { + return send(msg, flags); + } +#endif + + ZMQ_CPP11_DEPRECATED( + "from 4.3.1, use recv taking a mutable_buffer and recv_flags") + size_t recv(void *buf_, size_t len_, int flags_ = 0) + { + int nbytes = zmq_recv(_handle, buf_, len_, flags_); + if (nbytes >= 0) + return static_cast(nbytes); + if (zmq_errno() == EAGAIN) + return 0; + throw error_t(); + } + + ZMQ_CPP11_DEPRECATED( + "from 4.3.1, use recv taking a reference to message_t and recv_flags") + bool recv(message_t *msg_, int flags_ = 0) + { + int nbytes = zmq_msg_recv(msg_->handle(), _handle, flags_); + if (nbytes >= 0) + return true; + if (zmq_errno() == EAGAIN) + return false; + throw error_t(); + } + +#ifdef ZMQ_CPP11 + ZMQ_NODISCARD + recv_buffer_result_t recv(mutable_buffer buf, + recv_flags flags = recv_flags::none) + { + const int nbytes = + zmq_recv(_handle, buf.data(), buf.size(), static_cast(flags)); + if (nbytes >= 0) { + return recv_buffer_size{ + (std::min)(static_cast(nbytes), buf.size()), + static_cast(nbytes)}; + } + if (zmq_errno() == EAGAIN) + return {}; + throw error_t(); + } + + ZMQ_NODISCARD + recv_result_t recv(message_t &msg, recv_flags flags = recv_flags::none) + { + const int nbytes = + zmq_msg_recv(msg.handle(), _handle, static_cast(flags)); + if (nbytes >= 0) { + assert(msg.size() == static_cast(nbytes)); + return static_cast(nbytes); + } + if (zmq_errno() == EAGAIN) + return {}; + throw error_t(); + } +#endif + +#if defined(ZMQ_BUILD_DRAFT_API) && ZMQ_VERSION >= ZMQ_MAKE_VERSION(4, 2, 0) + void join(const char *group) + { + int rc = zmq_join(_handle, group); + if (rc != 0) + throw error_t(); + } + + void leave(const char *group) + { + int rc = zmq_leave(_handle, group); + if (rc != 0) + throw error_t(); + } +#endif + + ZMQ_NODISCARD void *handle() ZMQ_NOTHROW { return _handle; } + ZMQ_NODISCARD const void *handle() const ZMQ_NOTHROW { return _handle; } + + ZMQ_EXPLICIT operator bool() const ZMQ_NOTHROW { return _handle != ZMQ_NULLPTR; } + // note: non-const operator bool can be removed once + // operator void* is removed from socket_t + ZMQ_EXPLICIT operator bool() ZMQ_NOTHROW { return _handle != ZMQ_NULLPTR; } + + protected: + void *_handle; + + private: + void set_option(int option_, const void *optval_, size_t optvallen_) + { + int rc = zmq_setsockopt(_handle, option_, optval_, optvallen_); + if (rc != 0) + throw error_t(); + } + + void get_option(int option_, void *optval_, size_t *optvallen_) const + { + int rc = zmq_getsockopt(_handle, option_, optval_, optvallen_); + if (rc != 0) + throw error_t(); + } +}; +} // namespace detail + +struct from_handle_t +{ + struct _private + { + }; // disabling use other than with from_handle + ZMQ_CONSTEXPR_FN ZMQ_EXPLICIT from_handle_t(_private /*p*/) ZMQ_NOTHROW {} +}; + +ZMQ_CONSTEXPR_VAR from_handle_t from_handle = + from_handle_t(from_handle_t::_private()); + +// A non-owning nullable reference to a socket. +// The reference is invalidated on socket close or destruction. +class socket_ref : public detail::socket_base +{ + public: + socket_ref() ZMQ_NOTHROW : detail::socket_base() {} +#ifdef ZMQ_CPP11 + socket_ref(std::nullptr_t) ZMQ_NOTHROW : detail::socket_base() {} +#endif + socket_ref(from_handle_t /*fh*/, void *handle) ZMQ_NOTHROW + : detail::socket_base(handle) + { + } +}; + +#ifdef ZMQ_CPP11 +inline bool operator==(socket_ref sr, std::nullptr_t /*p*/) ZMQ_NOTHROW +{ + return sr.handle() == nullptr; +} +inline bool operator==(std::nullptr_t /*p*/, socket_ref sr) ZMQ_NOTHROW +{ + return sr.handle() == nullptr; +} +inline bool operator!=(socket_ref sr, std::nullptr_t /*p*/) ZMQ_NOTHROW +{ + return !(sr == nullptr); +} +inline bool operator!=(std::nullptr_t /*p*/, socket_ref sr) ZMQ_NOTHROW +{ + return !(sr == nullptr); +} +#endif + +inline bool operator==(const detail::socket_base& a, const detail::socket_base& b) ZMQ_NOTHROW +{ + return std::equal_to()(a.handle(), b.handle()); +} +inline bool operator!=(const detail::socket_base& a, const detail::socket_base& b) ZMQ_NOTHROW +{ + return !(a == b); +} +inline bool operator<(const detail::socket_base& a, const detail::socket_base& b) ZMQ_NOTHROW +{ + return std::less()(a.handle(), b.handle()); +} +inline bool operator>(const detail::socket_base& a, const detail::socket_base& b) ZMQ_NOTHROW +{ + return b < a; +} +inline bool operator<=(const detail::socket_base& a, const detail::socket_base& b) ZMQ_NOTHROW +{ + return !(a > b); +} +inline bool operator>=(const detail::socket_base& a, const detail::socket_base& b) ZMQ_NOTHROW +{ + return !(a < b); +} + +} // namespace zmq + +#ifdef ZMQ_CPP11 +namespace std +{ +template<> struct hash +{ + size_t operator()(zmq::socket_ref sr) const ZMQ_NOTHROW + { + return hash()(sr.handle()); + } +}; +} // namespace std +#endif + +namespace zmq +{ +class socket_t : public detail::socket_base +{ + friend class monitor_t; + + public: + socket_t() ZMQ_NOTHROW : detail::socket_base(ZMQ_NULLPTR), ctxptr(ZMQ_NULLPTR) {} + + socket_t(context_t &context_, int type_) : + detail::socket_base(zmq_socket(context_.handle(), type_)), + ctxptr(context_.handle()) + { + if (_handle == ZMQ_NULLPTR) + throw error_t(); + } + +#ifdef ZMQ_CPP11 + socket_t(context_t &context_, socket_type type_) : + socket_t(context_, static_cast(type_)) + { + } +#endif + +#ifdef ZMQ_HAS_RVALUE_REFS + socket_t(socket_t &&rhs) ZMQ_NOTHROW : detail::socket_base(rhs._handle), + ctxptr(rhs.ctxptr) + { + rhs._handle = ZMQ_NULLPTR; + rhs.ctxptr = ZMQ_NULLPTR; + } + socket_t &operator=(socket_t &&rhs) ZMQ_NOTHROW + { + close(); + std::swap(_handle, rhs._handle); + std::swap(ctxptr, rhs.ctxptr); + return *this; + } +#endif + + ~socket_t() ZMQ_NOTHROW { close(); } + + operator void *() ZMQ_NOTHROW { return _handle; } + + operator void const *() const ZMQ_NOTHROW { return _handle; } + + void close() ZMQ_NOTHROW + { + if (_handle == ZMQ_NULLPTR) + // already closed + return; + int rc = zmq_close(_handle); + ZMQ_ASSERT(rc == 0); + _handle = ZMQ_NULLPTR; + ctxptr = ZMQ_NULLPTR; + } + + void swap(socket_t &other) ZMQ_NOTHROW + { + std::swap(_handle, other._handle); + std::swap(ctxptr, other.ctxptr); + } + + operator socket_ref() ZMQ_NOTHROW { return socket_ref(from_handle, _handle); } + + private: + void *ctxptr; + + socket_t(const socket_t &) ZMQ_DELETED_FUNCTION; + void operator=(const socket_t &) ZMQ_DELETED_FUNCTION; + + // used by monitor_t + socket_t(void *context_, int type_) : + detail::socket_base(zmq_socket(context_, type_)), ctxptr(context_) + { + if (_handle == ZMQ_NULLPTR) + throw error_t(); + if (ctxptr == ZMQ_NULLPTR) + throw error_t(); + } +}; + +inline void swap(socket_t &a, socket_t &b) ZMQ_NOTHROW +{ + a.swap(b); +} + +ZMQ_DEPRECATED("from 4.3.1, use proxy taking socket_t objects") +inline void proxy(void *frontend, void *backend, void *capture) +{ + int rc = zmq_proxy(frontend, backend, capture); + if (rc != 0) + throw error_t(); +} + +inline void +proxy(socket_ref frontend, socket_ref backend, socket_ref capture = socket_ref()) +{ + int rc = zmq_proxy(frontend.handle(), backend.handle(), capture.handle()); + if (rc != 0) + throw error_t(); +} + +#ifdef ZMQ_HAS_PROXY_STEERABLE +ZMQ_DEPRECATED("from 4.3.1, use proxy_steerable taking socket_t objects") +inline void +proxy_steerable(void *frontend, void *backend, void *capture, void *control) +{ + int rc = zmq_proxy_steerable(frontend, backend, capture, control); + if (rc != 0) + throw error_t(); +} + +inline void proxy_steerable(socket_ref frontend, + socket_ref backend, + socket_ref capture, + socket_ref control) +{ + int rc = zmq_proxy_steerable(frontend.handle(), backend.handle(), + capture.handle(), control.handle()); + if (rc != 0) + throw error_t(); +} +#endif + +class monitor_t +{ + public: + monitor_t() : _socket(), _monitor_socket() {} + + virtual ~monitor_t() { close(); } + +#ifdef ZMQ_HAS_RVALUE_REFS + monitor_t(monitor_t &&rhs) ZMQ_NOTHROW : _socket(), _monitor_socket() + { + std::swap(_socket, rhs._socket); + std::swap(_monitor_socket, rhs._monitor_socket); + } + + monitor_t &operator=(monitor_t &&rhs) ZMQ_NOTHROW + { + close(); + _socket = socket_ref(); + std::swap(_socket, rhs._socket); + std::swap(_monitor_socket, rhs._monitor_socket); + return *this; + } +#endif + + + void + monitor(socket_t &socket, std::string const &addr, int events = ZMQ_EVENT_ALL) + { + monitor(socket, addr.c_str(), events); + } + + void monitor(socket_t &socket, const char *addr_, int events = ZMQ_EVENT_ALL) + { + init(socket, addr_, events); + while (true) { + check_event(-1); + } + } + + void init(socket_t &socket, std::string const &addr, int events = ZMQ_EVENT_ALL) + { + init(socket, addr.c_str(), events); + } + + void init(socket_t &socket, const char *addr_, int events = ZMQ_EVENT_ALL) + { + int rc = zmq_socket_monitor(socket.handle(), addr_, events); + if (rc != 0) + throw error_t(); + + _socket = socket; + _monitor_socket = socket_t(socket.ctxptr, ZMQ_PAIR); + _monitor_socket.connect(addr_); + + on_monitor_started(); + } + + bool check_event(int timeout = 0) + { + assert(_monitor_socket); + + zmq::pollitem_t items[] = { + {_monitor_socket.handle(), 0, ZMQ_POLLIN, 0}, + }; + + #ifdef ZMQ_CPP11 + zmq::poll(&items[0], 1, std::chrono::milliseconds(timeout)); + #else + zmq::poll(&items[0], 1, timeout); + #endif + + return process_event(items[0].revents); + } + +#ifdef ZMQ_EVENT_MONITOR_STOPPED + void abort() + { + if (_socket) + zmq_socket_monitor(_socket.handle(), ZMQ_NULLPTR, 0); + + _socket = socket_ref(); + } + + virtual void on_monitor_stopped() {} +#endif + virtual void on_monitor_started() {} + virtual void on_event_connected(const zmq_event_t &event_, const char *addr_) + { + (void) event_; + (void) addr_; + } + virtual void on_event_connect_delayed(const zmq_event_t &event_, + const char *addr_) + { + (void) event_; + (void) addr_; + } + virtual void on_event_connect_retried(const zmq_event_t &event_, + const char *addr_) + { + (void) event_; + (void) addr_; + } + virtual void on_event_listening(const zmq_event_t &event_, const char *addr_) + { + (void) event_; + (void) addr_; + } + virtual void on_event_bind_failed(const zmq_event_t &event_, const char *addr_) + { + (void) event_; + (void) addr_; + } + virtual void on_event_accepted(const zmq_event_t &event_, const char *addr_) + { + (void) event_; + (void) addr_; + } + virtual void on_event_accept_failed(const zmq_event_t &event_, const char *addr_) + { + (void) event_; + (void) addr_; + } + virtual void on_event_closed(const zmq_event_t &event_, const char *addr_) + { + (void) event_; + (void) addr_; + } + virtual void on_event_close_failed(const zmq_event_t &event_, const char *addr_) + { + (void) event_; + (void) addr_; + } + virtual void on_event_disconnected(const zmq_event_t &event_, const char *addr_) + { + (void) event_; + (void) addr_; + } +#if ZMQ_VERSION >= ZMQ_MAKE_VERSION(4, 2, 3) + virtual void on_event_handshake_failed_no_detail(const zmq_event_t &event_, + const char *addr_) + { + (void) event_; + (void) addr_; + } + virtual void on_event_handshake_failed_protocol(const zmq_event_t &event_, + const char *addr_) + { + (void) event_; + (void) addr_; + } + virtual void on_event_handshake_failed_auth(const zmq_event_t &event_, + const char *addr_) + { + (void) event_; + (void) addr_; + } + virtual void on_event_handshake_succeeded(const zmq_event_t &event_, + const char *addr_) + { + (void) event_; + (void) addr_; + } +#elif ZMQ_VERSION >= ZMQ_MAKE_VERSION(4, 2, 1) + virtual void on_event_handshake_failed(const zmq_event_t &event_, + const char *addr_) + { + (void) event_; + (void) addr_; + } + virtual void on_event_handshake_succeed(const zmq_event_t &event_, + const char *addr_) + { + (void) event_; + (void) addr_; + } +#endif + virtual void on_event_unknown(const zmq_event_t &event_, const char *addr_) + { + (void) event_; + (void) addr_; + } + + protected: + bool process_event(short events) + { + zmq::message_t eventMsg; + + if (events & ZMQ_POLLIN) { + int rc = zmq_msg_recv(eventMsg.handle(), _monitor_socket.handle(), 0); + if (rc == -1 && zmq_errno() == ETERM) + return false; + assert(rc != -1); + + } else { + return false; + } + +#if ZMQ_VERSION_MAJOR >= 4 + const char *data = static_cast(eventMsg.data()); + zmq_event_t msgEvent; + memcpy(&msgEvent.event, data, sizeof(uint16_t)); + data += sizeof(uint16_t); + memcpy(&msgEvent.value, data, sizeof(int32_t)); + zmq_event_t *event = &msgEvent; +#else + zmq_event_t *event = static_cast(eventMsg.data()); +#endif + +#ifdef ZMQ_NEW_MONITOR_EVENT_LAYOUT + zmq::message_t addrMsg; + int rc = zmq_msg_recv(addrMsg.handle(), _monitor_socket.handle(), 0); + if (rc == -1 && zmq_errno() == ETERM) { + return false; + } + + assert(rc != -1); + std::string address = addrMsg.to_string(); +#else + // Bit of a hack, but all events in the zmq_event_t union have the same layout so this will work for all event types. + std::string address = event->data.connected.addr; +#endif + +#ifdef ZMQ_EVENT_MONITOR_STOPPED + if (event->event == ZMQ_EVENT_MONITOR_STOPPED) { + on_monitor_stopped(); + return false; + } + +#endif + + switch (event->event) { + case ZMQ_EVENT_CONNECTED: + on_event_connected(*event, address.c_str()); + break; + case ZMQ_EVENT_CONNECT_DELAYED: + on_event_connect_delayed(*event, address.c_str()); + break; + case ZMQ_EVENT_CONNECT_RETRIED: + on_event_connect_retried(*event, address.c_str()); + break; + case ZMQ_EVENT_LISTENING: + on_event_listening(*event, address.c_str()); + break; + case ZMQ_EVENT_BIND_FAILED: + on_event_bind_failed(*event, address.c_str()); + break; + case ZMQ_EVENT_ACCEPTED: + on_event_accepted(*event, address.c_str()); + break; + case ZMQ_EVENT_ACCEPT_FAILED: + on_event_accept_failed(*event, address.c_str()); + break; + case ZMQ_EVENT_CLOSED: + on_event_closed(*event, address.c_str()); + break; + case ZMQ_EVENT_CLOSE_FAILED: + on_event_close_failed(*event, address.c_str()); + break; + case ZMQ_EVENT_DISCONNECTED: + on_event_disconnected(*event, address.c_str()); + break; +#if ZMQ_VERSION >= ZMQ_MAKE_VERSION(4, 3, 0) || (defined(ZMQ_BUILD_DRAFT_API) && ZMQ_VERSION >= ZMQ_MAKE_VERSION(4, 2, 3)) + case ZMQ_EVENT_HANDSHAKE_FAILED_NO_DETAIL: + on_event_handshake_failed_no_detail(*event, address.c_str()); + break; + case ZMQ_EVENT_HANDSHAKE_FAILED_PROTOCOL: + on_event_handshake_failed_protocol(*event, address.c_str()); + break; + case ZMQ_EVENT_HANDSHAKE_FAILED_AUTH: + on_event_handshake_failed_auth(*event, address.c_str()); + break; + case ZMQ_EVENT_HANDSHAKE_SUCCEEDED: + on_event_handshake_succeeded(*event, address.c_str()); + break; +#elif defined(ZMQ_BUILD_DRAFT_API) && ZMQ_VERSION >= ZMQ_MAKE_VERSION(4, 2, 1) + case ZMQ_EVENT_HANDSHAKE_FAILED: + on_event_handshake_failed(*event, address.c_str()); + break; + case ZMQ_EVENT_HANDSHAKE_SUCCEED: + on_event_handshake_succeed(*event, address.c_str()); + break; +#endif + default: + on_event_unknown(*event, address.c_str()); + break; + } + + return true; + } + + socket_ref monitor_socket() {return _monitor_socket;} + + private: + monitor_t(const monitor_t &) ZMQ_DELETED_FUNCTION; + void operator=(const monitor_t &) ZMQ_DELETED_FUNCTION; + + socket_ref _socket; + socket_t _monitor_socket; + + void close() ZMQ_NOTHROW + { + if (_socket) + zmq_socket_monitor(_socket.handle(), ZMQ_NULLPTR, 0); + _monitor_socket.close(); + } +}; + +#if defined(ZMQ_BUILD_DRAFT_API) && defined(ZMQ_CPP11) && defined(ZMQ_HAVE_POLLER) + +// polling events +enum class event_flags : short +{ + none = 0, + pollin = ZMQ_POLLIN, + pollout = ZMQ_POLLOUT, + pollerr = ZMQ_POLLERR, + pollpri = ZMQ_POLLPRI +}; + +constexpr event_flags operator|(event_flags a, event_flags b) noexcept +{ + return detail::enum_bit_or(a, b); +} +constexpr event_flags operator&(event_flags a, event_flags b) noexcept +{ + return detail::enum_bit_and(a, b); +} +constexpr event_flags operator^(event_flags a, event_flags b) noexcept +{ + return detail::enum_bit_xor(a, b); +} +constexpr event_flags operator~(event_flags a) noexcept +{ + return detail::enum_bit_not(a); +} + +struct no_user_data; + +// layout compatible with zmq_poller_event_t +template struct poller_event +{ + socket_ref socket; + ::zmq::fd_t fd; + T *user_data; + event_flags events; +}; + +template class poller_t +{ + public: + using event_type = poller_event; + + poller_t() : poller_ptr(zmq_poller_new()) + { + if (!poller_ptr) + throw error_t(); + } + + template< + typename Dummy = void, + typename = + typename std::enable_if::value, Dummy>::type> + void add(zmq::socket_ref socket, event_flags events, T *user_data) + { + add_impl(socket, events, user_data); + } + + void add(zmq::socket_ref socket, event_flags events) + { + add_impl(socket, events, nullptr); + } + + template< + typename Dummy = void, + typename = + typename std::enable_if::value, Dummy>::type> + void add(fd_t fd, event_flags events, T *user_data) + { + add_impl(fd, events, user_data); + } + + void add(fd_t fd, event_flags events) { add_impl(fd, events, nullptr); } + + void remove(zmq::socket_ref socket) + { + if (0 != zmq_poller_remove(poller_ptr.get(), socket.handle())) { + throw error_t(); + } + } + + void remove(fd_t fd) + { + if (0 != zmq_poller_remove_fd(poller_ptr.get(), fd)) { + throw error_t(); + } + } + + void modify(zmq::socket_ref socket, event_flags events) + { + if (0 + != zmq_poller_modify(poller_ptr.get(), socket.handle(), + static_cast(events))) { + throw error_t(); + } + } + + void modify(fd_t fd, event_flags events) + { + if (0 + != zmq_poller_modify_fd(poller_ptr.get(), fd, + static_cast(events))) { + throw error_t(); + } + } + + template + size_t wait_all(Sequence &poller_events, + const std::chrono::milliseconds timeout) + { + static_assert(std::is_same::value, + "Sequence::value_type must be of poller_t::event_type"); + int rc = zmq_poller_wait_all( + poller_ptr.get(), + reinterpret_cast(poller_events.data()), + static_cast(poller_events.size()), + static_cast(timeout.count())); + if (rc > 0) + return static_cast(rc); + +#if ZMQ_VERSION >= ZMQ_MAKE_VERSION(4, 2, 3) + if (zmq_errno() == EAGAIN) +#else + if (zmq_errno() == ETIMEDOUT) +#endif + return 0; + + throw error_t(); + } + +#if ZMQ_VERSION >= ZMQ_MAKE_VERSION(4, 3, 3) + size_t size() const noexcept + { + int rc = zmq_poller_size(const_cast(poller_ptr.get())); + ZMQ_ASSERT(rc >= 0); + return static_cast(std::max(rc, 0)); + } +#endif + + private: + struct destroy_poller_t + { + void operator()(void *ptr) noexcept + { + int rc = zmq_poller_destroy(&ptr); + ZMQ_ASSERT(rc == 0); + } + }; + + std::unique_ptr poller_ptr; + + void add_impl(zmq::socket_ref socket, event_flags events, T *user_data) + { + if (0 + != zmq_poller_add(poller_ptr.get(), socket.handle(), user_data, + static_cast(events))) { + throw error_t(); + } + } + + void add_impl(fd_t fd, event_flags events, T *user_data) + { + if (0 + != zmq_poller_add_fd(poller_ptr.get(), fd, user_data, + static_cast(events))) { + throw error_t(); + } + } +}; +#endif // defined(ZMQ_BUILD_DRAFT_API) && defined(ZMQ_CPP11) && defined(ZMQ_HAVE_POLLER) + +inline std::ostream &operator<<(std::ostream &os, const message_t &msg) +{ + return os << msg.str(); +} + +} // namespace zmq + +#endif // __ZMQ_HPP_INCLUDED__ \ No newline at end of file diff --git a/include/zmq_addon.hpp b/include/zmq_addon.hpp new file mode 100644 index 00000000..291c2b93 --- /dev/null +++ b/include/zmq_addon.hpp @@ -0,0 +1,848 @@ +/* + Copyright (c) 2016-2017 ZeroMQ community + Copyright (c) 2016 VOCA AS / Harald NΓΈkland + + 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. +*/ + +#ifndef __ZMQ_ADDON_HPP_INCLUDED__ +#define __ZMQ_ADDON_HPP_INCLUDED__ + +#include "zmq.hpp" + +#include +#include +#include +#include +#ifdef ZMQ_CPP11 +#include +#include +#include + +namespace zmq +{ + // socket ref or native file descriptor for poller + class poller_ref_t + { + public: + enum RefType + { + RT_SOCKET, + RT_FD + }; + + poller_ref_t() : poller_ref_t(socket_ref{}) + {} + + poller_ref_t(const zmq::socket_ref& socket) : data{RT_SOCKET, socket, {}} + {} + + poller_ref_t(zmq::fd_t fd) : data{RT_FD, {}, fd} + {} + + size_t hash() const ZMQ_NOTHROW + { + std::size_t h = 0; + hash_combine(h, std::get<0>(data)); + hash_combine(h, std::get<1>(data)); + hash_combine(h, std::get<2>(data)); + return h; + } + + bool operator == (const poller_ref_t& o) const ZMQ_NOTHROW + { + return data == o.data; + } + + private: + template + static void hash_combine(std::size_t& seed, const T& v) ZMQ_NOTHROW + { + std::hash hasher; + seed ^= hasher(v) + 0x9e3779b9 + (seed<<6) + (seed>>2); + } + + std::tuple data; + + }; // class poller_ref_t + +} // namespace zmq + +// std::hash<> specialization for std::unordered_map +template <> struct std::hash +{ + size_t operator()(const zmq::poller_ref_t& ref) const ZMQ_NOTHROW + { + return ref.hash(); + } +}; +#endif // ZMQ_CPP11 + +namespace zmq +{ +#ifdef ZMQ_CPP11 + +namespace detail +{ +template +recv_result_t +recv_multipart_n(socket_ref s, OutputIt out, size_t n, recv_flags flags) +{ + size_t msg_count = 0; + message_t msg; + while (true) { + if ZMQ_CONSTEXPR_IF (CheckN) { + if (msg_count >= n) + throw std::runtime_error( + "Too many message parts in recv_multipart_n"); + } + if (!s.recv(msg, flags)) { + // zmq ensures atomic delivery of messages + assert(msg_count == 0); + return {}; + } + ++msg_count; + const bool more = msg.more(); + *out++ = std::move(msg); + if (!more) + break; + } + return msg_count; +} + +inline bool is_little_endian() +{ + const uint16_t i = 0x01; + return *reinterpret_cast(&i) == 0x01; +} + +inline void write_network_order(unsigned char *buf, const uint32_t value) +{ + if (is_little_endian()) { + ZMQ_CONSTEXPR_VAR uint32_t mask = (std::numeric_limits::max)(); + *buf++ = static_cast((value >> 24) & mask); + *buf++ = static_cast((value >> 16) & mask); + *buf++ = static_cast((value >> 8) & mask); + *buf++ = static_cast(value & mask); + } else { + std::memcpy(buf, &value, sizeof(value)); + } +} + +inline uint32_t read_u32_network_order(const unsigned char *buf) +{ + if (is_little_endian()) { + return (static_cast(buf[0]) << 24) + + (static_cast(buf[1]) << 16) + + (static_cast(buf[2]) << 8) + + static_cast(buf[3]); + } else { + uint32_t value; + std::memcpy(&value, buf, sizeof(value)); + return value; + } +} +} // namespace detail + +/* Receive a multipart message. + + Writes the zmq::message_t objects to OutputIterator out. + The out iterator must handle an unspecified number of writes, + e.g. by using std::back_inserter. + + Returns: the number of messages received or nullopt (on EAGAIN). + Throws: if recv throws. Any exceptions thrown + by the out iterator will be propagated and the message + may have been only partially received with pending + message parts. It is adviced to close this socket in that event. +*/ +template +ZMQ_NODISCARD recv_result_t recv_multipart(socket_ref s, + OutputIt out, + recv_flags flags = recv_flags::none) +{ + return detail::recv_multipart_n(s, std::move(out), 0, flags); +} + +/* Receive a multipart message. + + Writes at most n zmq::message_t objects to OutputIterator out. + If the number of message parts of the incoming message exceeds n + then an exception will be thrown. + + Returns: the number of messages received or nullopt (on EAGAIN). + Throws: if recv throws. Throws std::runtime_error if the number + of message parts exceeds n (exactly n messages will have been written + to out). Any exceptions thrown + by the out iterator will be propagated and the message + may have been only partially received with pending + message parts. It is adviced to close this socket in that event. +*/ +template +ZMQ_NODISCARD recv_result_t recv_multipart_n(socket_ref s, + OutputIt out, + size_t n, + recv_flags flags = recv_flags::none) +{ + return detail::recv_multipart_n(s, std::move(out), n, flags); +} + +/* Send a multipart message. + + The range must be a ForwardRange of zmq::message_t, + zmq::const_buffer or zmq::mutable_buffer. + The flags may be zmq::send_flags::sndmore if there are + more message parts to be sent after the call to this function. + + Returns: the number of messages sent (exactly msgs.size()) or nullopt (on EAGAIN). + Throws: if send throws. Any exceptions thrown + by the msgs range will be propagated and the message + may have been only partially sent. It is adviced to close this socket in that event. +*/ +template::value + && (std::is_same, message_t>::value + || detail::is_buffer>::value)>::type +#endif + > +send_result_t +send_multipart(socket_ref s, Range &&msgs, send_flags flags = send_flags::none) +{ + using std::begin; + using std::end; + auto it = begin(msgs); + const auto end_it = end(msgs); + size_t msg_count = 0; + while (it != end_it) { + const auto next = std::next(it); + const auto msg_flags = + flags | (next == end_it ? send_flags::none : send_flags::sndmore); + if (!s.send(*it, msg_flags)) { + // zmq ensures atomic delivery of messages + assert(it == begin(msgs)); + return {}; + } + ++msg_count; + it = next; + } + return msg_count; +} + +/* Encode a multipart message. + + The range must be a ForwardRange of zmq::message_t. A + zmq::multipart_t or STL container may be passed for encoding. + + Returns: a zmq::message_t holding the encoded multipart data. + + Throws: std::range_error is thrown if the size of any single part + can not fit in an unsigned 32 bit integer. + + The encoding is compatible with that used by the CZMQ function + zmsg_encode(), see https://rfc.zeromq.org/spec/50/. + Each part consists of a size followed by the data. + These are placed contiguously into the output message. A part of + size less than 255 bytes will have a single byte size value. + Larger parts will have a five byte size value with the first byte + set to 0xFF and the remaining four bytes holding the size of the + part's data. +*/ +template::value + && (std::is_same, message_t>::value + || detail::is_buffer>::value)>::type +#endif + > +message_t encode(const Range &parts) +{ + size_t mmsg_size = 0; + + // First pass check sizes + for (const auto &part : parts) { + const size_t part_size = part.size(); + if (part_size > (std::numeric_limits::max)()) { + // Size value must fit into uint32_t. + throw std::range_error("Invalid size, message part too large"); + } + const size_t count_size = + part_size < (std::numeric_limits::max)() ? 1 : 5; + mmsg_size += part_size + count_size; + } + + message_t encoded(mmsg_size); + unsigned char *buf = encoded.data(); + for (const auto &part : parts) { + const uint32_t part_size = static_cast(part.size()); + const unsigned char *part_data = + static_cast(part.data()); + + if (part_size < (std::numeric_limits::max)()) { + // small part + *buf++ = (unsigned char) part_size; + } else { + // big part + *buf++ = (std::numeric_limits::max)(); + detail::write_network_order(buf, part_size); + buf += sizeof(part_size); + } + std::memcpy(buf, part_data, part_size); + buf += part_size; + } + + assert(static_cast(buf - encoded.data()) == mmsg_size); + return encoded; +} + +/* Decode an encoded message to multiple parts. + + The given output iterator must be a ForwardIterator to a container + holding zmq::message_t such as a zmq::multipart_t or various STL + containers. + + Returns the ForwardIterator advanced once past the last decoded + part. + + Throws: a std::out_of_range is thrown if the encoded part sizes + lead to exceeding the message data bounds. + + The decoding assumes the message is encoded in the manner + performed by zmq::encode(), see https://rfc.zeromq.org/spec/50/. + */ +template OutputIt decode(const message_t &encoded, OutputIt out) +{ + const unsigned char *source = encoded.data(); + const unsigned char *const limit = source + encoded.size(); + + while (source < limit) { + size_t part_size = *source++; + if (part_size == (std::numeric_limits::max)()) { + if (static_cast(limit - source) < sizeof(uint32_t)) { + throw std::out_of_range( + "Malformed encoding, overflow in reading size"); + } + part_size = detail::read_u32_network_order(source); + // the part size is allowed to be less than 0xFF + source += sizeof(uint32_t); + } + + if (static_cast(limit - source) < part_size) { + throw std::out_of_range("Malformed encoding, overflow in reading part"); + } + *out = message_t(source, part_size); + ++out; + source += part_size; + } + + assert(source == limit); + return out; +} + +#endif + + +#ifdef ZMQ_HAS_RVALUE_REFS + +/* + This class handles multipart messaging. It is the C++ equivalent of zmsg.h, + which is part of CZMQ (the high-level C binding). Furthermore, it is a major + improvement compared to zmsg.hpp, which is part of the examples in the ØMQ + Guide. Unnecessary copying is avoided by using move semantics to efficiently + add/remove parts. +*/ +class multipart_t +{ + private: + std::deque m_parts; + + public: + typedef std::deque::value_type value_type; + + typedef std::deque::iterator iterator; + typedef std::deque::const_iterator const_iterator; + + typedef std::deque::reverse_iterator reverse_iterator; + typedef std::deque::const_reverse_iterator const_reverse_iterator; + + // Default constructor + multipart_t() {} + + // Construct from socket receive + multipart_t(socket_ref socket) { recv(socket); } + + // Construct from memory block + multipart_t(const void *src, size_t size) { addmem(src, size); } + + // Construct from string + multipart_t(const std::string &string) { addstr(string); } + + // Construct from message part + multipart_t(message_t &&message) { add(std::move(message)); } + + // Move constructor + multipart_t(multipart_t &&other) ZMQ_NOTHROW { m_parts = std::move(other.m_parts); } + + // Move assignment operator + multipart_t &operator=(multipart_t &&other) ZMQ_NOTHROW + { + m_parts = std::move(other.m_parts); + return *this; + } + + // Destructor + virtual ~multipart_t() { clear(); } + + message_t &operator[](size_t n) { return m_parts[n]; } + + const message_t &operator[](size_t n) const { return m_parts[n]; } + + message_t &at(size_t n) { return m_parts.at(n); } + + const message_t &at(size_t n) const { return m_parts.at(n); } + + iterator begin() { return m_parts.begin(); } + + const_iterator begin() const { return m_parts.begin(); } + + const_iterator cbegin() const { return m_parts.cbegin(); } + + reverse_iterator rbegin() { return m_parts.rbegin(); } + + const_reverse_iterator rbegin() const { return m_parts.rbegin(); } + + iterator end() { return m_parts.end(); } + + const_iterator end() const { return m_parts.end(); } + + const_iterator cend() const { return m_parts.cend(); } + + reverse_iterator rend() { return m_parts.rend(); } + + const_reverse_iterator rend() const { return m_parts.rend(); } + + // Delete all parts + void clear() { m_parts.clear(); } + + // Get number of parts + size_t size() const { return m_parts.size(); } + + // Check if number of parts is zero + bool empty() const { return m_parts.empty(); } + + // Receive multipart message from socket + bool recv(socket_ref socket, int flags = 0) + { + clear(); + bool more = true; + while (more) { + message_t message; +#ifdef ZMQ_CPP11 + if (!socket.recv(message, static_cast(flags))) + return false; +#else + if (!socket.recv(&message, flags)) + return false; +#endif + more = message.more(); + add(std::move(message)); + } + return true; + } + + // Send multipart message to socket + bool send(socket_ref socket, int flags = 0) + { + flags &= ~(ZMQ_SNDMORE); + bool more = size() > 0; + while (more) { + message_t message = pop(); + more = size() > 0; +#ifdef ZMQ_CPP11 + if (!socket.send(message, static_cast( + (more ? ZMQ_SNDMORE : 0) | flags))) + return false; +#else + if (!socket.send(message, (more ? ZMQ_SNDMORE : 0) | flags)) + return false; +#endif + } + clear(); + return true; + } + + // Concatenate other multipart to front + void prepend(multipart_t &&other) + { + while (!other.empty()) + push(other.remove()); + } + + // Concatenate other multipart to back + void append(multipart_t &&other) + { + while (!other.empty()) + add(other.pop()); + } + + // Push memory block to front + void pushmem(const void *src, size_t size) + { + m_parts.push_front(message_t(src, size)); + } + + // Push memory block to back + void addmem(const void *src, size_t size) + { + m_parts.push_back(message_t(src, size)); + } + + // Push string to front + void pushstr(const std::string &string) + { + m_parts.push_front(message_t(string.data(), string.size())); + } + + // Push string to back + void addstr(const std::string &string) + { + m_parts.push_back(message_t(string.data(), string.size())); + } + + // Push type (fixed-size) to front + template void pushtyp(const T &type) + { + static_assert(!std::is_same::value, + "Use pushstr() instead of pushtyp()"); + m_parts.push_front(message_t(&type, sizeof(type))); + } + + // Push type (fixed-size) to back + template void addtyp(const T &type) + { + static_assert(!std::is_same::value, + "Use addstr() instead of addtyp()"); + m_parts.push_back(message_t(&type, sizeof(type))); + } + + // Push message part to front + void push(message_t &&message) { m_parts.push_front(std::move(message)); } + + // Push message part to back + void add(message_t &&message) { m_parts.push_back(std::move(message)); } + + // Alias to allow std::back_inserter() + void push_back(message_t &&message) { m_parts.push_back(std::move(message)); } + + // Pop string from front + std::string popstr() + { + std::string string(m_parts.front().data(), m_parts.front().size()); + m_parts.pop_front(); + return string; + } + + // Pop type (fixed-size) from front + template T poptyp() + { + static_assert(!std::is_same::value, + "Use popstr() instead of poptyp()"); + if (sizeof(T) != m_parts.front().size()) + throw std::runtime_error( + "Invalid type, size does not match the message size"); + T type = *m_parts.front().data(); + m_parts.pop_front(); + return type; + } + + // Pop message part from front + message_t pop() + { + message_t message = std::move(m_parts.front()); + m_parts.pop_front(); + return message; + } + + // Pop message part from back + message_t remove() + { + message_t message = std::move(m_parts.back()); + m_parts.pop_back(); + return message; + } + + // get message part from front + const message_t &front() { return m_parts.front(); } + + // get message part from back + const message_t &back() { return m_parts.back(); } + + // Get pointer to a specific message part + const message_t *peek(size_t index) const { return &m_parts[index]; } + + // Get a string copy of a specific message part + std::string peekstr(size_t index) const + { + std::string string(m_parts[index].data(), m_parts[index].size()); + return string; + } + + // Peek type (fixed-size) from front + template T peektyp(size_t index) const + { + static_assert(!std::is_same::value, + "Use peekstr() instead of peektyp()"); + if (sizeof(T) != m_parts[index].size()) + throw std::runtime_error( + "Invalid type, size does not match the message size"); + T type = *m_parts[index].data(); + return type; + } + + // Create multipart from type (fixed-size) + template static multipart_t create(const T &type) + { + multipart_t multipart; + multipart.addtyp(type); + return multipart; + } + + // Copy multipart + multipart_t clone() const + { + multipart_t multipart; + for (size_t i = 0; i < size(); i++) + multipart.addmem(m_parts[i].data(), m_parts[i].size()); + return multipart; + } + + // Dump content to string + std::string str() const + { + std::stringstream ss; + for (size_t i = 0; i < m_parts.size(); i++) { + const unsigned char *data = m_parts[i].data(); + size_t size = m_parts[i].size(); + + // Dump the message as text or binary + bool isText = true; + for (size_t j = 0; j < size; j++) { + if (data[j] < 32 || data[j] > 127) { + isText = false; + break; + } + } + ss << "\n[" << std::dec << std::setw(3) << std::setfill('0') << size + << "] "; + if (size >= 1000) { + ss << "... (too big to print)"; + continue; + } + for (size_t j = 0; j < size; j++) { + if (isText) + ss << static_cast(data[j]); + else + ss << std::hex << std::setw(2) << std::setfill('0') + << static_cast(data[j]); + } + } + return ss.str(); + } + + // Check if equal to other multipart + bool equal(const multipart_t *other) const ZMQ_NOTHROW + { + return *this == *other; + } + + bool operator==(const multipart_t &other) const ZMQ_NOTHROW + { + if (size() != other.size()) + return false; + for (size_t i = 0; i < size(); i++) + if (at(i) != other.at(i)) + return false; + return true; + } + + bool operator!=(const multipart_t &other) const ZMQ_NOTHROW + { + return !(*this == other); + } + +#ifdef ZMQ_CPP11 + + // Return single part message_t encoded from this multipart_t. + message_t encode() const { return zmq::encode(*this); } + + // Decode encoded message into multiple parts and append to self. + void decode_append(const message_t &encoded) + { + zmq::decode(encoded, std::back_inserter(*this)); + } + + // Return a new multipart_t containing the decoded message_t. + static multipart_t decode(const message_t &encoded) + { + multipart_t tmp; + zmq::decode(encoded, std::back_inserter(tmp)); + return tmp; + } + +#endif + + private: + // Disable implicit copying (moving is more efficient) + multipart_t(const multipart_t &other) ZMQ_DELETED_FUNCTION; + void operator=(const multipart_t &other) ZMQ_DELETED_FUNCTION; +}; // class multipart_t + +inline std::ostream &operator<<(std::ostream &os, const multipart_t &msg) +{ + return os << msg.str(); +} + +#endif // ZMQ_HAS_RVALUE_REFS + +#if defined(ZMQ_BUILD_DRAFT_API) && defined(ZMQ_CPP11) && defined(ZMQ_HAVE_POLLER) +class active_poller_t +{ + public: + active_poller_t() = default; + ~active_poller_t() = default; + + active_poller_t(const active_poller_t &) = delete; + active_poller_t &operator=(const active_poller_t &) = delete; + + active_poller_t(active_poller_t &&src) = default; + active_poller_t &operator=(active_poller_t &&src) = default; + + using handler_type = std::function; + + void add(zmq::socket_ref socket, event_flags events, handler_type handler) + { + const poller_ref_t ref{socket}; + + if (!handler) + throw std::invalid_argument("null handler in active_poller_t::add (socket)"); + auto ret = handlers.emplace( + ref, std::make_shared(std::move(handler))); + if (!ret.second) + throw error_t(EINVAL); // already added + try { + base_poller.add(socket, events, ret.first->second.get()); + need_rebuild = true; + } + catch (...) { + // rollback + handlers.erase(ref); + throw; + } + } + + void add(fd_t fd, event_flags events, handler_type handler) + { + const poller_ref_t ref{fd}; + + if (!handler) + throw std::invalid_argument("null handler in active_poller_t::add (fd)"); + auto ret = handlers.emplace( + ref, std::make_shared(std::move(handler))); + if (!ret.second) + throw error_t(EINVAL); // already added + try { + base_poller.add(fd, events, ret.first->second.get()); + need_rebuild = true; + } + catch (...) { + // rollback + handlers.erase(ref); + throw; + } + } + + void remove(zmq::socket_ref socket) + { + base_poller.remove(socket); + handlers.erase(socket); + need_rebuild = true; + } + + void remove(fd_t fd) + { + base_poller.remove(fd); + handlers.erase(fd); + need_rebuild = true; + } + + void modify(zmq::socket_ref socket, event_flags events) + { + base_poller.modify(socket, events); + } + + void modify(fd_t fd, event_flags events) + { + base_poller.modify(fd, events); + } + + size_t wait(std::chrono::milliseconds timeout) + { + if (need_rebuild) { + poller_events.resize(handlers.size()); + poller_handlers.clear(); + poller_handlers.reserve(handlers.size()); + for (const auto &handler : handlers) { + poller_handlers.push_back(handler.second); + } + need_rebuild = false; + } + const auto count = base_poller.wait_all(poller_events, timeout); + std::for_each(poller_events.begin(), + poller_events.begin() + static_cast(count), + [](decltype(base_poller)::event_type &event) { + assert(event.user_data != nullptr); + (*event.user_data)(event.events); + }); + return count; + } + + ZMQ_NODISCARD bool empty() const noexcept { return handlers.empty(); } + + size_t size() const noexcept { return handlers.size(); } + + private: + bool need_rebuild{false}; + + poller_t base_poller{}; + + std::unordered_map> handlers{}; + + std::vector poller_events{}; + std::vector> poller_handlers{}; +}; // class active_poller_t +#endif // defined(ZMQ_BUILD_DRAFT_API) && defined(ZMQ_CPP11) && defined(ZMQ_HAVE_POLLER) + + +} // namespace zmq + +#endif // __ZMQ_ADDON_HPP_INCLUDED__ \ No newline at end of file diff --git a/src/llama.cpp b/src/llama.cpp index 3fb8132f..af62c18f 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -1,3 +1,5 @@ +#include "zmq_addon.hpp" + #include "llama-impl.h" #include "llama-vocab.h" #include "llama-sampling.h" @@ -84,6 +86,7 @@ #include #include #include +#include #if defined(_MSC_VER) #pragma warning(disable: 4244 4267) // possible loss of data @@ -93,6 +96,22 @@ #define LLAMA_MAX_LAYERS 512 #define LLAMA_MAX_EXPERTS 160 // DeepSeekV2 +#define timer(name) auto _timer_##name = Timer(#name) + +struct Timer { + const char* name; + int64_t start_time; + bool enable_timer = false; + Timer(const char* name) : name(name), start_time(ggml_time_us()) {} + ~Timer() { + if (enable_timer) { + int64_t end_time = ggml_time_us(); + LLAMA_LOG_INFO("Time to run %s: %lld ms\n", name, (end_time - start_time)/1000); + } + } +}; + + // // helpers // @@ -2540,6 +2559,9 @@ struct llama_hparams { static_assert(std::is_trivially_copyable::value, "llama_hparams must be trivially copyable"); struct llama_cparams { + uint32_t n_world; + uint32_t rank; + uint32_t n_layer_window; uint32_t n_ctx; // context size used during inference uint32_t n_batch; uint32_t n_ubatch; @@ -2730,6 +2752,8 @@ struct llama_ubatch { llama_token * token; // [n_tokens] float * embd; // [n_embd, n_tokens] + float * backend_embd; // [n_embd, n_tokens] + float * out_embd; // [n_embd] llama_pos * pos; // [n_tokens] int32_t * n_seq_id; // [n_seqs] llama_seq_id ** seq_id; // [n_seqs] @@ -2957,6 +2981,8 @@ struct llama_sbatch { // buffers for the ubatch std::vector ubatch_token; std::vector ubatch_embd; + std::vector ubatch_backend_embd; + std::vector ubatch_out_embd; std::vector ubatch_pos; std::vector ubatch_n_seq_id; std::vector ubatch_seq_id; @@ -2975,6 +3001,8 @@ struct llama_sbatch { } ubatch_token.resize(!has_embd ? n_ubatch : 0); ubatch_embd.resize(has_embd ? n_embd * n_ubatch : 0); + ubatch_backend_embd.resize(n_embd * n_tokens); + ubatch_out_embd.resize(n_embd); ubatch_pos.resize(n_ubatch); ubatch_n_seq_id.resize(n_ubatch); ubatch_seq_id.resize(n_ubatch); @@ -2986,6 +3014,8 @@ struct llama_sbatch { /*n_seqs =*/ 0, /*token =*/ !has_embd ? ubatch_token.data() : nullptr, /*embd =*/ has_embd ? ubatch_embd.data() : nullptr, + /*backend_embd =*/ ubatch_backend_embd.data(), + /*out_embd =*/ ubatch_out_embd.data(), /*pos =*/ ubatch_pos.data(), /*n_seq_id =*/ ubatch_n_seq_id.data(), /*seq_id =*/ ubatch_seq_id.data(), @@ -3169,8 +3199,8 @@ struct llama_sbatch { void from_batch(const llama_batch & batch, const size_t n_embd, const bool simple_split = false, const bool logits_all = false) { GGML_ASSERT(batch.n_tokens >= 0); - this->batch = &batch; - this->n_embd = n_embd; + this->batch = &batch; + this->n_embd = n_embd; this->logits_all = logits_all; n_tokens = batch.n_tokens; @@ -3184,10 +3214,10 @@ struct llama_sbatch { if (simple_split) { seq.resize(1); llama_sbatch_seq & s = seq[0]; - s.n_seq_id = 0; - s.seq_id = nullptr; - s.offset = 0; - s.length = n_tokens; + s.n_seq_id = 0; + s.seq_id = nullptr; + s.offset = 0; + s.length = n_tokens; s.all_seq_id = batch.all_seq_id; return; } @@ -3265,7 +3295,9 @@ struct llama_context { , t_load_us(model.t_load_us) {} ~llama_context() { - ggml_backend_sched_free(sched); + for (ggml_backend_sched_t s : sched) { + ggml_backend_sched_free(s); + } for (ggml_backend_t backend : backends) { ggml_backend_free(backend); @@ -3335,27 +3367,40 @@ struct llama_context { // memory buffers used to evaluate the model std::vector buf_compute_meta; - ggml_backend_sched_t sched = nullptr; + std::vector sched = {}; ggml_abort_callback abort_callback = nullptr; void * abort_callback_data = nullptr; // input tensors - struct ggml_tensor * inp_tokens; // I32 [n_batch] - struct ggml_tensor * inp_embd; // F32 [n_embd, n_batch] - struct ggml_tensor * inp_pos; // I32 [n_batch] - struct ggml_tensor * inp_out_ids; // I32 [n_outputs] - struct ggml_tensor * inp_KQ_mask; // F32 [kv_size, n_batch] - struct ggml_tensor * inp_KQ_mask_swa; // F32 [kv_size, n_batch] - struct ggml_tensor * inp_K_shift; // I32 [kv_size] - struct ggml_tensor * inp_mean; // F32 [n_batch, n_batch] - struct ggml_tensor * inp_cls; // I32 [n_batch] - struct ggml_tensor * inp_s_copy; // I32 [kv_size] - struct ggml_tensor * inp_s_mask; // F32 [1, n_kv] - struct ggml_tensor * inp_s_seq; // I32 [n_kv, n_batch] + struct ggml_tensor * inp_tokens; // I32 [n_batch] + struct ggml_tensor * inp_embd; // F32 [n_embd, n_batch] + struct ggml_tensor * backend_embd; // F32 [n_embd, n_tokens] + struct ggml_tensor * out_embd; // F32 [n_embd, n_outputs] + struct ggml_tensor * inp_pos; // I32 [n_batch] + struct ggml_tensor * inp_out_ids; // I32 [n_outputs] + struct ggml_tensor * inp_KQ_mask; // F32 [kv_size, n_batch] + struct ggml_tensor * inp_KQ_mask_swa; // F32 [kv_size, n_batch] + struct ggml_tensor * inp_K_shift; // I32 [kv_size] + struct ggml_tensor * inp_mean; // F32 [n_batch, n_batch] + struct ggml_tensor * inp_cls; // I32 [n_batch] + struct ggml_tensor * inp_s_copy; // I32 [kv_size] + struct ggml_tensor * inp_s_mask; // F32 [1, n_kv] + struct ggml_tensor * inp_s_seq; // I32 [n_kv, n_batch] struct ggml_tensor * inp_pos_bucket; // I32 [n_batch|n_kv, n_batch] struct ggml_tensor * inp_embd_enc; // F32 [n_embd, n_outputs_enc] struct ggml_tensor * inp_KQ_mask_cross; // F32 [n_outputs_enc, n_batch] + + // sockets + std::string master_ip = "localhost"; + std::string next_node_ip = "localhost"; + uint32_t data_port = 9000; + uint32_t signal_port = 10000; + zmq::context_t * sock_context = nullptr; + zmq::socket_t * send_socket = nullptr; + zmq::socket_t * recv_socket = nullptr; + zmq::socket_t * master_socket = nullptr; + zmq::socket_t * signal_socket = nullptr; }; struct llama_lora_weight { @@ -3565,6 +3610,31 @@ static size_t llama_get_device_memory(const llama_model & model, int device) { GGML_UNUSED(device); } +static bool this_layer_is_mine( + uint32_t layer_id, + uint32_t n_world, + uint32_t my_rank, + uint32_t n_layer_window) { + return (layer_id / n_layer_window) % n_world == my_rank; +} + +static int64_t map_layer_to_local_id( + int64_t layer_id, + uint32_t n_world, + uint32_t my_rank, + uint32_t n_layer_window) { + if (!this_layer_is_mine(layer_id, n_world, my_rank, n_layer_window)) { + return -1; + } + // map layer_id to kvcache_id. + // example: For n_world=2 and n_layer_window=4, rank 0 handles layers 0-3, 8-11, 16-19, while rank 1 handles layers 4-7, 12-15, 20-23. + // on rank 0, layer_id should map to kvcache_id as follows: 0-3 -> 0-3, 8-11 -> 4-7, 16-19 -> 8-11. + // on rank 1, layer_id should map to kvcache_id as follows: 4-7 -> 0-3, 12-15 -> 4-7, 20-23 -> 8-11. + int64_t cycle_size = n_world * n_layer_window; + int64_t local_offset = (layer_id / cycle_size) * n_layer_window; + return (layer_id % cycle_size) % n_layer_window + local_offset; +} + // // kv cache helpers // @@ -3576,36 +3646,41 @@ static bool llama_kv_cache_init( ggml_type type_v, uint32_t kv_size, bool offload) { - const llama_model & model = ctx->model; - const llama_cparams & cparams = ctx->cparams; - + const llama_model & model = ctx->model; + const llama_cparams & cparams = ctx->cparams; const struct llama_hparams & hparams = model.hparams; - - const int64_t n_layer = hparams.n_layer; + const int64_t n_layer = hparams.n_layer; + const uint32_t n_world = cparams.n_world; + const uint32_t my_rank = cparams.rank; + const uint32_t n_layer_window = cparams.n_layer_window; cache.has_shift = false; - cache.recurrent = llama_model_is_recurrent(&model); cache.v_trans = !cache.recurrent && !cparams.flash_attn; - cache.head = 0; cache.size = kv_size; cache.used = 0; - cache.type_k = type_k; cache.type_v = type_v; - cache.cells.clear(); cache.cells.resize(kv_size); // count used buffer types std::map buft_layer_count; - if (offload) { - for (int64_t i = 0; i < n_layer; ++i) { - buft_layer_count[model.buft_layer[i].buft]++; + int64_t local_i; + uint32_t my_layers = 0; + + for (int64_t i = 0; i < n_layer; ++i) { + if (!this_layer_is_mine(i, n_world, my_rank, n_layer_window)) { + continue; + } + my_layers++; + local_i = map_layer_to_local_id(i, n_world, my_rank, n_layer_window); + if (offload) { + buft_layer_count[model.buft_layer[local_i].buft]++; + } else { + buft_layer_count[llama_default_buffer_type_cpu(model, true)]++; } - } else { - buft_layer_count[llama_default_buffer_type_cpu(model, true)] = n_layer; } // create a context for each buffer type @@ -3626,16 +3701,21 @@ static bool llama_kv_cache_init( cache.ctxs.push_back(ctx); } - cache.k_l.reserve(n_layer); - cache.v_l.reserve(n_layer); + cache.k_l.reserve(my_layers); + cache.v_l.reserve(my_layers); for (int i = 0; i < (int) n_layer; i++) { + if (!this_layer_is_mine(i, n_world, my_rank, n_layer_window)) { + continue; + } + int local_i = map_layer_to_local_id(i, n_world, my_rank, n_layer_window); + const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(i) + hparams.n_embd_k_s(); const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(i) + hparams.n_embd_v_s(); - struct ggml_context * ctx = offload ? ctx_map.at(model.buft_layer[i].buft) : cache.ctxs.front(); - ggml_tensor * k = ggml_new_tensor_1d(ctx, type_k, n_embd_k_gqa*kv_size); - ggml_tensor * v = ggml_new_tensor_1d(ctx, type_v, n_embd_v_gqa*kv_size); + struct ggml_context * ctx = offload ? ctx_map.at(model.buft_layer[local_i].buft) : cache.ctxs.front(); + ggml_tensor * k = ggml_new_tensor_1d(ctx, type_k, n_embd_k_gqa*kv_size); + ggml_tensor * v = ggml_new_tensor_1d(ctx, type_v, n_embd_v_gqa*kv_size); ggml_format_name(k, "cache_k_l%d", i); ggml_format_name(v, "cache_v_l%d", i); cache.k_l.push_back(k); @@ -6939,10 +7019,12 @@ static void llm_load_print_meta(llama_model_loader & ml, llama_model & model) { static bool llm_load_tensors( llama_model_loader & ml, llama_model & model, + uint32_t n_world, + uint32_t my_rank, + uint32_t n_layer_window, int n_gpu_layers, enum llama_split_mode split_mode, int main_gpu, - const float * tensor_split, bool use_mlock, llama_progress_callback progress_callback, void * progress_callback_user_data) { @@ -6955,95 +7037,55 @@ static bool llm_load_tensors( throw std::runtime_error(format("invalid value for main_gpu: %d (available devices: %d)", main_gpu, llama_get_device_count(model))); } - model.split_mode = split_mode; - model.main_gpu = main_gpu; - model.n_gpu_layers = n_gpu_layers; + model.split_mode = split_mode; + model.main_gpu = main_gpu; + model.n_gpu_layers = n_gpu_layers; - const int n_layer = hparams.n_layer; - const int i_gpu_start = std::max((int) hparams.n_layer - n_gpu_layers, (int) 0); + const int n_layer = hparams.n_layer; bool use_mmap_buffer = true; // there is very little benefit to offloading the input layer, so always keep it on the CPU model.buft_input = llama_default_buffer_type_cpu(model, true); - //model.buft_input = llama_default_buffer_type_offload(main_gpu); - model.buft_layer.resize(n_layer); + int my_layers = 0; + for (int i = 0; i < n_layer; ++i) { + if (this_layer_is_mine(i, n_world, my_rank, n_layer_window)) { + my_layers++; + } + } + model.buft_layer.resize(my_layers); - // assign cpu layers - for (int i = 0; i < i_gpu_start; ++i) { - model.buft_layer[i] = llama_default_buffer_type_cpu(model, true); + for (int i = 0; i < n_layer; ++i) { + if (this_layer_is_mine(i, n_world, my_rank, n_layer_window)) { + int local_i = map_layer_to_local_id(i, n_world, my_rank, n_layer_window); + if (i % (int)n_layer_window >= (int)n_layer_window - n_gpu_layers) { + LLAMA_LOG_INFO("Layer %i assigned to gpu (cache index %i)\n", i, local_i); + model.buft_layer[local_i] = llama_default_buffer_type_offload(model, main_gpu); + } else { + LLAMA_LOG_INFO("Layer %i assigned to cpu (cache index %i)\n", i, local_i); + model.buft_layer[local_i] = llama_default_buffer_type_cpu(model, true); + } + } } - if (split_mode == LLAMA_SPLIT_MODE_LAYER) { - // calculate the split points - int device_count = llama_get_device_count(model); - bool all_zero = tensor_split == nullptr || std::all_of(tensor_split, tensor_split + device_count, [](float x) { return x == 0.0f; }); - std::vector splits(device_count); - if (all_zero) { - // default split, by free memory - for (int i = 0; i < device_count; ++i) { - splits[i] = llama_get_device_memory(model, i); - } - } else { - std::copy(tensor_split, tensor_split + device_count, splits.begin()); - } - - // sum and normalize the splits to get the split points - float split_sum = 0.0f; - for (int i = 0; i < device_count; ++i) { - split_sum += splits[i]; - splits[i] = split_sum; - } - for (int i = 0; i < device_count; ++i) { - splits[i] /= split_sum; - } - - // assign the repeating layers to the devices according to the splits - int act_gpu_layers = std::min(n_gpu_layers, (int)n_layer + 1); - for (int i = i_gpu_start; i < n_layer; ++i) { - int layer_gpu = std::upper_bound(splits.begin(), splits.begin() + device_count, float(i - i_gpu_start)/act_gpu_layers) - splits.begin(); - model.buft_layer[i] = llama_default_buffer_type_offload(model, layer_gpu); - } - // assign the output layer - if (n_gpu_layers > n_layer) { - int layer_gpu = std::upper_bound(splits.begin(), splits.begin() + device_count, float(act_gpu_layers - 1)/act_gpu_layers) - splits.begin(); - model.buft_output = llama_default_buffer_type_offload(model, layer_gpu); - } else { - model.buft_output = llama_default_buffer_type_cpu(model, true); - } + // assign the output layer + if (my_rank == 0 && n_gpu_layers > (int)n_layer_window) { + LLAMA_LOG_INFO("Layer output assigned to gpu\n"); + model.buft_output = llama_default_buffer_type_offload(model, main_gpu); } else { - ggml_backend_buffer_type_t split_buft; - if (split_mode == LLAMA_SPLIT_MODE_ROW) { - split_buft = llama_default_buffer_type_split(model, main_gpu, tensor_split); - } else { - // LLAMA_SPLIT_MODE_NONE or LLAMA_SPLIT_MODE_LAYER in backends where it is not supported - split_buft = llama_default_buffer_type_offload(model, main_gpu); - } - // assign the repeating layers - for (int i = i_gpu_start; i < n_layer; ++i) { - model.buft_layer[i] = { - split_buft, - llama_default_buffer_type_offload(model, main_gpu) - }; - } - // assign the output layer - if (n_gpu_layers > n_layer) { - model.buft_output = { - split_buft, - llama_default_buffer_type_offload(model, main_gpu) - }; - } else { - model.buft_output = llama_default_buffer_type_cpu(model, true); - } + LLAMA_LOG_INFO("Layer output assigned to cpu\n"); + model.buft_output = llama_default_buffer_type_cpu(model, true); } // count used buffer types std::map buft_layer_count; - buft_layer_count[model.buft_input.buft]++; - buft_layer_count[model.buft_input.buft_matrix]++; - buft_layer_count[model.buft_output.buft]++; - buft_layer_count[model.buft_output.buft_matrix]++; - for (int i = 0; i < n_layer; ++i) { + if (my_rank == 0) { + buft_layer_count[model.buft_input.buft]++; + buft_layer_count[model.buft_input.buft_matrix]++; + buft_layer_count[model.buft_output.buft]++; + buft_layer_count[model.buft_output.buft_matrix]++; + } + for (int i = 0; i < my_layers; ++i) { buft_layer_count[model.buft_layer[i].buft]++; buft_layer_count[model.buft_layer[i].buft_matrix]++; } @@ -7069,8 +7111,6 @@ static bool llm_load_tensors( model.ctxs.push_back(ctx); } - LLAMA_LOG_INFO("%s: ggml ctx size = %7.2f MiB\n", __func__, model.ctxs.size()*ctx_size/1024.0/1024.0); - // create tensors for the weights { // note: cast to int64_t since we will use these for the tensor dimensions @@ -7094,14 +7134,20 @@ static bool llm_load_tensors( throw std::runtime_error("model has expert layers but no expert layers are used"); } - ggml_context * ctx_input = ctx_map.at(model.buft_input.buft); - ggml_context * ctx_output = ctx_map.at(model.buft_output.buft); - ggml_context * ctx_output_split = ctx_map.at(model.buft_output.buft_matrix); + ggml_context * ctx_input = nullptr; + ggml_context * ctx_output = nullptr; + ggml_context * ctx_output_split = nullptr; - auto ctx_for_layer = [&](int i) { return ctx_map.at(model.buft_layer[i].buft); }; - auto ctx_for_layer_split = [&](int i) { return ctx_map.at(model.buft_layer[i].buft_matrix); }; + if (my_rank == 0) { + ctx_input = ctx_map.at(model.buft_input.buft); + ctx_output = ctx_map.at(model.buft_output.buft); + ctx_output_split = ctx_map.at(model.buft_output.buft_matrix); + } - model.layers.resize(n_layer); + auto ctx_for_layer = [&](int local_i) { return ctx_map.at(model.buft_layer[local_i].buft); }; + auto ctx_for_layer_split = [&](int local_i) { return ctx_map.at(model.buft_layer[local_i].buft_matrix); }; + + model.layers.resize(my_layers); const auto tn = LLM_TN(model.arch); switch (model.arch) { @@ -7111,24 +7157,31 @@ static bool llm_load_tensors( case LLM_ARCH_GRANITE: case LLM_ARCH_GRANITE_MOE: { - model.tok_embd = ml.create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}); + if (my_rank == 0) { + model.tok_embd = ml.create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}); - // output - { - model.output_norm = ml.create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}); - model.output = ml.create_tensor(ctx_output_split, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_NOT_REQUIRED); + // output + { + model.output_norm = ml.create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}); + model.output = ml.create_tensor(ctx_output_split, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_NOT_REQUIRED); - // if output is NULL, init from the input tok embed - if (model.output == NULL) { - model.output = ml.create_tensor(ctx_output, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_DUPLICATED); + // if output is NULL, init from the input tok embed + if (model.output == NULL) { + model.output = ml.create_tensor(ctx_output, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_DUPLICATED); + } } } for (int i = 0; i < n_layer; ++i) { - ggml_context * ctx_layer = ctx_for_layer(i); - ggml_context * ctx_split = ctx_for_layer_split(i); + if (!this_layer_is_mine(i, n_world, my_rank, n_layer_window)) { + continue; + } - auto & layer = model.layers[i]; + int local_i = map_layer_to_local_id(i, n_world, my_rank, n_layer_window); + ggml_context * ctx_layer = ctx_for_layer(local_i); + ggml_context * ctx_split = ctx_for_layer_split(local_i); + + auto & layer = model.layers[local_i]; layer.attn_norm = ml.create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}); @@ -8882,8 +8935,6 @@ static bool llm_load_tensors( } } - ml.done_getting_tensors(); - ml.init_mappings(true, use_mlock ? &model.mlock_mmaps : nullptr); model.mappings.reserve(ml.mappings.size()); @@ -8904,7 +8955,12 @@ static bool llm_load_tensors( // check if this backend device supports buffer_from_host_ptr // when using a host buffer as the CPU bakcend buffer, use the CPU device to prioritize using buffer_from_host_ptr over the host buffer - ggml_backend_dev_t dev = ggml_backend_buft_get_device(buft == llama_default_buffer_type_cpu(model, true) ? ggml_backend_cpu_buffer_type() : buft); + ggml_backend_buffer_type_t buft_type = buft == llama_default_buffer_type_cpu(model, true) ? ggml_backend_cpu_buffer_type() : buft; + if (buft_type == nullptr) { + continue; + } + ggml_backend_dev_t dev = ggml_backend_buft_get_device(buft_type); + bool buffer_from_host_ptr_supported = false; if (dev) { ggml_backend_dev_props props; @@ -8931,8 +8987,7 @@ static bool llm_load_tensors( model.bufs.push_back(buf); bufs.emplace(idx, buf); } - } - else { + } else { ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft); if (buf == nullptr) { throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft))); @@ -8962,25 +9017,6 @@ static bool llm_load_tensors( ctx_bufs.emplace_back(ctx, bufs); } - if (llama_supports_gpu_offload()) { - const int n_gpu = std::min(n_gpu_layers, int(hparams.n_layer)); - - LLAMA_LOG_INFO("%s: offloading %d repeating layers to GPU\n", __func__, n_gpu); - if (n_gpu_layers > (int) hparams.n_layer) { - LLAMA_LOG_INFO("%s: offloading non-repeating layers to GPU\n", __func__); - } - - const int max_backend_supported_layers = hparams.n_layer + 1; - const int max_offloadable_layers = hparams.n_layer + 1; - - LLAMA_LOG_INFO("%s: offloaded %d/%d layers to GPU\n", __func__, std::min(n_gpu_layers, max_offloadable_layers), max_backend_supported_layers); - } - - // print memory requirements - for (ggml_backend_buffer_t buf : model.bufs) { - LLAMA_LOG_INFO("%s: %10s buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf) / 1024.0 / 1024.0); - } - // populate tensors_by_name for (ggml_context * ctx : model.ctxs) { for (auto * cur = ggml_get_first_tensor(ctx); cur != NULL; cur = ggml_get_next_tensor(ctx, cur)) { @@ -9061,8 +9097,8 @@ static int llama_model_load(const std::string & fname, llama_model & model, llam #endif if (!llm_load_tensors( - ml, model, params.n_gpu_layers, params.split_mode, params.main_gpu, params.tensor_split, params.use_mlock, - params.progress_callback, params.progress_callback_user_data + ml, model, params.n_world, params.rank, params.n_layer_window, params.n_gpu_layers, params.split_mode, + params.main_gpu, params.use_mlock, params.progress_callback, params.progress_callback_user_data )) { return -2; } @@ -9146,15 +9182,18 @@ static void llm_build_kv_store( int32_t n_tokens, int32_t kv_head, const llm_build_cb & cb, - int64_t il) { - const int64_t n_ctx = cparams.n_ctx; - - const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il); - const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa(il); + int il) { + const int64_t n_ctx = cparams.n_ctx; + const uint32_t n_world = cparams.n_world; + const uint32_t my_rank = cparams.rank; + const uint32_t n_layer_window = cparams.n_layer_window; + const int local_il = map_layer_to_local_id(il, n_world, my_rank, n_layer_window); + const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il); + const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa(il); GGML_ASSERT(kv.size == n_ctx); - struct ggml_tensor * k_cache_view = ggml_view_1d(ctx, kv.k_l[il], n_tokens*n_embd_k_gqa, ggml_row_size(kv.k_l[il]->type, n_embd_k_gqa)*kv_head); + struct ggml_tensor * k_cache_view = ggml_view_1d(ctx, kv.k_l[local_il], n_tokens*n_embd_k_gqa, ggml_row_size(kv.k_l[local_il]->type, n_embd_k_gqa)*kv_head); cb(k_cache_view, "k_cache_view", il); // note: storing RoPE-ed version of K in the KV cache @@ -9165,12 +9204,12 @@ static void llm_build_kv_store( struct ggml_tensor * v_cache_view = nullptr; if (cparams.flash_attn) { - v_cache_view = ggml_view_1d(ctx, kv.v_l[il], n_tokens*n_embd_v_gqa, ggml_row_size(kv.v_l[il]->type, n_embd_v_gqa)*kv_head); + v_cache_view = ggml_view_1d(ctx, kv.v_l[local_il], n_tokens*n_embd_v_gqa, ggml_row_size(kv.v_l[local_il]->type, n_embd_v_gqa)*kv_head); } else { // note: the V cache is transposed when not using flash attention - v_cache_view = ggml_view_2d(ctx, kv.v_l[il], n_tokens, n_embd_v_gqa, - ( n_ctx)*ggml_element_size(kv.v_l[il]), - (kv_head)*ggml_element_size(kv.v_l[il])); + v_cache_view = ggml_view_2d(ctx, kv.v_l[local_il], n_tokens, n_embd_v_gqa, + ( n_ctx)*ggml_element_size(kv.v_l[local_il]), + (kv_head)*ggml_element_size(kv.v_l[local_il])); v_cur = ggml_transpose(ctx, v_cur); } @@ -9505,10 +9544,14 @@ static struct ggml_tensor * llm_build_kqv( float kq_scale, const llm_build_cb & cb, int il) { - const llama_model & model = lctx.model; - const llama_hparams & hparams = lctx.model.hparams; - const llama_cparams & cparams = lctx.cparams; - + const llama_model & model = lctx.model; + const llama_hparams & hparams = lctx.model.hparams; + const llama_cparams & cparams = lctx.cparams; + const uint32_t n_world = cparams.n_world; + const uint32_t my_rank = cparams.rank; + const uint32_t n_layer_window = cparams.n_layer_window; + + const int local_il = map_layer_to_local_id(il, n_world, my_rank, n_layer_window); const int64_t n_ctx = cparams.n_ctx; const int64_t n_head = hparams.n_head(il); const int64_t n_head_kv = hparams.n_head_kv(il); @@ -9521,10 +9564,10 @@ static struct ggml_tensor * llm_build_kqv( cb(q, "q", il); struct ggml_tensor * k = - ggml_view_3d(ctx, kv.k_l[il], + ggml_view_3d(ctx, kv.k_l[local_il], n_embd_head_k, n_kv, n_head_kv, - ggml_row_size(kv.k_l[il]->type, n_embd_k_gqa), - ggml_row_size(kv.k_l[il]->type, n_embd_head_k), + ggml_row_size(kv.k_l[local_il]->type, n_embd_k_gqa), + ggml_row_size(kv.k_l[local_il]->type, n_embd_head_k), 0); cb(k, "k", il); @@ -9536,10 +9579,10 @@ static struct ggml_tensor * llm_build_kqv( // split cached v into n_head heads (not transposed) struct ggml_tensor * v = - ggml_view_3d(ctx, kv.v_l[il], + ggml_view_3d(ctx, kv.v_l[local_il], n_embd_head_v, n_kv, n_head_kv, - ggml_row_size(kv.v_l[il]->type, n_embd_v_gqa), - ggml_row_size(kv.v_l[il]->type, n_embd_head_v), + ggml_row_size(kv.v_l[local_il]->type, n_embd_v_gqa), + ggml_row_size(kv.v_l[local_il]->type, n_embd_head_v), 0); cb(v, "v", il); @@ -9588,10 +9631,10 @@ static struct ggml_tensor * llm_build_kqv( // split cached v into n_head heads struct ggml_tensor * v = - ggml_view_3d(ctx, kv.v_l[il], + ggml_view_3d(ctx, kv.v_l[local_il], n_kv, n_embd_head_v, n_head_kv, - ggml_element_size(kv.v_l[il])*n_ctx, - ggml_element_size(kv.v_l[il])*n_ctx*n_embd_head_v, + ggml_element_size(kv.v_l[local_il])*n_ctx, + ggml_element_size(kv.v_l[local_il])*n_ctx*n_embd_head_v, 0); cb(v, "v", il); @@ -10096,6 +10139,8 @@ struct llm_build_context { lctx.inp_tokens = nullptr; lctx.inp_embd = nullptr; + lctx.backend_embd = nullptr; + lctx.out_embd = nullptr; lctx.inp_pos = nullptr; lctx.inp_out_ids = nullptr; lctx.inp_KQ_mask = nullptr; @@ -10146,7 +10191,7 @@ struct llm_build_context { for (auto * backend : lctx.backends) { // Figure out which backend KV cache belongs to if (ggml_backend_supports_buft(backend, lctx.model.buft_layer[il].buft)) { - ggml_backend_sched_set_tensor_backend(lctx.sched, tmp, backend); + ggml_backend_sched_set_tensor_backend(lctx.sched.at(0), tmp, backend); // todo. break; } } @@ -10230,9 +10275,6 @@ struct llm_build_context { i += nm - 1; } - - //LLAMA_LOG_INFO("gf->n_nodes = %d\n", gf->n_nodes); - return gf; } @@ -10243,6 +10285,29 @@ struct llm_build_context { return lctx.inp_pos; } + struct ggml_tensor * llm_build_backend_embd( + struct ggml_context * ctx, + struct llama_context & lctx, + const llama_hparams & hparams, + const llama_ubatch & batch, + const llm_build_cb & cb) { + lctx.backend_embd = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hparams.n_embd, batch.n_tokens); + cb(lctx.backend_embd, "backend_embd", -1); + ggml_set_input(lctx.backend_embd); + return lctx.backend_embd; + } + + struct ggml_tensor * llm_build_out_embd( + struct ggml_context * ctx, + struct llama_context & lctx, + const llama_hparams & hparams, + const llm_build_cb & cb) { + lctx.out_embd = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hparams.n_embd, n_outputs); + cb(lctx.out_embd, "out_embd", -1); + ggml_set_input(lctx.out_embd); + return lctx.out_embd; + } + struct ggml_tensor * build_rope_factors(int il) { // choose long/short freq factors based on the context size const auto n_ctx_pre_seq = cparams.n_ctx / cparams.n_seq_max; @@ -10427,20 +10492,32 @@ struct llm_build_context { return lctx.inp_KQ_mask_cross; } - struct ggml_cgraph * build_llama() { - struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model), false); + std::vector build_llama() { + // create a vector to hold sub-graphs + std::vector sub_gfs; + struct ggml_cgraph * sub_gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model), false); // mutable variable, needed during the last layer of the computation to skip unused tokens int32_t n_tokens = this->n_tokens; const int64_t n_embd_head = hparams.n_embd_head_v; - GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); - GGML_ASSERT(n_embd_head == hparams.n_rot); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k); + GGML_ASSERT(n_embd_head == hparams.n_rot); - struct ggml_tensor * cur; - struct ggml_tensor * inpL; + struct ggml_tensor * cur = nullptr; + struct ggml_tensor * inpL = nullptr; + struct ggml_tensor * inpB = nullptr; - inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); + const uint32_t n_world = this->cparams.n_world; + const uint32_t my_rank = this->cparams.rank; + const uint32_t n_layer_window = this->cparams.n_layer_window; + + if (my_rank == 0) { + // inp_embd - contains the input embedding + inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); + // inpB - contains the output embedding from other nodes + inpB = llm_build_backend_embd(ctx0, lctx, hparams, batch, cb); + } // inp_pos - contains the positions struct ggml_tensor * inp_pos = build_inp_pos(); @@ -10450,38 +10527,63 @@ struct llm_build_context { const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; for (int il = 0; il < n_layer; ++il) { - struct ggml_tensor * inpSA = inpL; + if (!this_layer_is_mine(il, n_world, my_rank, n_layer_window)) { + // if we have an active sub-graph, add it to the list + if (sub_gf != nullptr && inpL != nullptr) { + ggml_build_forward_expand(sub_gf, cur); + sub_gfs.push_back(sub_gf); + sub_gf = nullptr; + } + if (inpL != inpB) { + inpL = inpB; + } + continue; + } + + int local_il = map_layer_to_local_id(il, n_world, my_rank, n_layer_window); + + if (inpL == nullptr) { + inpB = llm_build_backend_embd(ctx0, lctx, hparams, batch, cb); + inpL = inpB; + } + + // start a new sub-graph + if (sub_gf == nullptr) { + sub_gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model), false); + } + + struct ggml_tensor * inpSA = inpL; // use for shortcut // norm cur = llm_build_norm(ctx0, inpL, hparams, - model.layers[il].attn_norm, NULL, + model.layers[local_il].attn_norm, NULL, LLM_NORM_RMS, cb, il); cb(cur, "attn_norm", il); // self-attention { // rope freq factors for llama3; may return nullptr for llama2 and other models - struct ggml_tensor * rope_factors = build_rope_factors(il); + struct ggml_tensor * rope_factors = build_rope_factors(local_il); // compute Q and K and RoPE them - struct ggml_tensor * Qcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wq, cur); + struct ggml_tensor * Qcur = llm_build_lora_mm(lctx, ctx0, model.layers[local_il].wq, cur); cb(Qcur, "Qcur", il); - if (model.layers[il].bq) { - Qcur = ggml_add(ctx0, Qcur, model.layers[il].bq); + if (model.layers[local_il].bq) { + Qcur = ggml_add(ctx0, Qcur, model.layers[local_il].bq); cb(Qcur, "Qcur", il); } - struct ggml_tensor * Kcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, cur); + struct ggml_tensor * Kcur = llm_build_lora_mm(lctx, ctx0, model.layers[local_il].wk, cur); cb(Kcur, "Kcur", il); - if (model.layers[il].bk) { - Kcur = ggml_add(ctx0, Kcur, model.layers[il].bk); + if (model.layers[local_il].bk) { + Kcur = ggml_add(ctx0, Kcur, model.layers[local_il].bk); cb(Kcur, "Kcur", il); } - struct ggml_tensor * Vcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, cur); + struct ggml_tensor * Vcur = llm_build_lora_mm(lctx, ctx0, model.layers[local_il].wv, cur); cb(Vcur, "Vcur", il); - if (model.layers[il].bv) { - Vcur = ggml_add(ctx0, Vcur, model.layers[il].bv); + if (model.layers[local_il].bv) { + Vcur = ggml_add(ctx0, Vcur, model.layers[local_il].bv); cb(Vcur, "Vcur", il); } @@ -10499,8 +10601,8 @@ struct llm_build_context { ); cb(Kcur, "Kcur", il); - cur = llm_build_kv(ctx0, lctx, kv_self, gf, - model.layers[il].wo, model.layers[il].bo, + cur = llm_build_kv(ctx0, lctx, kv_self, sub_gf, + model.layers[local_il].wo, model.layers[local_il].bo, Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, kq_scale, cb, il); } @@ -10517,35 +10619,35 @@ struct llm_build_context { cur = ggml_scale(ctx0, cur, hparams.f_residual_scale); } - struct ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + struct ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); // shortcut cb(ffn_inp, "ffn_inp", il); // feed-forward network - if (model.layers[il].ffn_gate_inp == nullptr) { + if (model.layers[local_il].ffn_gate_inp == nullptr) { cur = llm_build_norm(ctx0, ffn_inp, hparams, - model.layers[il].ffn_norm, NULL, + model.layers[local_il].ffn_norm, NULL, LLM_NORM_RMS, cb, il); cb(cur, "ffn_norm", il); cur = llm_build_ffn(ctx0, lctx, cur, - model.layers[il].ffn_up, model.layers[il].ffn_up_b, NULL, - model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, NULL, - model.layers[il].ffn_down, model.layers[il].ffn_down_b, NULL, + model.layers[local_il].ffn_up, model.layers[local_il].ffn_up_b, NULL, + model.layers[local_il].ffn_gate, model.layers[local_il].ffn_gate_b, NULL, + model.layers[local_il].ffn_down, model.layers[local_il].ffn_down_b, NULL, NULL, LLM_FFN_SILU, LLM_FFN_PAR, cb, il); cb(cur, "ffn_out", il); } else { // MoE branch cur = llm_build_norm(ctx0, ffn_inp, hparams, - model.layers[il].ffn_norm, NULL, + model.layers[local_il].ffn_norm, NULL, LLM_NORM_RMS, cb, il); cb(cur, "ffn_norm", il); cur = llm_build_moe_ffn(ctx0, lctx, cur, - model.layers[il].ffn_gate_inp, - model.layers[il].ffn_up_exps, - model.layers[il].ffn_gate_exps, - model.layers[il].ffn_down_exps, + model.layers[local_il].ffn_gate_inp, + model.layers[local_il].ffn_up_exps, + model.layers[local_il].ffn_gate_exps, + model.layers[local_il].ffn_down_exps, n_expert, n_expert_used, LLM_FFN_SILU, true, false, 0.0, @@ -10558,7 +10660,7 @@ struct llm_build_context { cur = ggml_scale(ctx0, cur, hparams.f_residual_scale); } - cur = ggml_add(ctx0, cur, ffn_inp); + cur = ggml_add(ctx0, cur, ffn_inp); // shortcut cb(cur, "ffn_out", il); cur = lctx.cvec.apply_to(ctx0, cur, il); @@ -10568,26 +10670,43 @@ struct llm_build_context { inpL = cur; } - cur = inpL; - - cur = llm_build_norm(ctx0, cur, hparams, - model.output_norm, NULL, - LLM_NORM_RMS, cb, -1); - cb(cur, "result_norm", -1); - - // lm_head - cur = llm_build_lora_mm(lctx, ctx0, model.output, cur); - - // For Granite architecture - if (hparams.f_logit_scale) { - cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale); + // add the last active sub-graph to the list + if (sub_gf != nullptr){ + ggml_build_forward_expand(sub_gf, cur); + sub_gfs.push_back(sub_gf); + sub_gf = nullptr; } - cb(cur, "result_output", -1); + // output norm and lm_head + if (my_rank == 0) { + // start a new sub-graph for the output + sub_gf = ggml_new_graph_custom(ctx0, llama_model_max_nodes(model), false); - ggml_build_forward_expand(gf, cur); + cur = llm_build_out_embd(ctx0, lctx, hparams, cb); - return gf; + // cur = ggml_get_rows(ctx0, cur, inp_out_ids); + + cur = llm_build_norm(ctx0, cur, hparams, + model.output_norm, NULL, + LLM_NORM_RMS, cb, -1); + cb(cur, "result_norm", -1); + + // lm_head + cur = llm_build_lora_mm(lctx, ctx0, model.output, cur); + + // For Granite architecture + if (hparams.f_logit_scale) { + cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale); + } + + cb(cur, "result_output", -1); + + ggml_build_forward_expand(sub_gf, cur); + sub_gfs.push_back(sub_gf); + sub_gf = nullptr; + } + + return sub_gfs; } struct ggml_cgraph * build_baichuan() { @@ -16211,7 +16330,7 @@ static struct ggml_cgraph * llama_build_graph_defrag(llama_context & lctx, const llama_ubatch dummy = {}; dummy.equal_seqs = true; - llm_build_cb cb = [&](struct ggml_tensor * , const char * , int ) { }; + llm_build_cb cb = [&](struct ggml_tensor * , const char * , int) { }; struct llm_build_context llm(lctx, dummy, cb, false); @@ -16241,16 +16360,34 @@ static struct ggml_cgraph * llama_build_graph_k_shift(llama_context & lctx) { return result; } -static struct ggml_cgraph * llama_build_graph( +static uint32_t map_layer_to_subgf_id(uint32_t i, uint32_t my_rank, uint32_t n_world, uint32_t n_layer, uint32_t n_layer_window) { + uint32_t global_layer_offset = my_rank * n_layer_window; + uint32_t step = n_world * n_layer_window; + if (i < n_layer) { + uint32_t relative_layer = i % step; + if (relative_layer >= global_layer_offset && relative_layer < global_layer_offset + n_layer_window) { + return i / step; + } + } + return -1; +} + +static std::vector llama_build_graph( llama_context & lctx, const llama_ubatch & batch, bool worst_case) { - const auto & model = lctx.model; + const auto & model = lctx.model; + const uint32_t n_world = lctx.cparams.n_world; + const uint32_t my_rank = lctx.cparams.rank; + const uint32_t n_layer_window = lctx.cparams.n_layer_window; + const uint32_t n_layer = lctx.model.hparams.n_layer; // this callback allows us to apply custom logic to each tensor (e.g. ggml-alloc, offloading, etc.) llm_build_cb cb = [&](struct ggml_tensor * cur, const char * name, int il) { + int sub_gf_id = 0; if (il >= 0) { ggml_format_name(cur, "%s-%d", name, il); + sub_gf_id = map_layer_to_subgf_id(il, my_rank, n_world, n_layer, n_layer_window); } else { ggml_set_name(cur, name); } @@ -16258,19 +16395,20 @@ static struct ggml_cgraph * llama_build_graph( if (!lctx.cparams.offload_kqv) { if (strcmp(name, "kqv_merged_cont") == 0) { // all nodes between the KV store and the attention output are run on the CPU - ggml_backend_sched_set_tensor_backend(lctx.sched, cur, lctx.backend_cpu); + ggml_backend_sched_set_tensor_backend(lctx.sched[sub_gf_id], cur, lctx.backend_cpu); } } // norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends // FIXME: fix in ggml_backend_sched - const bool full_offload = lctx.model.n_gpu_layers > (int)lctx.model.hparams.n_layer; + const bool full_offload = lctx.model.n_gpu_layers > (int)n_layer_window; if (batch.n_tokens < 32 || full_offload) { if (il != -1 && strcmp(name, "norm") == 0) { for (auto * backend : lctx.backends) { - if (ggml_backend_supports_buft(backend, lctx.model.buft_layer[il].buft) && + int local_id = map_layer_to_local_id(il, n_world, my_rank, n_layer_window); + if (ggml_backend_supports_buft(backend, lctx.model.buft_layer[local_id].buft) && (ggml_backend_supports_op(backend, cur) || ggml_backend_offload_op(backend, cur))) { - ggml_backend_sched_set_tensor_backend(lctx.sched, cur, backend); + ggml_backend_sched_set_tensor_backend(lctx.sched[sub_gf_id], cur, backend); break; } } @@ -16278,12 +16416,14 @@ static struct ggml_cgraph * llama_build_graph( } }; - struct ggml_cgraph * result = NULL; + std::vector result = {}; struct llm_build_context llm(lctx, batch, cb, worst_case); llm.init(); + GGML_ASSERT(model.arch == LLM_ARCH_LLAMA && "this model is currently not supported"); + switch (model.arch) { case LLM_ARCH_LLAMA: case LLM_ARCH_GRANITE: @@ -16293,181 +16433,181 @@ static struct ggml_cgraph * llama_build_graph( } break; case LLM_ARCH_BAICHUAN: { - result = llm.build_baichuan(); + result.push_back(llm.build_baichuan()); } break; case LLM_ARCH_FALCON: { - result = llm.build_falcon(); + result.push_back(llm.build_falcon()); } break; case LLM_ARCH_GROK: { - result = llm.build_grok(); + result.push_back(llm.build_grok()); } break; case LLM_ARCH_STARCODER: { - result = llm.build_starcoder(); + result.push_back(llm.build_starcoder()); } break; case LLM_ARCH_REFACT: { - result = llm.build_refact(); + result.push_back(llm.build_refact()); } break; case LLM_ARCH_BERT: case LLM_ARCH_JINA_BERT_V2: case LLM_ARCH_NOMIC_BERT: { - result = llm.build_bert(); + result.push_back(llm.build_bert()); } break; case LLM_ARCH_BLOOM: { - result = llm.build_bloom(); + result.push_back(llm.build_bloom()); } break; case LLM_ARCH_MPT: { - result = llm.build_mpt(); + result.push_back(llm.build_mpt()); } break; case LLM_ARCH_STABLELM: { - result = llm.build_stablelm(); + result.push_back(llm.build_stablelm()); } break; case LLM_ARCH_QWEN: { - result = llm.build_qwen(); + result.push_back(llm.build_qwen()); } break; case LLM_ARCH_QWEN2: { - result = llm.build_qwen2(); + result.push_back(llm.build_qwen2()); } break; case LLM_ARCH_QWEN2MOE: { - result = llm.build_qwen2moe(); + result.push_back(llm.build_qwen2moe()); } break; case LLM_ARCH_PHI2: { - result = llm.build_phi2(); + result.push_back(llm.build_phi2()); } break; case LLM_ARCH_PHI3: { - result = llm.build_phi3(); + result.push_back(llm.build_phi3()); } break; case LLM_ARCH_PLAMO: { - result = llm.build_plamo(); + result.push_back(llm.build_plamo()); } break; case LLM_ARCH_GPT2: { - result = llm.build_gpt2(); + result.push_back(llm.build_gpt2()); } break; case LLM_ARCH_CODESHELL: { - result = llm.build_codeshell(); + result.push_back(llm.build_codeshell()); } break; case LLM_ARCH_ORION: { - result = llm.build_orion(); + result.push_back(llm.build_orion()); } break; case LLM_ARCH_INTERNLM2: { - result = llm.build_internlm2(); + result.push_back(llm.build_internlm2()); } break; case LLM_ARCH_MINICPM: { - result = llm.build_minicpm(); + result.push_back(llm.build_minicpm()); } break; case LLM_ARCH_MINICPM3: { - result = llm.build_minicpm3(); + result.push_back(llm.build_minicpm3()); } break; case LLM_ARCH_GEMMA: { - result = llm.build_gemma(); + result.push_back(llm.build_gemma()); } break; case LLM_ARCH_GEMMA2: { - result = llm.build_gemma2(); + result.push_back(llm.build_gemma2()); } break; case LLM_ARCH_STARCODER2: { - result = llm.build_starcoder2(); + result.push_back(llm.build_starcoder2()); } break; case LLM_ARCH_MAMBA: { - result = llm.build_mamba(); + result.push_back(llm.build_mamba()); } break; case LLM_ARCH_XVERSE: { - result = llm.build_xverse(); + result.push_back(llm.build_xverse()); } break; case LLM_ARCH_COMMAND_R: { - result = llm.build_command_r(); + result.push_back(llm.build_command_r()); } break; case LLM_ARCH_DBRX: { - result = llm.build_dbrx(); + result.push_back(llm.build_dbrx()); } break; case LLM_ARCH_OLMO: { - result = llm.build_olmo(); + result.push_back(llm.build_olmo()); } break; case LLM_ARCH_OLMOE: { - result = llm.build_olmoe(); + result.push_back(llm.build_olmoe()); } break; case LLM_ARCH_OPENELM: { - result = llm.build_openelm(); + result.push_back(llm.build_openelm()); } break; case LLM_ARCH_GPTNEOX: { - result = llm.build_gptneox(); + result.push_back(llm.build_gptneox()); } break; case LLM_ARCH_ARCTIC: { - result = llm.build_arctic(); + result.push_back(llm.build_arctic()); } break; case LLM_ARCH_DEEPSEEK2: { - result = llm.build_deepseek2(); + result.push_back(llm.build_deepseek2()); } break; case LLM_ARCH_CHATGLM: { - result = llm.build_chatglm(); + result.push_back(llm.build_chatglm()); } break; case LLM_ARCH_BITNET: { - result = llm.build_bitnet(); + result.push_back(llm.build_bitnet()); } break; case LLM_ARCH_T5: { if (lctx.is_encoding) { - result = llm.build_t5_encoder(); + result.push_back(llm.build_t5_encoder()); } else { - result = llm.build_t5_decoder(); + result.push_back(llm.build_t5_decoder()); } } break; case LLM_ARCH_T5ENCODER: { - result = llm.build_t5_encoder(); + result.push_back(llm.build_t5_encoder()); } break; case LLM_ARCH_JAIS: { - result = llm.build_jais(); + result.push_back(llm.build_jais()); } break; case LLM_ARCH_NEMOTRON: { - result = llm.build_nemotron(); + result.push_back(llm.build_nemotron()); } break; case LLM_ARCH_EXAONE: { - result = llm.build_exaone(); + result.push_back(llm.build_exaone()); } break; case LLM_ARCH_RWKV6: { - result = llm.build_rwkv6(); + result.push_back(llm.build_rwkv6()); } break; case LLM_ARCH_CHAMELEON: { - result = llm.build_chameleon(); + result.push_back(llm.build_chameleon()); } break; default: GGML_ABORT("fatal error"); @@ -16475,7 +16615,8 @@ static struct ggml_cgraph * llama_build_graph( // add on pooling layer if (lctx.cparams.embeddings) { - result = llm.append_pooling(result); + struct ggml_cgraph * pooled_sub_gf = llm.append_pooling(result.back()); + result.back() = pooled_sub_gf; } llm.free(); @@ -16531,6 +16672,29 @@ static int32_t llama_relative_position_bucket(llama_pos x, llama_pos y, uint64_t return relative_bucket; } +static void llama_view_data(float * data, int n=10) { + if (data != nullptr && n > 0) { + for (int i = 0; i < n; i++) LLAMA_LOG_INFO("%f ", data[i]); + LLAMA_LOG_INFO("\n"); + fflush(stdout); + } +} + +static void llama_view_tensor(struct ggml_tensor * tensor, int n=10, bool is_int32 = false) { + if (tensor != nullptr && tensor->data != nullptr && n > 0) { + LLAMA_LOG_INFO("[%s] ", tensor->name); + if (is_int32) { + int32_t *data = (int32_t*) tensor->data; + for (int i = 0; i < n; i++) LLAMA_LOG_INFO("%i ", data[i]); + } else { + float *data = (float*) tensor->data; + for (int i = 0; i < n; i++) LLAMA_LOG_INFO("%f ", data[i]); + } + LLAMA_LOG_INFO("\n"); + fflush(stdout); + } +} + static void llama_set_inputs(llama_context & lctx, const llama_ubatch & batch) { // // set input data @@ -16553,13 +16717,27 @@ static void llama_set_inputs(llama_context & lctx, const llama_ubatch & batch) { ggml_backend_tensor_set(lctx.inp_embd, batch.embd, 0, n_tokens*n_embd*ggml_element_size(lctx.inp_embd)); } + if (batch.backend_embd && lctx.backend_embd && lctx.backend_embd->data != nullptr) { + const int64_t n_embd = lctx.backend_embd->ne[0]; + const int64_t n_tokens = lctx.backend_embd->ne[1]; + + ggml_backend_tensor_set(lctx.backend_embd, batch.backend_embd, 0, n_tokens*n_embd*ggml_element_size(lctx.backend_embd)); + } + + if (batch.out_embd && lctx.out_embd) { + const int64_t n_embd = lctx.out_embd->ne[0]; + const int64_t n_output = lctx.out_embd->ne[1]; + + ggml_backend_tensor_set(lctx.out_embd, batch.out_embd, 0, n_output*n_embd*ggml_element_size(lctx.out_embd)); + } + if (batch.pos && lctx.inp_pos) { const int64_t n_tokens = batch.n_tokens; ggml_backend_tensor_set(lctx.inp_pos, batch.pos, 0, n_tokens*ggml_element_size(lctx.inp_pos)); } - if (hparams.causal_attn || cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) { + if (lctx.inp_out_ids && (hparams.causal_attn || cparams.pooling_type == LLAMA_POOLING_TYPE_NONE)) { GGML_ASSERT(lctx.inp_out_ids && "every model that can must skip unused outputs"); const int64_t n_tokens = batch.n_tokens; @@ -16949,8 +17127,8 @@ static size_t llama_output_reserve(llama_context & lctx, size_t n_outputs) { const bool has_logits = !cparams.embeddings; const bool has_embd = cparams.embeddings && (cparams.pooling_type == LLAMA_POOLING_TYPE_NONE); - const size_t logits_size = has_logits ? n_vocab*n_outputs_max : 0; - const size_t embd_size = has_embd ? n_embd*n_outputs_max : 0; + const size_t logits_size = has_logits ? n_vocab * n_outputs_max : 0; + const size_t embd_size = has_embd ? n_embd * n_outputs_max : 0; if (lctx.output_ids.empty()) { // init, never resized afterwards @@ -16983,12 +17161,12 @@ static size_t llama_output_reserve(llama_context & lctx, size_t n_outputs) { float * output_base = (float *) ggml_backend_buffer_get_base(lctx.buf_output); - lctx.logits = has_logits ? output_base : nullptr; - lctx.embd = has_embd ? output_base + logits_size : nullptr; + lctx.logits = has_logits ? output_base : nullptr; + lctx.embd = has_embd ? output_base + logits_size : nullptr; - lctx.output_size = n_outputs_max; - lctx.logits_size = logits_size; - lctx.embd_size = embd_size; + lctx.output_size = n_outputs_max; + lctx.logits_size = logits_size; + lctx.embd_size = embd_size; // set all ids as invalid (negative) std::fill(lctx.output_ids.begin(), lctx.output_ids.end(), -1); @@ -17041,6 +17219,7 @@ static void llama_output_reorder(struct llama_context * ctx) { static void llama_graph_compute( llama_context & lctx, ggml_cgraph * gf, + ggml_backend_sched_t sched, int n_threads, ggml_threadpool * threadpool) { if (lctx.backend_cpu != nullptr) { @@ -17053,12 +17232,107 @@ static void llama_graph_compute( set_n_threads_fn.second(set_n_threads_fn.first, n_threads); } - auto err = ggml_backend_sched_graph_compute_async(lctx.sched, gf); + auto err = ggml_backend_sched_graph_compute_async(sched, gf); if (err != GGML_STATUS_SUCCESS) { LLAMA_LOG_ERROR("%s: ggml_backend_sched_graph_compute_async failed with error %d\n", __func__, err); } +} - // fprintf(stderr, "splits: %d\n", ggml_backend_sched_get_n_splits(lctx.sched)); +struct input_tensors { + ggml_tensor * sub_gf_out = nullptr; + ggml_tensor * inp_pos = nullptr; +}; + +struct sync_meta { + int32_t n_tokens = 0; +}; + +static void llama_send_meta(zmq::socket_t & socket, struct sync_meta * meta) { + GGML_ASSERT(meta != nullptr); + try { + std::vector send_msgs; + + GGML_ASSERT(meta->n_tokens != 0); + send_msgs.emplace_back("n_tokens", strlen("n_tokens")); + send_msgs.emplace_back(&(meta->n_tokens), sizeof(meta->n_tokens)); + + zmq::send_multipart(socket, send_msgs); + } catch (const zmq::error_t& e) { + LLAMA_LOG_INFO("Failed to send meta data: %s\n", e.what()); + } +} + +static void llama_send_tensors(zmq::socket_t & socket, struct input_tensors * tensors) { + try { + std::vector send_msgs; + size_t buff_size = 0; + + send_msgs.emplace_back("sub_gf_out", strlen("sub_gf_out")); + send_msgs.emplace_back(tensors->sub_gf_out->ne, sizeof(tensors->sub_gf_out->ne)); + buff_size = tensors->sub_gf_out->ne[0] * tensors->sub_gf_out->ne[1] * sizeof(float); + send_msgs.emplace_back(tensors->sub_gf_out->data, buff_size); + + if (tensors->inp_pos) { + send_msgs.emplace_back("inp_pos", strlen("inp_pos")); + send_msgs.emplace_back(tensors->inp_pos->ne, sizeof(tensors->inp_pos->ne[0])); + buff_size = tensors->inp_pos->ne[0] * sizeof(int32_t); + send_msgs.emplace_back(tensors->inp_pos->data, buff_size); + } + + zmq::send_multipart(socket, send_msgs); + } catch (const zmq::error_t& e) { + LLAMA_LOG_INFO("Failed to send tensor data: %s\n", e.what()); + } +} + +static int llama_recv_meta(zmq::socket_t & socket, struct sync_meta * meta) { + socket.setsockopt(ZMQ_RCVTIMEO, 1000); + + std::vector recv_msgs; + if (!zmq::recv_multipart(socket, std::back_inserter(recv_msgs))) { + return -1; + } + + socket.setsockopt(ZMQ_RCVTIMEO, -1); + + for (size_t i = 0; i < recv_msgs.size(); i += 2) { + std::string key = recv_msgs[i].to_string(); + zmq::message_t & data_msg = recv_msgs[i + 1]; + + if (key == "n_tokens") { + GGML_ASSERT(data_msg.size() == sizeof(meta->n_tokens)); + std::memcpy(&(meta->n_tokens), data_msg.data(), sizeof(meta->n_tokens)); + } + } + return 0; +} + +static void llama_recv_tensors(zmq::socket_t & socket, input_tensors * tensors) { + std::vector recv_msgs; + if (!zmq::recv_multipart(socket, std::back_inserter(recv_msgs))) { + LLAMA_LOG_INFO("Failed to receive tensor data.\n"); + } + + for (size_t i = 0; i < recv_msgs.size(); i += 3) { + std::string key = recv_msgs[i].to_string(); + zmq::message_t &dims_msg = recv_msgs[i + 1]; + zmq::message_t &data_msg = recv_msgs[i + 2]; + + if (key == "sub_gf_out" && tensors->sub_gf_out) { + int64_t * dims = static_cast(dims_msg.data()); + size_t buff_size = dims[0] * dims[1] * sizeof(float); + GGML_ASSERT(dims[0] == tensors->sub_gf_out->ne[0]); + GGML_ASSERT(dims[1] == tensors->sub_gf_out->ne[1]); + GGML_ASSERT(data_msg.size() == buff_size); + std::memcpy(tensors->sub_gf_out->data, data_msg.data(), buff_size); + } else if (key == "inp_pos" && tensors->inp_pos) { + int64_t * dims = static_cast(dims_msg.data()); + size_t buff_size = dims[0] * sizeof(int32_t); + GGML_ASSERT(dims[0] == tensors->inp_pos->ne[0]); + GGML_ASSERT(data_msg.size() == buff_size); + std::memcpy(tensors->inp_pos->data, data_msg.data(), buff_size); + } + } } // decode a batch of tokens by evaluating the transformer @@ -17073,20 +17347,19 @@ static void llama_graph_compute( static int llama_decode_internal( llama_context & lctx, llama_batch batch_all) { // TODO: rename back to batch + const auto & model = lctx.model; + const auto & hparams = model.hparams; + const auto & cparams = lctx.cparams; + const uint32_t n_world = cparams.n_world; + const uint32_t my_rank = cparams.rank; lctx.is_encoding = false; const uint32_t n_tokens_all = batch_all.n_tokens; - - if (n_tokens_all == 0) { - LLAMA_LOG_ERROR("%s: n_tokens == 0\n", __func__); - return -1; + if (my_rank != 0) { + batch_all.token = nullptr; } - const auto & model = lctx.model; - const auto & hparams = model.hparams; - const auto & cparams = lctx.cparams; - - GGML_ASSERT((!batch_all.token && batch_all.embd) || (batch_all.token && !batch_all.embd)); // NOLINT + GGML_ASSERT(!(my_rank == 0 && n_tokens_all == 0) && "n_tokens == 0 on master node"); if (batch_all.token) { for (uint32_t i = 0; i < n_tokens_all; ++i) { @@ -17133,17 +17406,31 @@ static int llama_decode_internal( n_outputs = 1; } + sync_meta meta; + if (my_rank != 0) { + if (llama_recv_meta(*lctx.recv_socket, &meta) == -1) { + return -1; + } + batch_all.n_tokens = meta.n_tokens; + } + + if (my_rank != n_world - 1) { + meta.n_tokens = batch_all.n_tokens; + llama_send_meta(*lctx.send_socket, &meta); + } + lctx.sbatch.from_batch(batch_all, n_embd, /* simple_split */ !kv_self.recurrent, /* logits_all */ n_outputs == n_tokens_all); // reserve output buffer - if (llama_output_reserve(lctx, n_outputs) < n_outputs) { + if (my_rank == 0 && llama_output_reserve(lctx, n_outputs) < n_outputs) { LLAMA_LOG_ERROR("%s: could not reserve space for batch with %u outputs\n", __func__, n_outputs); return -2; }; - while (lctx.sbatch.n_tokens > 0) { + { // assume there is only one batch + // while (lctx.sbatch.n_tokens > 0) { // handle multiple batches llama_ubatch ubatch; if (kv_self.recurrent) { if (embd_pooled) { @@ -17160,9 +17447,9 @@ static int llama_decode_internal( const uint32_t n_tokens = ubatch.n_tokens; // count the outputs in this u_batch - { - int32_t n_outputs_new = 0; - + int32_t n_outputs_new = 0; + + if (my_rank == 0) { if (n_outputs == n_tokens_all) { n_outputs_new = n_tokens; } else { @@ -17171,11 +17458,13 @@ static int llama_decode_internal( n_outputs_new += (int32_t) (ubatch.output[i] != 0); } } - - // needs to happen before the graph is built - lctx.n_outputs = n_outputs_new; + } else { + n_outputs_new += 1; } + // needs to happen before the graph is built + lctx.n_outputs = n_outputs_new; + int n_threads = n_tokens == 1 ? cparams.n_threads : cparams.n_threads_batch; ggml_threadpool_t threadpool = n_tokens == 1 ? lctx.threadpool : lctx.threadpool_batch; @@ -17187,7 +17476,7 @@ static int llama_decode_internal( // if we have enough unused cells before the current head -> // better to start searching from the beginning of the cache, hoping to fill it - if (kv_self.head > kv_self.used + 2*n_tokens) { + if (kv_self.head > kv_self.used + 2 * n_tokens) { kv_self.head = 0; } @@ -17201,21 +17490,26 @@ static int llama_decode_internal( // if we start defragmenting the cache, the benefit from this will be more important const uint32_t pad = llama_kv_cache_get_padding(cparams); kv_self.n = std::min(kv_self.size, std::max(pad, GGML_PAD(llama_kv_cache_cell_max(kv_self), pad))); - //kv_self.n = llama_kv_cache_cell_max(kv_self); } } - //printf("kv_self.n = %5d, kv_self.used = %5d, kv_self.head = %5d\n", kv_self.n, kv_self.used, kv_self.head); + for (size_t i = 0; i < lctx.sched.size(); i++) { + ggml_backend_sched_reset(lctx.sched[i]); + ggml_backend_sched_set_eval_callback(lctx.sched[i], lctx.cparams.cb_eval, lctx.cparams.cb_eval_user_data); + } - ggml_backend_sched_reset(lctx.sched); - ggml_backend_sched_set_eval_callback(lctx.sched, lctx.cparams.cb_eval, lctx.cparams.cb_eval_user_data); - - ggml_cgraph * gf = llama_build_graph(lctx, ubatch, false); + std::vector gf = llama_build_graph(lctx, ubatch, false); // the output is always the last tensor in the graph - struct ggml_tensor * res = ggml_graph_node(gf, -1); - struct ggml_tensor * embd = ggml_graph_node(gf, -2); - + struct ggml_tensor * res = nullptr; + struct ggml_tensor * embd = nullptr; + struct ggml_tensor * sub_gf_out = nullptr; + const int64_t n_embd = hparams.n_embd; + if (my_rank == 0) { + res = ggml_graph_node(gf.back(), -1); + embd = ggml_graph_node(gf.back(), -2); + } + if (lctx.n_outputs == 0) { // no output res = nullptr; @@ -17223,24 +17517,90 @@ static int llama_decode_internal( } else if (cparams.embeddings) { res = nullptr; // do not extract logits for embedding case embd = nullptr; - for (int i = ggml_graph_n_nodes(gf) - 1; i >= 0; --i) { - if (strcmp(ggml_graph_node(gf, i)->name, "result_embd_pooled") == 0) { - embd = ggml_graph_node(gf, i); + for (int i = ggml_graph_n_nodes(gf.back()) - 1; i >= 0; --i) { + if (strcmp(ggml_graph_node(gf.back(), i)->name, "result_embd_pooled") == 0) { + embd = ggml_graph_node(gf.back(), i); break; } } GGML_ASSERT(embd != nullptr && "missing embeddings tensor"); } else { embd = nullptr; // do not extract embeddings when not needed - GGML_ASSERT(strcmp(res->name, "result_output") == 0 && "missing result_output tensor"); } - // LLAMA_LOG_INFO("graph build time: %.3f ms (%d nodes, %d leafs)\n", (ggml_time_us() - t_start_us)/1000.0, gf->n_nodes, gf->n_leafs); + + GGML_ASSERT(lctx.sched.size() == gf.size()); + for (size_t i = 0; i < (size_t)lctx.sched.size(); ++i) { + ggml_backend_sched_alloc_graph(lctx.sched[i], gf[i]); + } - ggml_backend_sched_alloc_graph(lctx.sched, gf); + const uint32_t n_layer = hparams.n_layer; + const char * layer_str = nullptr; + int cur_l = -1; + bool is_output = false; + bool is_last_l = false; + GGML_ASSERT(my_rank == 0 || n_world > 1); + + for (size_t i = 0; i < (size_t)gf.size(); ++i) { + ggml_cgraph * sub_gf = gf[i]; - llama_set_inputs(lctx, ubatch); + if (n_world > 1 && !(my_rank == 0 && i == 0) && !(my_rank == 0 && is_last_l)) { + { // receive data from previous nodes + timer(llama_recv_tensors); + input_tensors tensors; + const bool is_out_embd = my_rank == 0 && i == (size_t)gf.size() - 1; + tensors.sub_gf_out = is_out_embd ? lctx.out_embd : lctx.backend_embd; + tensors.inp_pos = lctx.inp_pos; + llama_recv_tensors(*lctx.recv_socket, &tensors); + + is_last_l = my_rank == 0 && i == (size_t)gf.size() - 1; + size_t buff_size = tensors.sub_gf_out->ne[0] * tensors.sub_gf_out->ne[1] * ggml_element_size(tensors.sub_gf_out); + if (!is_last_l) { + memcpy(ubatch.backend_embd, tensors.sub_gf_out->data, buff_size); + } else { + memcpy(ubatch.out_embd, tensors.sub_gf_out->data, buff_size); + } + if (my_rank != 0 && i == 0) { + buff_size = tensors.inp_pos->ne[0] * ggml_element_size(tensors.inp_pos); + memcpy(ubatch.pos, tensors.inp_pos->data, buff_size); + } + } + } - llama_graph_compute(lctx, gf, n_threads, threadpool); + llama_set_inputs(lctx, ubatch); + + { // compute graph + timer(llama_graph_compute); + llama_graph_compute(lctx, sub_gf, lctx.sched[i], n_threads, threadpool); + } + + sub_gf_out = ggml_graph_node(sub_gf, -1); + is_output = strcmp(sub_gf_out->name, "result_output") == 0; + if (is_output) { + break; + } + + layer_str = strchr(sub_gf_out->name, '-') + 1; + cur_l = std::atoi(layer_str); + is_last_l = (cur_l == static_cast(n_layer) - 1); + + // send the result to the next node (or the master) + if (n_world == 1 || (my_rank == 0 && is_last_l)) { + size_t buf_size = sub_gf_out->ne[0]*sub_gf_out->ne[1]*sizeof(float); + float * embd_buf = is_last_l ? ubatch.out_embd : ubatch.backend_embd; + memcpy(embd_buf, sub_gf_out->data, buf_size); + } else { + input_tensors tensors; + tensors.sub_gf_out = sub_gf_out; + if (i == 0 && !is_last_l && my_rank != n_world - 1) { + tensors.inp_pos = lctx.inp_pos; + const size_t buff_size = ubatch.n_tokens * ggml_element_size(tensors.inp_pos); + memcpy(tensors.inp_pos->data, ubatch.pos, buff_size); + } + const bool is_to_master = my_rank !=0 && is_last_l; + zmq::socket_t * s = is_to_master ? lctx.master_socket : lctx.send_socket; + llama_send_tensors(*s, &tensors); + } + } // update the kv ring buffer { @@ -17252,30 +17612,27 @@ static int llama_decode_internal( } } - // plot the computation graph in dot format (for debugging purposes) - //if (n_past%100 == 0) { - // ggml_graph_dump_dot(gf, NULL, "llama.dot"); - //} - // extract logits if (res) { - ggml_backend_t backend_res = ggml_backend_sched_get_tensor_backend(lctx.sched, res); + ggml_backend_t backend_res = ggml_backend_sched_get_tensor_backend(lctx.sched.back(), res); GGML_ASSERT(backend_res != nullptr); GGML_ASSERT(lctx.logits != nullptr); - float * logits_out = lctx.logits + n_outputs_prev*n_vocab; + float * logits_out = lctx.logits + n_outputs_prev * n_vocab; const int32_t n_outputs_new = lctx.n_outputs; if (n_outputs_new) { GGML_ASSERT( n_outputs_prev + n_outputs_new <= n_outputs); - GGML_ASSERT((n_outputs_prev + n_outputs_new)*n_vocab <= (int64_t) lctx.logits_size); - ggml_backend_tensor_get_async(backend_res, res, logits_out, 0, n_outputs_new*n_vocab*sizeof(float)); + GGML_ASSERT((n_outputs_prev + n_outputs_new) * n_vocab <= (int64_t) lctx.logits_size); + ggml_backend_tensor_get_async(backend_res, res, logits_out, 0, n_outputs_new * n_vocab * sizeof(float)); } } // extract embeddings if (embd) { - ggml_backend_t backend_embd = ggml_backend_sched_get_tensor_backend(lctx.sched, embd); + throw std::runtime_error("embd is currently not supported"); + + ggml_backend_t backend_embd = ggml_backend_sched_get_tensor_backend(lctx.sched.back(), embd); // todo. GGML_ASSERT(backend_embd != nullptr); switch (cparams.pooling_type) { @@ -17331,10 +17688,9 @@ static int llama_decode_internal( n_outputs_prev += lctx.n_outputs; } - // set output mappings - { + if (my_rank == 0) { + // set output mappings bool sorted_output = true; - GGML_ASSERT(lctx.sbatch.out_ids.size() == n_outputs); for (size_t i = 0; i < n_outputs; ++i) { @@ -17348,29 +17704,28 @@ static int llama_decode_internal( if (sorted_output) { lctx.sbatch.out_ids.clear(); } + + // set to total number of outputs in the batch, for use in llama_get_logits_ith + lctx.n_outputs = n_outputs; } - // set to total number of outputs in the batch, for use in llama_get_logits_ith - lctx.n_outputs = n_outputs; - // wait for the computation to finish (automatically done when obtaining the model output) - //llama_synchronize(&lctx); + // llama_synchronize(&lctx); // decide if we need to defrag the kv cache if (cparams.causal_attn && cparams.defrag_thold >= 0.0f) { const float fragmentation = kv_self.n >= 128 ? 1.0f - float(kv_self.used)/float(kv_self.n) : 0.0f; - // queue defragmentation for next llama_kv_cache_update if (fragmentation > cparams.defrag_thold) { - //LLAMA_LOG_INFO("fragmentation: %.2f\n", fragmentation); - llama_kv_cache_defrag(kv_self); } } // Reset state for the next token before backend sync, to allow the CPU activities in the reset to // overlap with device computation. - ggml_backend_sched_reset(lctx.sched); + for (size_t i = 0; i < (size_t)lctx.sched.size(); ++i) { + ggml_backend_sched_reset(lctx.sched[i]); + } return 0; } @@ -17445,10 +17800,10 @@ static int llama_encode_internal( GGML_ASSERT(n_threads > 0); - ggml_backend_sched_reset(lctx.sched); - ggml_backend_sched_set_eval_callback(lctx.sched, lctx.cparams.cb_eval, lctx.cparams.cb_eval_user_data); + ggml_backend_sched_reset(lctx.sched.at(0)); // todo. + ggml_backend_sched_set_eval_callback(lctx.sched.at(0), lctx.cparams.cb_eval, lctx.cparams.cb_eval_user_data); // todo. - ggml_cgraph * gf = llama_build_graph(lctx, ubatch, false); + std::vector gf = llama_build_graph(lctx, ubatch, false); // the output embeddings after the final encoder normalization struct ggml_tensor * embd = nullptr; @@ -17456,29 +17811,29 @@ static int llama_encode_internal( // there are two cases here if (llama_model_has_decoder(&lctx.model)) { // first case is an encoder-decoder T5 model where embeddings are passed to decoder - embd = ggml_graph_node(gf, -1); + embd = ggml_graph_node(gf.back(), -1); // todo. GGML_ASSERT(strcmp(embd->name, "result_norm") == 0 && "missing result_output tensor"); } else { // second case is an encoder-only T5 model if (cparams.embeddings) { // only output embeddings if required - embd = ggml_graph_node(gf, -1); + embd = ggml_graph_node(gf.back(), -1); // todo. if (strcmp(embd->name, "result_embd_pooled") != 0) { - embd = ggml_graph_node(gf, -2); + embd = ggml_graph_node(gf.back(), -2); // todo. } GGML_ASSERT(strcmp(embd->name, "result_embd_pooled") == 0 && "missing embeddings tensor"); } } - ggml_backend_sched_alloc_graph(lctx.sched, gf); + ggml_backend_sched_alloc_graph(lctx.sched.at(0), gf.at(0)); // todo. llama_set_inputs(lctx, ubatch); - llama_graph_compute(lctx, gf, n_threads, threadpool); + llama_graph_compute(lctx, gf.at(0), lctx.sched.at(0), n_threads, threadpool); // todo. // extract embeddings if (embd) { - ggml_backend_t backend_embd = ggml_backend_sched_get_tensor_backend(lctx.sched, embd); + ggml_backend_t backend_embd = ggml_backend_sched_get_tensor_backend(lctx.sched.at(0), embd); // todo. GGML_ASSERT(backend_embd != nullptr); if (llama_model_has_decoder(&lctx.model)) { @@ -17545,7 +17900,7 @@ static int llama_encode_internal( // Reset state for the next token before backend sync, to allow the CPU activities in the reset to // overlap with device computation. - ggml_backend_sched_reset(lctx.sched); + ggml_backend_sched_reset(lctx.sched.at(0)); // todo. return 0; } @@ -17759,11 +18114,11 @@ static void llama_kv_cache_defrag_internal(struct llama_context & lctx) { #else // ggml_graph defrag - ggml_backend_sched_reset(lctx.sched); + ggml_backend_sched_reset(lctx.sched.at(0)); // todo. ggml_cgraph * gf = llama_build_graph_defrag(lctx, ids); - llama_graph_compute(lctx, gf, lctx.cparams.n_threads, lctx.threadpool); + llama_graph_compute(lctx, gf, lctx.sched.at(0), lctx.cparams.n_threads, lctx.threadpool); // todo. #endif //const int64_t t_end = ggml_time_us(); @@ -17781,15 +18136,15 @@ static void llama_kv_cache_update_internal(struct llama_context & lctx) { } { - ggml_backend_sched_reset(lctx.sched); + ggml_backend_sched_reset(lctx.sched.at(0)); // todo. ggml_cgraph * gf = llama_build_graph_k_shift(lctx); - ggml_backend_sched_alloc_graph(lctx.sched, gf); + ggml_backend_sched_alloc_graph(lctx.sched.at(0), gf); // todo. llama_set_k_shift(lctx); - llama_graph_compute(lctx, gf, lctx.cparams.n_threads, lctx.threadpool); + llama_graph_compute(lctx, gf, lctx.sched.at(0), lctx.cparams.n_threads, lctx.threadpool); // todo. need_reserve = true; } @@ -17821,12 +18176,18 @@ static void llama_kv_cache_update_internal(struct llama_context & lctx) { uint32_t n_seqs = 1; // TODO: worst-case number of sequences uint32_t n_tokens = std::min(lctx.cparams.n_ctx, lctx.cparams.n_ubatch); llama_token token = llama_token_bos(&lctx.model); // not actually used by llama_build_graph, but required to choose between token and embedding inputs graph - llama_ubatch ubatch = { true, n_tokens, n_tokens / n_seqs, n_seqs, &token, nullptr, nullptr, nullptr, nullptr, nullptr}; - ggml_cgraph * gf = llama_build_graph(lctx, ubatch, true); + llama_ubatch ubatch = { true, n_tokens, n_tokens / n_seqs, n_seqs, &token, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr}; + std::vector gf = llama_build_graph(lctx, ubatch, true); // initialize scheduler with the worst-case graph - ggml_backend_sched_reset(lctx.sched); - if (!ggml_backend_sched_reserve(lctx.sched, gf)) { + ggml_backend_sched_reset(lctx.sched.at(0)); // todo. + + bool ok = true; + GGML_ASSERT(lctx.sched.size() == gf.size()); + for (size_t i = 0; i < gf.size(); ++i) { + ok = ok & ggml_backend_sched_reserve(lctx.sched[i], gf[i]); + } + if (!ok) { LLAMA_LOG_ERROR("%s: failed to allocate compute buffers\n", __func__); } } @@ -18923,6 +19284,9 @@ void llama_lora_adapter_free(struct llama_lora_adapter * adapter) { // struct llama_model_params llama_model_default_params() { struct llama_model_params result = { + /*.n_world =*/ 1, + /*.rank =*/ 0, + /*.n_layer_window =*/ 32, /*.n_gpu_layers =*/ 0, /*.split_mode =*/ LLAMA_SPLIT_MODE_LAYER, /*.main_gpu =*/ 0, @@ -18947,6 +19311,11 @@ struct llama_model_params llama_model_default_params() { struct llama_context_params llama_context_default_params() { struct llama_context_params result = { + /*.n_world =*/ 1, + /*.rank =*/ 0, + /*.n_layer_window =*/ 32, + /*.master_ip =*/ nullptr, + /*.next_node_ip =*/ nullptr, /*.n_ctx =*/ 512, /*.n_batch =*/ 2048, /*.n_ubatch =*/ 512, @@ -19067,9 +19436,7 @@ int64_t llama_time_us(void) { return ggml_time_us(); } -struct llama_model * llama_load_model_from_file( - const char * path_model, - struct llama_model_params params) { +struct llama_model * llama_load_model_from_file(const char * path_model, struct llama_model_params params) { ggml_time_init(); llama_model * model = new llama_model; @@ -19140,6 +19507,74 @@ void llama_free_model(struct llama_model * model) { delete model; } +static uint32_t map_rank_to_port(uint32_t rank, uint32_t data_port) { + return data_port + rank; +} + +void llama_init_sockets(struct llama_context * ctx, uint32_t n_world, uint32_t my_rank) { + if (n_world == 1) { + return; + } + + ctx->sock_context = new zmq::context_t(2); + ctx->send_socket = new zmq::socket_t(*ctx->sock_context, zmq::socket_type::push); + ctx->recv_socket = new zmq::socket_t(*ctx->sock_context, zmq::socket_type::pull); + ctx->signal_socket = new zmq::socket_t(*ctx->sock_context, zmq::socket_type::pull); + + if (my_rank != 0 && my_rank != (n_world - 1)) { + ctx->master_socket = new zmq::socket_t(*ctx->sock_context, zmq::socket_type::push); + } else if (my_rank == (n_world - 1)) { + ctx->master_socket = ctx->send_socket; // No need for master_socket in the last rank, reuse send_socket for communication + } + + const uint32_t next_rank = (my_rank + 1) % n_world; + std::string recv_endp = "tcp://*:" + std::to_string(map_rank_to_port(my_rank, ctx->data_port)); + std::string send_endp = "tcp://" + ctx->next_node_ip + ":" + std::to_string(map_rank_to_port(next_rank, ctx->data_port)); + std::string master_endp = "tcp://" + ctx->master_ip + ":" + std::to_string(map_rank_to_port(0, ctx->data_port)); + std::string signal_endp = "tcp://*:" + std::to_string(map_rank_to_port(my_rank, ctx->signal_port)); + + try { + ctx->recv_socket->bind(recv_endp); + ctx->signal_socket->bind(signal_endp); + ctx->send_socket->connect(send_endp); + if (ctx->master_socket && my_rank != (n_world - 1)) { + ctx->master_socket->connect(master_endp); + } + } catch (const zmq::error_t &e) { + LLAMA_LOG_INFO("Error binding/connecting recv socket to endpoint: %s", e.what()); + exit(1); + } +} + +void llama_free_sockets(struct llama_context * ctx, char ** msg) { + const uint32_t n_world = ctx->cparams.n_world; + const uint32_t my_rank = ctx->cparams.rank; + const uint32_t next_rank = (my_rank + 1) % n_world; + + if (n_world == 1) { + return; + } + + zmq::socket_t signal_sender(*ctx->sock_context, zmq::socket_type::push); + std::string endp = "tcp://" + ctx->next_node_ip + ":" + std::to_string(map_rank_to_port(next_rank, ctx->signal_port)); + signal_sender.connect(endp); + + if (my_rank == 0) { + signal_sender.send(zmq::message_t("STOP", 4), zmq::send_flags::dontwait); + } + + zmq::message_t recv_msg; + if (ctx->signal_socket->recv(recv_msg, zmq::recv_flags::none)) { + std::string msg_str(static_cast(recv_msg.data()), recv_msg.size()); + if (msg_str == "STOP" && my_rank != 0) { + signal_sender.send(recv_msg, zmq::send_flags::dontwait); + } + delete[] *msg; // clean up the old message if any + *msg = new char[msg_str.size() + 1]; + std::strcpy(*msg, msg_str.c_str()); + } +} + struct llama_context * llama_new_context_with_model( struct llama_model * model, struct llama_context_params params) { @@ -19179,6 +19614,9 @@ struct llama_context * llama_new_context_with_model( const auto & hparams = model->hparams; auto & cparams = ctx->cparams; + cparams.n_world = params.n_world; + cparams.rank = params.rank; + cparams.n_layer_window = params.n_layer_window; cparams.n_seq_max = std::max(1u, params.n_seq_max); cparams.n_threads = params.n_threads; cparams.n_threads_batch = params.n_threads_batch; @@ -19249,6 +19687,10 @@ struct llama_context * llama_new_context_with_model( cparams.causal_attn = params.attention_type == LLAMA_ATTENTION_TYPE_CAUSAL; } + LLAMA_LOG_INFO("\n"); + LLAMA_LOG_INFO("%s: n_world = %u\n", __func__, cparams.n_world); + LLAMA_LOG_INFO("%s: rank = %u\n", __func__, cparams.rank); + LLAMA_LOG_INFO("%s: n_layer_win= %u\n", __func__, cparams.n_layer_window); LLAMA_LOG_INFO("%s: n_ctx = %u\n", __func__, cparams.n_ctx); LLAMA_LOG_INFO("%s: n_batch = %u\n", __func__, cparams.n_batch); LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch); @@ -19256,6 +19698,13 @@ struct llama_context * llama_new_context_with_model( LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base); LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale); + ctx->master_ip = params.master_ip; + ctx->next_node_ip = params.next_node_ip; + + LLAMA_LOG_INFO("\n"); + LLAMA_LOG_INFO("%s: master_ip = %s\n", __func__, ctx->master_ip.c_str()); + LLAMA_LOG_INFO("%s: next_node_ip = %s\n", __func__, ctx->next_node_ip.c_str()); + ctx->abort_callback = params.abort_callback; ctx->abort_callback_data = params.abort_callback_data; @@ -19470,8 +19919,9 @@ struct llama_context * llama_new_context_with_model( ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f)); } - // graph outputs buffer - { + // graph outputs buffer, reserve for rank 0 only + const uint32_t my_rank = params.rank; + if (my_rank == 0) { // resized during inference when a batch uses more outputs if (llama_output_reserve(*ctx, params.n_seq_max) < params.n_seq_max) { LLAMA_LOG_ERROR("%s: failed to reserve initial output buffer\n", __func__); @@ -19533,21 +19983,26 @@ struct llama_context * llama_new_context_with_model( } } - ctx->sched = ggml_backend_sched_new(ctx->backends.data(), backend_buft.data(), ctx->backends.size(), max_nodes, pipeline_parallel); - - if (pipeline_parallel) { - LLAMA_LOG_INFO("%s: pipeline parallelism enabled (n_copies=%d)\n", __func__, ggml_backend_sched_get_n_copies(ctx->sched)); + for (int i = 0; i < 10; ++i) { + ctx->sched.push_back(ggml_backend_sched_new(ctx->backends.data(), backend_buft.data(), ctx->backends.size(), max_nodes, pipeline_parallel)); } - + // build worst-case graph uint32_t n_seqs = 1; // TODO: worst-case number of sequences uint32_t n_tokens = std::min(cparams.n_ctx, cparams.n_ubatch); llama_token token = llama_token_bos(&ctx->model); // not actually used by llama_build_graph, but required to choose between token and embedding inputs graph - llama_ubatch ubatch = { true, n_tokens, n_tokens / n_seqs, n_seqs, &token, nullptr, nullptr, nullptr, nullptr, nullptr}; - ggml_cgraph * gf = llama_build_graph(*ctx, ubatch, true); + llama_ubatch ubatch = { true, n_tokens, n_tokens / n_seqs, n_seqs, &token, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr}; + std::vector gf = llama_build_graph(*ctx, ubatch, true); + + ctx->sched.resize(gf.size()); // initialize scheduler with the worst-case graph - if (!ggml_backend_sched_reserve(ctx->sched, gf)) { + bool ok = true; + GGML_ASSERT(ctx->sched.size() == gf.size()); + for (size_t i = 0; i < gf.size(); ++i) { + ok = ok & ggml_backend_sched_reserve(ctx->sched[i], gf[i]); + } + if (!ok) { LLAMA_LOG_ERROR("%s: failed to allocate compute buffers\n", __func__); llama_free(ctx); return nullptr; @@ -19556,18 +20011,17 @@ struct llama_context * llama_new_context_with_model( for (size_t i = 0; i < ctx->backends.size(); i++) { ggml_backend_t backend = ctx->backends[i]; ggml_backend_buffer_type_t buft = backend_buft[i]; - size_t size = ggml_backend_sched_get_buffer_size(ctx->sched, backend); - if (size > 1) { - LLAMA_LOG_INFO("%s: %10s compute buffer size = %8.2f MiB\n", __func__, + + size_t total_size = 0; + for (size_t j = 0; j < ctx->sched.size(); j++) { + total_size += ggml_backend_sched_get_buffer_size(ctx->sched.at(j), backend); + } + if (total_size > 1) { + LLAMA_LOG_INFO("%s: %10s compute buffer size = %8.2f MiB (in total)\n", __func__, ggml_backend_buft_name(buft), - size / 1024.0 / 1024.0); + total_size / 1024.0 / 1024.0); } } - - // note: the number of splits during measure is higher than during inference due to the kv shift - int n_splits = ggml_backend_sched_get_n_splits(ctx->sched); - LLAMA_LOG_INFO("%s: graph nodes = %d\n", __func__, ggml_graph_n_nodes(gf)); - LLAMA_LOG_INFO("%s: graph splits = %d\n", __func__, n_splits); } } @@ -21139,16 +21593,11 @@ int32_t llama_encode( int32_t llama_decode( struct llama_context * ctx, struct llama_batch batch) { - const int ret = llama_decode_internal(*ctx, batch); - if (ret < 0) { - LLAMA_LOG_ERROR("%s: failed to decode, ret = %d\n", __func__, ret); - } - - return ret; + return llama_decode_internal(*ctx, batch); } void llama_synchronize(struct llama_context * ctx) { - ggml_backend_sched_synchronize(ctx->sched); + ggml_backend_sched_synchronize(ctx->sched.at(0)); // todo. // FIXME: if multiple single tokens are evaluated without a synchronization, // the stats will be added to the prompt evaluation stats