diff --git a/Makefile b/Makefile index fe3f3853..0a97e269 100644 --- a/Makefile +++ b/Makefile @@ -3,6 +3,8 @@ BUILD_TARGETS = \ llama-perplexity \ llama-server \ llama-cli \ + llama-speculative \ + llama-gguf-split \ profile-tool # BUILD_TARGETS = \ @@ -951,7 +953,8 @@ OBJ_LLAMA = \ src/llama-grammar.o \ src/llama-sampling.o \ src/unicode.o \ - src/unicode-data.o + src/unicode-data.o \ + src/network-utils.o \ OBJ_COMMON = \ common/profiler.o \ @@ -961,6 +964,7 @@ OBJ_COMMON = \ common/console.o \ common/ngram-cache.o \ common/sampling.o \ + common/speculative.o \ common/train.o \ common/build-info.o \ common/json-schema-to-grammar.o @@ -1140,6 +1144,11 @@ src/unicode-data.o: \ src/unicode-data.cpp \ src/unicode-data.h $(CXX) $(CXXFLAGS) -c $< -o $@ + +src/network-utils.o: \ + src/network-utils.cpp \ + src/network-utils.h + $(CXX) $(CXXFLAGS) -c $< -o $@ src/llama.o: \ src/llama.cpp \ @@ -1148,6 +1157,7 @@ src/llama.o: \ src/llama-grammar.h \ src/llama-sampling.h \ src/unicode.h \ + src/network-utils.h \ include/llama.h \ ggml/include/ggml-cuda.h \ ggml/include/ggml-metal.h \ @@ -1231,6 +1241,13 @@ common/json-schema-to-grammar.o: \ common/json-schema-to-grammar.h $(CXX) $(CXXFLAGS) -c $< -o $@ +# speculative +common/speculative.o: \ + common/speculative.cpp \ + common/speculative.h \ + include/llama.h + $(CXX) $(CXXFLAGS) -c $< -o $@ + common/train.o: \ common/train.cpp \ common/train.h diff --git a/README.md b/README.md index 2b11d98d..2f891a8b 100644 --- a/README.md +++ b/README.md @@ -34,26 +34,26 @@ And, if your devices are more powerful, you could unlock even more possibilities > Device D4 runs inside a Termux-simulated Linux. Device D1 reads disk data in random mode and D2~D4 read in sequential mode. -**Table 2:** Token latency for Llama models (w/o device selection). +**Table 2:** Token latency for Llama models (with device selection). | **Model** | **llama.cpp** | **exo** | **dllama** | **prima.cpp** | -|-----------------|---------------|-----------|------------|---------------| -| Llama 3-8B | **15 ms** | 263 ms | 459 ms | 54 ms | -| Llama 3-14B | **20 ms** | - | - | 65 ms | +|----------------|---------------|-----------|------------|---------------| +| Llama 3-8B | 15 ms | 263 ms | 459 ms | **15 ms** | +| Llama 3-14B | 20 ms | - | - | **20 ms** | | Llama 1-30B | 202 ms | - | - | **72 ms** | | Llama 3-45B | 328 ms | - | - | **233 ms** | | Llama 3-60B | 7965 ms | - | - | **468 ms** | | Llama 1-65B | 8807 ms | - | - | **569 ms** | | Llama 3-70B | 10120 ms | OOM | OOM | **674 ms** | -**Table 3:** Token latency for Qwen 2.5, QwQ, and DeepSeek R1 models (w/o device selection). +**Table 3:** Token latency for Qwen 2.5, QwQ, and DeepSeek R1 models (with device selection). | **Model** | **llama.cpp** | **exo** | **dllama** | **prima.cpp** | |-----------------------------------|---------------|---------------|------------|---------------| -| Qwen-2.5-7B | **14 ms** | 86 ms | - | 44 ms | -| DeepSeek-R1-Distill-Qwen-7B | **14 ms** | 68 ms | - | 52 ms | -| DeepSeek-R1-Distill-Llama-8B | **14 ms** | 77 ms | 435 ms | 59 ms | -| Qwen-2.5-14B | **23 ms** | 31710 ms | - | 65 ms | -| DeepSeek-R1-Distill-Qwen-14B | **24 ms** | 23475 ms | - | 76 ms | +| Qwen-2.5-7B | 14 ms | 86 ms | - | **14 ms** | +| DeepSeek-R1-Distill-Qwen-7B | 14 ms | 68 ms | - | **14 ms** | +| DeepSeek-R1-Distill-Llama-8B | 14 ms | 77 ms | 435 ms | **14 ms** | +| Qwen-2.5-14B | 23 ms | 31710 ms | - | **23 ms** | +| DeepSeek-R1-Distill-Qwen-14B | 24 ms | 23475 ms | - | **24 ms** | | Qwen-2.5-32B and QwQ-32B | 224 ms | OOM | - | **89 ms** | | DeepSeek-R1-Distill-Qwen-32B | 232 ms | OOM | - | **93 ms** | | DeepSeek-R1-Distill-Llama-70B | 10978 ms | OOM | - | **724 ms** | @@ -61,9 +61,9 @@ And, if your devices are more powerful, you could unlock even more possibilities > As video recording consumes some RAM, prima.cpp proactively reduces memory usage, resulting in slightly higher latency in the video compared to the table. -> In the old version (w/o device selection), each device is assigned at least one model layer. This would lead to a 1:1:29:1 split for Llama 3-8B, which makes prima.cpp slower than llama.cpp. +> ~~In the old version (w/o device selection), each device is assigned at least one model layer. This would lead to a 1:1:29:1 split for Llama 3-8B, which makes prima.cpp slower than llama.cpp.~~ > -> **New:** In the latest version (with device selection), we will have a 0:0:32:0 split and weak devices removed, then prima.cpp would become llama.cpp when serving small models. +> In the current version (with device selection), we will have a 32:0:0:0 split and weak devices removed, then prima.cpp would become llama.cpp when serving small models. ## 🔑 Key Features @@ -72,8 +72,10 @@ And, if your devices are more powerful, you could unlock even more possibilities - - **GPU & CPU Offloading:** If a device has a GPU, you can use both GPU and CPU for inference. For example, when VRAM is full, we can offload some model layers to RAM. - - **Piped-ring parallelism with prefetching:** Prefetch upcoming layer weights to overlap disk loading latency and use advanced piped-ring parallelism to prevent the "prefetch-release" effect. This new parallelism improves pipeline parallelism by using a ring structure and allows devices to run multiple cycles to predict a new token. - - **Heterogeneity-aware workload distribution:** A scheduler is designed to optimize workload distribution based on each device's computing power, disk speed, memory, and OS (the OS will affect the disk speed and the memory management strategy). It decides how many model layers a device should handle and how many should run on GPU (if available). -- - **Automatic device selection:** If there are weak devices and removing them would speed up inference, prima.cpp will automatically discover and remove them. +- - **Automatic device selection:** If there are weak devices and removing them would speed up inference, prima.cpp will automatically discover and remove them. This may retain some devices as proxy to prevent the socket connection from being blocked. - - **Quantization:** We now support Q4K, Q6K, Q80 and IQ1 quantization (GGUF format) and are exploring a Q4K-IQ1 hybrid for a better balance between performance and speed. +- - **Speculative decoding:** We now support speculative decoding, which can [further speed up by up to 80%.](https://github.com/Lizonghang/prima.cpp/discussions/29) +- **Dynamic batching**: We now support concurrent requests from multiple users and batch decoding. - **Support Models:** We now support hot models like the **Llama, Qwen (and QwQ), and DeepSeek series**. More will be added in future updates. - **Cross-Platform:** The cluster can consist of devices with different OSs, including macOS, Linux, Android, HarmonyOS, etc. Now, Android and HarmonyOS devices require Termux, and Windows support will be added in future update. @@ -120,7 +122,8 @@ Before using this project, ensure you have the following dependencies installed: **Linux (e.g., Ubuntu):** ```shell -sudo apt update -y && sudo apt install -y gcc-9 make cmake fio git wget libzmq3-dev +# Use apt in Linux and pkg in Termux +sudo apt update -y && sudo apt install -y gcc-9 make cmake fio git wget libzmq3-dev curl ``` For HiGHS, download and install from [source](https://github.com/ERGO-Code/HiGHS): @@ -132,12 +135,13 @@ mkdir build && cd build cmake .. make -j$(nproc) sudo make install +sudo ldconfig ``` **macOS:** ```shell -brew install gcc make cmake fio git wget highs zeromq +brew install gcc make cmake fio git wget highs zeromq curl ``` ### Build, Download, and Test @@ -201,7 +205,7 @@ graph LR; > **NOTE:** This ring communication is a communication overlay, not the physical topology. These devices are physically fully connected because they all connect to the same Wi-Fi. -> If possible, disable the firewall to prevent the ports needed (e.g., 9000, 10000) been blocked. +> If possible, disable the firewall to prevent the ports needed (e.g., 9000, 10000) been blocked, or you can use `--data-port` (9000, by default) and `--signal-port` (10000, by default) to customize the ports used. Take QwQ-32B as an example, run the following commands on the devices to launch distributed inference: @@ -222,6 +226,8 @@ Take QwQ-32B as an example, run the following commands on the devices to launch Once started, prima.cpp will profile each device and decide how much workload to assign, e.g., how many model layers each device should handle, and how many of them should run on GPU (if available). +> By default, the output layer runs on the CPU. However, if you have enough total VRAM, add `--keep-out-in-cuda` to the master to run it on the GPU. + ### Effecient Multi-Split GGUF Support in `prima.cpp` If your model is split into multiple GGUF files using `gguf-split`, for example by running: @@ -266,17 +272,15 @@ If a required tensor is missing from the specified splits, the loader will raise ### (Optional) Run with Prebuilt Docker Image Assume we have a host machine with at least 32 CPU cores, 32 GiB RAM, and 32 GiB VRAM. We simulate 4 homogeneous nodes using Docker containers, with each node allocated 8 CPU cores, 8 GiB RAM, and 8 GiB VRAM. Follow the below steps to get started: -1. Pull our prebuilt Docker image (e.g., [`prima.cpp:1.0.1-cuda`](https://hub.docker.com/repository/docker/lizonghango00o1/prima.cpp/general)) and run 4 containers: +1. Pull our prebuilt Docker image (e.g., [`prima.cpp:1.0.2-cuda`](https://hub.docker.com/repository/docker/lizonghango00o1/prima.cpp/general)) and run 4 containers: ```shell -sudo docker run -dit --name prima-v1 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="0-7" --network host --gpus all prima.cpp:1.0.1-cuda -sudo docker run -dit --name prima-v2 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="8-15" --network host --gpus all prima.cpp:1.0.1-cuda -sudo docker run -dit --name prima-v3 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="16-23" --network host --gpus all prima.cpp:1.0.1-cuda -sudo docker run -dit --name prima-v4 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="24-31" --network host --gpus all prima.cpp:1.0.1-cuda +sudo docker run -dit --name prima-v1 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="0-7" --network host --gpus all prima.cpp:1.0.2-cuda +sudo docker run -dit --name prima-v2 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="8-15" --network host --gpus all prima.cpp:1.0.2-cuda +sudo docker run -dit --name prima-v3 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="16-23" --network host --gpus all prima.cpp:1.0.2-cuda +sudo docker run -dit --name prima-v4 --memory=8gb --memory-swap=8gb --cpus 8 --cpuset-cpus="24-31" --network host --gpus all prima.cpp:1.0.2-cuda ``` -> If your host machine does not have a GPU, ignore the `--gpus all` option. - 2. Download the model file [`qwq-32b-q4_k_m.gguf`](https://huggingface.co/Qwen/QwQ-32B-GGUF) and copy it into each container: ```shell @@ -287,27 +291,27 @@ sudo docker cp qwq-32b-q4_k_m.gguf prima-v3:/root/prima.cpp/download/ sudo docker cp qwq-32b-q4_k_m.gguf prima-v4:/root/prima.cpp/download/ ``` -3. (Optional) Enter each container, rebuild prima.cpp if your host machine does not have a GPU: +3. Enter each container and build prima.cpp: ```shell -cd /root/prima.cpp && make clean -make -j$(nproc) # If not rank 0 -make USE_HIGHS=1 -j$(nproc) # If rank 0 +cd /root/prima.cpp +make GGML_CUDA=1 USE_HIGHS=1 -j$(nproc) # For rank 0 +make GGML_CUDA=1 -j$(nproc) # For other ranks ``` 4. Enter each container and launch the distributed inference: ```shell cd /root/prima.cpp -(prima-v1) ./llama-cli -m download/qwq-32b-q4_k_m.gguf -c 1024 -n 256 -p "what is edge AI?" --world 4 --rank 0 --prefetch --gpu-mem 8 -(prima-v2) ./llama-cli -m download/qwq-32b-q4_k_m.gguf -c 1024 --world 4 --rank 1 --prefetch --gpu-mem 8 -(prima-v3) ./llama-cli -m download/qwq-32b-q4_k_m.gguf -c 1024 --world 4 --rank 2 --prefetch --gpu-mem 8 -(prima-v4) ./llama-cli -m download/qwq-32b-q4_k_m.gguf -c 1024 --world 4 --rank 3 --prefetch --gpu-mem 8 +(prima-v1) ./llama-cli -m download/qwq-32b-q4_k_m.gguf --world 4 --rank 0 --prefetch --gpu-mem 8 -c 4096 -n 256 -p "what is edge AI?" +(prima-v2) ./llama-cli -m download/qwq-32b-q4_k_m.gguf --world 4 --rank 1 --prefetch --gpu-mem 8 +(prima-v3) ./llama-cli -m download/qwq-32b-q4_k_m.gguf --world 4 --rank 2 --prefetch --gpu-mem 8 +(prima-v4) ./llama-cli -m download/qwq-32b-q4_k_m.gguf --world 4 --rank 3 --prefetch --gpu-mem 8 ``` -> If your host machine does not have a GPU, ignore the `--gpu-mem` option. +> You can ignore `--gpu-mem` if you don't want to limit VRAM usage. -> If you update to the latest code, non-rank 0 nodes can omit `-c 1024`. +> Always use `git fetch` to update the local repository. ### Run in Server Mode You can run prima.cpp in server mode, by launching `llama-server` on the rank 0 device (with `--host` and `--port` specified) and `llama-cli` on the others. Here is an example with 2 devices: @@ -320,6 +324,8 @@ You can run prima.cpp in server mode, by launching `llama-server` on the rank 0 ./llama-cli -m download/qwq-32b-q4_k_m.gguf --world 2 --rank 1 --master 192.168.1.2 --next 192.168.1.2 --prefetch ``` +You can specify `-np 4 --cont-batching` when launching `llama-server` to enable concurrent requests. + After that, you can interact with the rank 0 device by calling the Chat Completion API: ```shell @@ -418,6 +424,10 @@ curl -X POST http://localhost:8080/v1/cancel \ -d '{"task_id": 0}' ``` +**9. How to use speculative decoding?** + +Please see "[Power prima.cpp with speculative decoding: Further speeds up by up to 80%](https://github.com/Lizonghang/prima.cpp/discussions/29)". + ## ❤️ Acknowledgment This project builds upon the incredible work from the open-source community, especially [ggml, gguf](https://github.com/ggml-org/ggml), and [llama.cpp](https://github.com/ggml-org/llama.cpp). We gratefully acknowledge their contributions. diff --git a/common/arg.cpp b/common/arg.cpp index 7bbdffb7..ab43ddba 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -627,12 +627,19 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, } ).set_examples({LLAMA_EXAMPLE_SPECULATIVE})); add_opt(llama_arg( - {"--draft"}, "N", - format("number of tokens to draft for speculative decoding (default: %d)", params.n_draft), + {"--draft-max", "--draft", "--draft-n"}, "N", + format("number of tokens to draft for speculative decoding (default: %d)", params.speculative.n_max), [](gpt_params & params, int value) { - params.n_draft = value; + params.speculative.n_max = value; } - ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP})); + ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER})); + add_opt(llama_arg( + {"--draft-min", "--draft-n-min"}, "N", + format("minimum number of draft tokens to use for speculative decoding (default: %d)", params.speculative.n_min), + [](gpt_params & params, int value) { + params.speculative.n_min = value; + } + ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER})); add_opt(llama_arg( {"-ps", "--p-split"}, "N", format("speculative decoding split probability (default: %.1f)", (double)params.p_split), @@ -640,6 +647,13 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, params.p_split = std::stof(value); } ).set_examples({LLAMA_EXAMPLE_SPECULATIVE})); + add_opt(llama_arg( + {"--draft-p-min"}, "P", + format("minimum speculative decoding probability (greedy) (default: %.1f)", (double)params.speculative.p_min), + [](gpt_params & params, const std::string & value) { + params.speculative.p_min = std::stof(value); + } + ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER})); add_opt(llama_arg( {"-lcs", "--lookup-cache-static"}, "FNAME", "path to static lookup cache to use for lookup decoding (not updated by generation)", @@ -659,6 +673,7 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, format("size of the prompt context (default: %d, 0 = loaded from model)", params.n_ctx), [](gpt_params & params, int value) { params.n_ctx = value; + params.speculative.n_ctx = value; } ).set_env("LLAMA_ARG_CTX_SIZE")); add_opt(llama_arg( @@ -770,6 +785,9 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, format("maximum GPU memory to use (default: %d)", params.gpu_mem), [](gpt_params & params, int value) { params.gpu_mem = value; // in GiB + if (value == 0) { + LOG_WRN("WARN: Set --gpu-mem to 0 may lead to errors during workload distribution.\n"); + } } ).set_env("LLAMA_ARG_CUDA_MEM")); add_opt(llama_arg( @@ -786,6 +804,14 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, params.force = true; } ).set_env("LLAMA_ARG_FORCE")); + add_opt(llama_arg( + {"--master-priority"}, "N", + format("priority to assign workload to the master (default: %f, set 1.01 to use master first, and 0.99 to offload to other devices)", params.master_priority), + [](gpt_params & params, const std::string & value) { + params.master_priority = std::stof(value); + } + ).set_env("LLAMA_ARG_MASTER_PRIORITY")); + // #ifdef GGML_USE_METAL // // warn: if the output layer weights are not kept in metal shared memory, its mmap-ed weight data // // could be released by the OS and reloaded repeatedly, which causes additional disk I/O latency. @@ -798,6 +824,17 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, // } // ).set_env("LLAMA_ARG_KEEP_INP_OUT_IN_METAL")); // #endif + +#ifdef GGML_USE_CUDA + add_opt(llama_arg( + {"--keep-out-in-cuda"}, + format("whether to compute the output layer on CUDA (default: %s)", params.keep_out_in_cuda ? "true" : "false"), + [](gpt_params & params) { + params.keep_out_in_cuda = true; + } + ).set_env("LLAMA_ARG_KEEP_INP_OUT_IN_CUDA")); +#endif + 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), @@ -1561,13 +1598,14 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, {"-ngld", "--gpu-layers-draft", "--n-gpu-layers-draft"}, "N", "number of layers to store in VRAM for the draft model", [](gpt_params & params, int value) { - params.n_gpu_layers_draft = value; + params.n_gpu_layers_draft = value; // TODO: remove + params.speculative.n_gpu_layers = value; if (!llama_supports_gpu_offload()) { fprintf(stderr, "warning: not compiled with GPU offload support, --gpu-layers-draft option will be ignored\n"); fprintf(stderr, "warning: see main README.md for information on enabling GPU BLAS support\n"); } } - ).set_examples({LLAMA_EXAMPLE_SPECULATIVE})); + ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER})); add_opt(llama_arg( {"-sm", "--split-mode"}, "{none,layer,row}", "how to split the model across multiple GPUs, one of:\n" @@ -1710,9 +1748,9 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, {"-md", "--model-draft"}, "FNAME", "draft model for speculative decoding (default: unused)", [](gpt_params & params, const std::string & value) { - params.model_draft = value; + params.speculative.model = value; } - ).set_examples({LLAMA_EXAMPLE_SPECULATIVE})); + ).set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER})); add_opt(llama_arg( {"-mu", "--model-url"}, "MODEL_URL", "model download url (default: unused)", @@ -2027,6 +2065,13 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, params.chat_template = value; } ).set_examples({LLAMA_EXAMPLE_MAIN, LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CHAT_TEMPLATE")); + add_opt(llama_arg( + {"-sys", "--system-prompt"}, "PROMPT", + "system prompt to use with model (if applicable, depending on chat template)", + [](gpt_params & params, const std::string & value) { + params.system_prompt = value; + } + ).set_examples({LLAMA_EXAMPLE_MAIN, LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_SYSTEM_PROMPT")); add_opt(llama_arg( {"-sps", "--slot-prompt-similarity"}, "SIMILARITY", format("how much the prompt of a request must match the prompt of a slot in order to use that slot (default: %.2f, 0.0 = disabled)\n", params.slot_prompt_similarity), diff --git a/common/common.cpp b/common/common.cpp index 2e56274f..4d216c1e 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -28,6 +28,7 @@ #include #include #include +#include #if defined(__APPLE__) && defined(__MACH__) #include @@ -846,6 +847,16 @@ static std::string vec_to_str(const std::vector & vec) { return oss.str(); } +static backend_type get_backend_type(const gpu_support & support) { + if (support.cuda) return BACKEND_CUDA; + if (support.metal) return BACKEND_METAL; + if (support.vulkan) return BACKEND_VULKAN; + if (support.kompute) return BACKEND_KOMPUTE; + if (support.gpublas) return BACKEND_GPUBLAS; + if (support.sycl) return BACKEND_SYCL; + return BACKEND_CPU; +} + static bool assign_layers_to_device( uint32_t n_world, const device_info * dev_info_set, @@ -971,11 +982,11 @@ static bool assign_layers_to_device( bool is_android = strcmp(dev.device_os, "Android") == 0; bool is_windows = strcmp(dev.device_os, "Windows") == 0; GGML_ASSERT(!is_windows && "Windows is not tested yet\n"); - + if ((is_macos && !dev.gpu_support.metal) || is_linux) { mem_budget[m] = dev.memory.available_physical; } else if (is_macos && dev.gpu_support.metal) { - mem_budget[m] = dev.gpu_props.memory_free; + mem_budget[m] = dev.gpu_props.memory_free + 1e-4; // to avoid division by zero } else if (is_android) { mem_budget[m] = dev.memory.available_physical + dev.memory.used_can_swap; } else { @@ -984,17 +995,35 @@ static bool assign_layers_to_device( } } - // initialize w_m proportionally to memory budget and n_m to 0 + // initialize w_m proportionally to memory budget float total_mem_budget = std::accumulate(mem_budget.begin(), mem_budget.end(), 0.0f); for (uint32_t m = 0; m < n_world; ++m) { w[m] = std::round(mem_budget[m] / total_mem_budget * n_layer); - n[m] = 0; + } + // no 0 is allowed in w, it must be at least 1 + for (uint32_t m = 0; m < n_world; ++m) { + if (w[m] == 0) { + w[m] = 1; + // find the maximum and decrease it by 1 + auto max_it = std::max_element(w.begin(), w.end()); + if (max_it != w.end() && *max_it > 1) { + *max_it -= 1; + } + } } // adjust w[m] to ensure L mod W = 0 int diff = n_layer - std::accumulate(w.begin(), w.end(), 0); auto device = (diff > 0) ? std::max_element(mem_budget.begin(), mem_budget.end()) : std::min_element(mem_budget.begin(), mem_budget.end()); w[std::distance(mem_budget.begin(), device)] += diff; + // initialize n_m to w_m (if there is GPU), assume all layers can run on GPU + for (uint32_t m = 0; m < n_world; ++m) { + if (dev_info_set[m].gpu_support.metal || dev_info_set[m].gpu_support.cuda) { + n[m] = w[m]; + } else { + n[m] = 0; + } + } // stores the actual read bandwidth (GB/s) for each device std::vector disk_speed(n_world, 0.0f); @@ -1051,8 +1080,7 @@ static bool assign_layers_to_device( bool is_windows = strcmp(dev.device_os, "Windows") == 0; GGML_ASSERT(!is_windows && "Windows is not tested yet\n"); - bool use_gpu = dev.gpu_support.metal || dev.gpu_support.cuda; - llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, use_gpu, m == 0, w[m] * k, n[m] * k); + llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, get_backend_type(dev.gpu_support), m, dev_info_set[0].model_bytes, w[m] > n[m], n[m] > 0); int l_m = w[m] * k; // total number of layers assigned to device m int l_m_gpu = n[m] * k; // number of layers assigned to device m that run on GPU @@ -1213,6 +1241,7 @@ static bool assign_layers_to_device( if (dev.gpu_support.metal && m == 0 && cparams.keep_out_in_metal) { vec_z_gpu[m] -= (double)bo / (double)(n_layer * b_prime); } + vec_z_gpu[m] = std::max(vec_z_gpu[m], 0.0f); } } @@ -1247,6 +1276,8 @@ static bool assign_layers_to_device( return cost * k; } ); + // apply priority to the head device + model.lp_.col_cost_[0] *= 1.0 / cparams.master_priority; // define the variable bounds model.lp_.col_lower_ = std::vector(n_world * 2, 0.0); @@ -1412,26 +1443,37 @@ static bool assign_layers_to_device( } // check the solution - bool has_free_gpu_memory = false, has_gpu_overload = false, has_cpu_overload = false; + bool has_free_gpu_memory = false, has_gpu_overload = false, has_cpu_overload = false, has_weak_device = false; for (uint32_t m = 0; m < n_world; ++m) { // if (!dev_gpu[m]) continue; uint32_t w_m = best_solution[m], n_m = best_solution[m + n_world]; + if (w_m == 1 && n_m == 0) { + // if the device is weak + has_weak_device = true; + LOG_INF("Device %d is weak, need to be removed: w_m = %d, n_m = %d\n", m, w_m, n_m); + } + if (dev_gpu[m]) { if (n_m < static_cast(std::floor(W * vec_z_gpu[m]))) { // if there is still free GPU memory has_free_gpu_memory = true; - } else if (w_m > n_m) { - // if the GPU is overloaded + LOG_INF("Device %d still has free GPU memory: w_m = %d, n_m = %d, W * vec_z_gpu[m]) = %d\n", + m, w_m, n_m, static_cast(std::floor(W * vec_z_gpu[m]))); + } + if (w_m > n_m) { + // if layers are offloaded to CPU has_gpu_overload = true; + LOG_INF("Device %d has GPU overload: w_m = %d, n_m = %d\n", m, w_m, n_m); } } else if (!in_set(m, M4)) { // if the CPU is overloaded has_cpu_overload = true; + LOG_INF("Device %d has CPU overload.\n", m); } } - if (has_free_gpu_memory && (has_gpu_overload || has_cpu_overload)) { + if (!has_weak_device && has_free_gpu_memory && (has_gpu_overload || has_cpu_overload)) { int worst_device = -1; float worst_speed = std::numeric_limits::max(); @@ -1503,13 +1545,24 @@ static bool assign_layers_to_device( float total_mem_budget = std::accumulate(mem_budget.begin(), mem_budget.end(), 0.0f); for (uint32_t m = 0; m < n_world; ++m) { w[m] = std::round(mem_budget[m] / total_mem_budget * n_layer); - n[m] = 0; + } + // no 0 is allowed in w, it must be at least 1 + for (uint32_t m = 0; m < n_world; ++m) { + if (w[m] == 0) { + w[m] = 1; + // find the maximum and decrease it by 1 + auto max_it = std::max_element(w.begin(), w.end()); + if (max_it != w.end() && *max_it > 1) { + *max_it -= 1; + } + } } // adjust w[m] to ensure L mod W = 0 int diff = n_layer - std::accumulate(w.begin(), w.end(), 0); auto device = (diff > 0) ? std::max_element(mem_budget.begin(), mem_budget.end()) : std::min_element(mem_budget.begin(), mem_budget.end()); w[std::distance(mem_budget.begin(), device)] += diff; + std::copy(w.begin(), w.end(), n_layer_window); std::vector vec_z_gpu(n_world, 0.0f); @@ -1517,8 +1570,7 @@ static bool assign_layers_to_device( for (uint32_t m = 0; m < n_world; ++m) { const device_info & dev = dev_info_set[m]; - bool use_gpu = dev.gpu_support.metal || dev.gpu_support.cuda; - llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, use_gpu, m == 0, w[m], n[m]); + llama_model_compute_buf_size(&c_cpu[m], &c_gpu[m], model, cparams, get_backend_type(dev.gpu_support), m, dev_info_set[0].model_bytes, false, true); if (dev.gpu_support.cuda || dev.gpu_support.metal) { int64_t required_mem = w[m] * b_prime; @@ -1676,6 +1728,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) { cparams.n_layer_window[0] = n_layers; mparams.n_layer_window[0] = n_layers; llama_context_n_layer_window(lctx)[0] = n_layers; + llama_update_context_with_rankworld(lctx, 0, 1, 0, 1); #if defined(GGML_USE_METAL) || defined(GGML_USE_CUDA) params.n_gpu_layers = std::min((int32_t)n_layers, params.n_gpu_layers); @@ -1717,13 +1770,15 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) { } // sychronize device profile to the master node + NodeType node_type = NodeType::NODE_TYPE_WORKER; + char is_forwarder[32] = {0}; if (my_rank == 0) { if (auto_schedule) { std::vector dev_info_set(n_world); dev_info_set[0] = dev_info; llama_gather_device_info(lctx, dev_info_set.data()); - device_print_props(dev_info_set.data(), n_world, model, cparams); + device_print_props (dev_info_set.data(), n_world, model, cparams); // assign layers to devices and remove weak devices if (!assign_layers_and_select_devices(n_world, dev_info_set, n_layer_window, n_gpu_layers, model, cparams)) { @@ -1733,7 +1788,7 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) { return iparams; } llama_bcast_layer_setup(lctx, n_layer_window, n_gpu_layers); - llama_rebuild_topo(lctx, n_layer_window, dev_info_set.data()); + llama_rebuild_topo (lctx, n_layer_window, dev_info_set.data(), &node_type, is_forwarder); } else { // use the user-defined n_layer_window std::copy(std::begin(params.n_layer_window), std::end(params.n_layer_window), n_layer_window); @@ -1741,16 +1796,16 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) { } } else { if (auto_schedule){ - llama_send_device_info(lctx, &dev_info); - llama_recv_layer_setup(lctx, n_layer_window, n_gpu_layers); - llama_rebuild_topo (lctx, n_layer_window, nullptr); + llama_send_device_info (lctx, &dev_info); + llama_recv_layer_setup (lctx, n_layer_window, n_gpu_layers); + llama_rebuild_topo (lctx, n_layer_window, nullptr, &node_type, is_forwarder); } else { - llama_recv_layer_setup(lctx, n_layer_window, n_gpu_layers); + llama_recv_layer_setup (lctx, n_layer_window, n_gpu_layers); } } // if this is a weak device, then exit - if (n_layer_window[my_rank] <= 0) { + if (node_type == NodeType::NODE_TYPE_EXIT) { LOG_INF("No layer is assigned to me, exit.\n"); llama_free(lctx); llama_free_model(model); @@ -1759,18 +1814,22 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) { // update my rank and n_world uint32_t update_rank = 0, update_n_world = 1; + uint32_t worker_rank = 0, n_worker = 1; std::vector n_layer_window_temp = {n_layer_window[0]}, n_gpu_layers_temp = {n_gpu_layers[0]}; for (uint32_t i = 1; i < n_world; i++) { - if (n_layer_window[i] <= 0) { + if (n_layer_window[i] <= 0 && is_forwarder[i] == 0) { continue; } - if (i <= my_rank) { - update_rank++; - } + if (i <= my_rank) update_rank++; update_n_world++; n_layer_window_temp.push_back(n_layer_window[i]); n_gpu_layers_temp.push_back(n_gpu_layers[i]); + + if (n_layer_window[i] > 0) { + if (i <= my_rank) worker_rank++; + n_worker++; + } } memset(n_layer_window, 0, n_world * sizeof(uint32_t)); @@ -1793,8 +1852,26 @@ struct llama_init_result llama_init_from_gpt_params(gpt_params & params) { params.n_world = update_n_world; n_world = update_n_world; - llama_update_context_with_rankworld(lctx, update_rank, update_n_world); - + llama_update_context_with_rankworld(lctx, update_rank, update_n_world, worker_rank, n_worker); + + if (node_type == NodeType::NODE_TYPE_FORWARDER) { + //just forward + LOG_INF("No layer is assigned to me, and I serve as a network proxy.\n"); + std::atomic should_exit{false}; + auto t = std::thread([lctx, &should_exit]() { + while(!should_exit) { + llama_forward_messages(lctx); + } + }); + char * stop_signal = nullptr; + llama_free_sockets(lctx, &stop_signal); // this will block until receive stop signal + + should_exit = true; + t.join(); + + exit(0); + } + // update n_layer_window and n_gpu_layers std::copy(std::begin(n_layer_window), std::end(n_layer_window), params.n_layer_window); std::copy(std::begin(n_layer_window), std::end(n_layer_window), cparams.n_layer_window); @@ -1940,17 +2017,20 @@ 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.rpc_servers = params.rpc_servers.c_str(); + + mparams.n_world = params.n_world; + mparams.rank = params.rank; + mparams.rpc_servers = params.rpc_servers.c_str(); mparams.gguf_splits = params.gguf_splits.c_str(); - mparams.main_gpu = params.main_gpu; - mparams.split_mode = params.split_mode; - mparams.tensor_split = params.tensor_split; - mparams.use_mmap = params.use_mmap; - mparams.use_mlock = params.use_mlock; - mparams.check_tensors = params.check_tensors; + mparams.main_gpu = params.main_gpu; + mparams.split_mode = params.split_mode; + mparams.tensor_split = params.tensor_split; + mparams.use_mmap = params.use_mmap; + mparams.use_mlock = params.use_mlock; + mparams.check_tensors = params.check_tensors; mparams.keep_out_in_metal = params.keep_out_in_metal; + mparams.keep_out_in_cuda = params.keep_out_in_cuda; + std::copy(std::begin(params.n_layer_window), std::end(params.n_layer_window), mparams.n_layer_window); if (params.kv_overrides.empty()) { mparams.kv_overrides = NULL; @@ -1990,7 +2070,9 @@ struct llama_context_params llama_context_params_from_gpt_params(const gpt_param cparams.rank = params.rank; cparams.prefetch = params.prefetch; cparams.force = params.force; + cparams.master_priority = params.master_priority; cparams.keep_out_in_metal = params.keep_out_in_metal; + cparams.keep_out_in_cuda = params.keep_out_in_cuda; cparams.n_gpu_layers = params.n_gpu_layers; cparams.n_cycles = params.n_cycles; std::copy(std::begin(params.n_layer_window), std::end(params.n_layer_window), cparams.n_layer_window); @@ -3044,7 +3126,7 @@ void yaml_dump_non_result_info(FILE * stream, const gpt_params & params, const l fprintf(stream, "mirostat_lr: %f # default: 0.1\n", sparams.mirostat_eta); fprintf(stream, "mlock: %s # default: false\n", params.use_mlock ? "true" : "false"); fprintf(stream, "model: %s # default: %s\n", params.model.c_str(), DEFAULT_MODEL_PATH); - fprintf(stream, "model_draft: %s # default:\n", params.model_draft.c_str()); + fprintf(stream, "model_draft: %s # default:\n", params.speculative.model.c_str()); fprintf(stream, "multiline_input: %s # default: false\n", params.multiline_input ? "true" : "false"); fprintf(stream, "n_gpu_layers: %d # default: -1\n", params.n_gpu_layers); fprintf(stream, "n_predict: %d # default: -1 (unlimited)\n", params.n_predict); diff --git a/common/common.h b/common/common.h index 546bf9c2..6eaeffe0 100644 --- a/common/common.h +++ b/common/common.h @@ -33,6 +33,8 @@ struct llama_lora_adapter_container : llama_lora_adapter_info { struct llama_lora_adapter * adapter; }; +using llama_tokens = std::vector; + // build info extern int LLAMA_BUILD_NUMBER; extern char const * LLAMA_COMMIT; @@ -141,12 +143,26 @@ struct gpt_sampler_params { std::string print() const; }; +struct common_params_speculative { + int32_t n_ctx = 0; // draft context size + int32_t n_max = 16; // maximum number of tokens to draft during speculative decoding + int32_t n_min = 5; // minimum number of draft tokens to use for speculative decoding + int32_t n_gpu_layers = -1; // number of layers to store in VRAM for the draft model (-1 - use default) + float p_split = 0.1f; // speculative decoding split probability + float p_min = 0.9f; // minimum speculative decoding probability (greedy) + + struct cpu_params cpuparams; + struct cpu_params cpuparams_batch; + + std::string model = ""; // draft model for speculative decoding // NOLINT +}; + struct gpt_params { int32_t n_world = 1; // number of devices to use int32_t rank = 0; // my rank for distributed inference uint32_t n_layer_window[32] = {0}; // layer window size on each node - std::string master_ip = "localhost"; // ip address of the master node - std::string next_node_ip = "localhost"; // ip address of my next node + std::string master_ip = "127.0.0.1"; // ip address of the master node + std::string next_node_ip = "127.0.0.1"; // ip address of my next node uint32_t data_port = 9000; // base data port for this node uint32_t signal_port = 10000; // base signal port for this node uint32_t master_data_port = 9000; // data port base for master node @@ -154,7 +170,9 @@ struct gpt_params { uint32_t next_node_signal_port = 10000; // signal port base for next node bool prefetch = false; // prefetch layer weights bool keep_out_in_metal = true; // whether to keep output weights in metal memory, true by default + bool keep_out_in_cuda = false; // whether to run the output layer on CUDA, false by default bool force = false; // force to start prefetching after computation + float master_priority = 1.01; // priority to assign workload to the master (set 1.01 to use master first, and 0.99 to offload to other devices) int32_t gpu_mem = 999.0; // gpu memory to use, in GiB int32_t n_cycles = 0; // number of cycles to output one token int32_t n_predict = -1; // new tokens to predict @@ -162,7 +180,6 @@ struct gpt_params { int32_t n_batch = 2048; // logical batch size for prompt processing (must be >=32 to use BLAS) int32_t n_ubatch = 512; // physical batch size for prompt processing (must be >=32 to use BLAS) int32_t n_keep = 0; // number of tokens to keep from initial prompt - int32_t n_draft = 5; // number of tokens to draft during speculative decoding int32_t n_chunks = -1; // max number of chunks to process (-1 = unlimited) int32_t n_parallel = 1; // number of parallel sequences to decode int32_t n_sequences = 1; // number of sequences to decode @@ -199,9 +216,9 @@ struct gpt_params { enum llama_attention_type attention_type = LLAMA_ATTENTION_TYPE_UNSPECIFIED; // attention type for embeddings struct gpt_sampler_params sparams; + struct common_params_speculative speculative; std::string model = ""; // model path // NOLINT - std::string model_draft = ""; // draft model for speculative decoding // NOLINT std::string model_alias = "unknown"; // model alias // NOLINT std::string model_url = ""; // model url to download // NOLINT std::string hf_token = ""; // HF token // NOLINT diff --git a/common/profiler.cpp b/common/profiler.cpp index 0da13824..e9862bfb 100644 --- a/common/profiler.cpp +++ b/common/profiler.cpp @@ -924,7 +924,7 @@ static void check_env_path() { setenv("PATH", update_env_path.c_str(), 1); } -static void external_fio_impl(float * read_bw, float * write_bw, bool op_rand, int n_threads) { +static void external_fio_impl(float * read_bw, float * write_bw, bool op_rand, int n_threads) { pid_t pid = getpid(); // avoid conflict with other processes std::string test_file = "fio_test_" + std::to_string(pid); @@ -1603,10 +1603,20 @@ static float device_disk_access_delay(struct device_info & dev_info, struct llam #if defined(GGML_USE_METAL) || defined(GGML_USE_CUDA) llama_kv_size(&cpu_kv_size, &gpu_kv_size, model, cparams, true); - llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, true, true, n_layers, n_gpu_layers); + + enum backend_type backend; +#if GGML_USE_METAL + backend = BACKEND_METAL; +#elif GGML_USE_CUDA + backend = BACKEND_CUDA; +#endif + llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, backend, 0, n_bytes, n_layers > n_gpu_layers, n_gpu_layers > 0); + #else llama_kv_size(&cpu_kv_size, &gpu_kv_size, model, cparams, false); - llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, false, true, n_layers, n_gpu_layers); + + enum backend_type backend = BACKEND_CPU; + llama_model_compute_buf_size(&cpu_compute_buf, &gpu_compute_buf, model, cparams, backend, 0, n_bytes, n_layers > n_gpu_layers, n_gpu_layers > 0); #endif double cpu_kv_size_gib = static_cast(cpu_kv_size) / 1024.0 / 1024.0 / 1024.0; // convert to GiB @@ -2621,7 +2631,7 @@ size_t serialize(const struct device_info * dev_info, char ** buffer) { return total_size; } -void deserialize(const char * buffer, struct device_info * dev_info) { +size_t deserialize(const char * buffer, struct device_info * dev_info) { const char * ptr = buffer; // rank @@ -2821,6 +2831,34 @@ void deserialize(const char * buffer, struct device_info * dev_info) { ptr += sizeof(float); memcpy(&dev_info->gpu_props.cuda_mem_cpy_delay, ptr, sizeof(float)); + ptr += sizeof(float); // no need to synchronize model flops and model params -} \ No newline at end of file + return ptr - buffer; +} + +void TopoRebuildHelperInfo::deserialize(const char * buffer) { + size_t buffer_size = ::deserialize(buffer, &dev_info); + if (buffer_size == 0) { + LOG_ERR("%s: failed to deserialize device info\n", __func__); + return; + } + memcpy(&is_forwarder, buffer + buffer_size, 1); +} + +size_t TopoRebuildHelperInfo::serialize(char ** buffer) const{ + size_t buffer_size = ::serialize(&dev_info, buffer); + char * buffer_ = (char *)malloc(buffer_size + 1); + + if (buffer_ == NULL) { + LOG_ERR("%s: failed to allocate %zu bytes for device info serialization\n", + __func__, buffer_size); + return 0; + } + + memcpy(buffer_, *buffer, buffer_size); + memcpy(buffer_ + buffer_size, &is_forwarder, 1); + free(*buffer); + *buffer = buffer_; + return buffer_size + 1; +} diff --git a/common/profiler.h b/common/profiler.h index 06741d6c..00407c96 100644 --- a/common/profiler.h +++ b/common/profiler.h @@ -293,10 +293,24 @@ struct model_bytes { int64_t nb_layer; int64_t nb_output; + // used to estimate the compute buffer size + int64_t nb_output_w; + int64_t nb_output_norm_w; + int64_t nb_attn_norm_w; + int64_t nb_attn_q_w; + int64_t nb_ffn_gate_w; + int64_t nb_ffn_down_w; + model_bytes() : - nb_input (0), - nb_layer (0), - nb_output(0) {} + nb_input (0), + nb_layer (0), + nb_output (0), + nb_output_w (0), + nb_output_norm_w(0), + nb_attn_norm_w (0), + nb_attn_q_w (0), + nb_ffn_gate_w (0), + nb_ffn_down_w (0) {} }; struct disk_props { @@ -346,6 +360,18 @@ struct device_info { model_bytes() {} }; +struct TopoRebuildHelperInfo{ + struct device_info dev_info; + char is_forwarder; + + TopoRebuildHelperInfo(): + dev_info(), + is_forwarder(0){} + + void deserialize(const char * buffer); + size_t serialize(char ** buffer) const; +}; + enum profiler_backend_type { PROFILER_BACKEND_TYPE_CPU = 0, PROFILER_BACKEND_TYPE_METAL = 1, @@ -389,6 +415,6 @@ int device_has_blas (void); int device_has_sycl (void); size_t serialize (const struct device_info * dev_info, char ** buffer); -void deserialize(const char * buffer, struct device_info * dev_info); +size_t deserialize(const char * buffer, struct device_info * dev_info); #endif // PROFILER_H diff --git a/common/sampling.cpp b/common/sampling.cpp index 3dc7f112..b8db920a 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -318,6 +318,45 @@ llama_token gpt_sampler_sample(struct gpt_sampler * gsmpl, struct llama_context return cur_p.data[cur_p.selected].id; } +std::vector gpt_sampler_sample_and_accept_n(struct gpt_sampler * gsmpl, struct llama_context * ctx, const std::vector & idxs, const llama_tokens & draft, bool grammar_first) { + GGML_ASSERT(idxs.size() == draft.size() + 1 && "idxs.size() must be draft.size() + 1"); + + std::vector result; + result.reserve(idxs.size()); + + size_t i = 0; + for (; i < draft.size(); i++) { + const llama_token id = gpt_sampler_sample(gsmpl, ctx, idxs[i], grammar_first); + + gpt_sampler_accept(gsmpl, id, true); + + result.push_back(id); + + if (draft[i] != id) { + break; + } + } + + if (i == draft.size()) { + const llama_token id = gpt_sampler_sample(gsmpl, ctx, idxs[i], grammar_first); + + gpt_sampler_accept(gsmpl, id, true); + + result.push_back(id); + } + + return result; +} + +std::vector gpt_sampler_sample_and_accept_n(struct gpt_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first) { + std::vector idxs(draft.size() + 1); + for (size_t i = 0; i < idxs.size(); ++i) { + idxs[i] = i; + } + + return gpt_sampler_sample_and_accept_n(gsmpl, ctx, idxs, draft, grammar_first); +} + uint32_t gpt_sampler_get_seed(const struct gpt_sampler * gsmpl) { return llama_sampler_get_seed(gsmpl->chain); } diff --git a/common/sampling.h b/common/sampling.h index d0e1a920..82d6b023 100644 --- a/common/sampling.h +++ b/common/sampling.h @@ -60,6 +60,27 @@ void gpt_perf_print(const struct llama_context * ctx, const struct gpt_sampler * // llama_token gpt_sampler_sample(struct gpt_sampler * gsmpl, struct llama_context * ctx, int idx, bool grammar_first = false); +// generalized version of gpt_sampler_sample +// +// will cross-reference the sampled tokens with a batch of draft tokens and accept those that match +// if the sampler disagrees at some point, we stop and return the accepted tokens up to now +// +// gpt_sampler_sample_n(gsmpl, ctx, { idx }, {}); +// +// is equivalent to +// +// gpt_sampler_sample(gsmpl, ctx, idx); +// gpt_sampler_accept(gsmpl, token, true); +// +// requires: idxs.size() == draft.size() + 1 +// +// returns at least 1 token, up to idxs.size() +// +std::vector gpt_sampler_sample_and_accept_n(struct gpt_sampler * gsmpl, struct llama_context * ctx, const std::vector & idxs, const llama_tokens & draft, bool grammar_first = false); + +// assume idxs == [ 0, 1, 2, ..., draft.size() ] +std::vector gpt_sampler_sample_and_accept_n(struct gpt_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first = false); + uint32_t gpt_sampler_get_seed(const struct gpt_sampler * gsmpl); // helpers diff --git a/common/speculative.cpp b/common/speculative.cpp new file mode 100644 index 00000000..8a03d08c --- /dev/null +++ b/common/speculative.cpp @@ -0,0 +1,271 @@ +#include "speculative.h" + +#include "log.h" +#include "common.h" +#include "sampling.h" + +#include + +#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 128 +#define SPEC_VOCAB_CHECK_START_TOKEN_ID 5 + +struct common_speculative { + struct llama_context * ctx; + struct gpt_sampler * smpl; + + llama_batch batch; + llama_tokens prompt; +}; + +struct common_speculative * common_speculative_init( + struct llama_context * ctx_dft) { + auto * result = new common_speculative { + /* .ctx = */ ctx_dft, + /* .smpl = */ nullptr, + /* .batch = */ llama_batch_init(llama_n_batch(ctx_dft), 0, 1), + /* .prompt = */ {}, + }; + + // TODO: optimize or pass from outside? +#if 0 + { + common_params_sampling params; + params.no_perf = false; + + params.top_k = 40; + params.top_p = 0.9; + + params.samplers = { + COMMON_SAMPLER_TYPE_TOP_K, + COMMON_SAMPLER_TYPE_TOP_P, + COMMON_SAMPLER_TYPE_INFILL, + }; + + result->smpl = gpt_sampler_init(llama_get_model(ctx_dft), params); + } +#else + { + gpt_sampler_params params; + params.no_perf = false; + + params.top_k = 10; + + params.samplers = { + GPT_SAMPLER_TYPE_TOP_K, + }; + + result->smpl = gpt_sampler_init(llama_get_model(ctx_dft), params); + } +#endif + + llama_update_context_with_rankworld(result->ctx, 0, 1, 0, 1); + + return result; +} + +void common_speculative_free(struct common_speculative * spec) { + gpt_sampler_free(spec->smpl); + + llama_batch_free(spec->batch); + + delete spec; +} + +bool common_speculative_are_compatible( + const struct llama_context * ctx_tgt, + const struct llama_context * ctx_dft) { + const struct llama_model * model_tgt = llama_get_model(ctx_tgt); + const struct llama_model * model_dft = llama_get_model(ctx_dft); + + const bool vocab_type_tgt = llama_vocab_type(model_tgt); + LOG_DBG("%s: vocab_type tgt: %d\n", __func__, vocab_type_tgt); + + const bool vocab_type_dft = llama_vocab_type(model_dft); + LOG_DBG("%s: vocab_type dft: %d\n", __func__, vocab_type_dft); + + if (vocab_type_tgt != vocab_type_dft) { + LOG_ERR("%s: draft model vocab type must match target model to use speculation but " + "vocab_type_dft = %d while vocab_type_tgt = %d\n", __func__, vocab_type_dft, vocab_type_tgt); + return false; + } + + if (llama_add_bos_token(model_tgt) != llama_add_bos_token(model_dft) || + llama_add_eos_token(model_tgt) != llama_add_eos_token(model_dft) || + llama_token_bos(model_tgt) != llama_token_bos(model_dft) || + llama_token_eos(model_tgt) != llama_token_eos(model_dft) + ) { + LOG_ERR("%s: draft model special tokens must match target model to use speculation\n", __func__); + return false; + } + + { + const int n_vocab_tgt = llama_n_vocab(model_tgt); + const int n_vocab_dft = llama_n_vocab(model_dft); + + const int vocab_diff = std::abs(n_vocab_tgt - n_vocab_dft); + + if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) { + LOG_ERR("%s: draft model vocab must closely match target model to use speculation but " + "target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n", + __func__, n_vocab_tgt, llama_n_vocab(model_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE); + return false; + } + + for (int i = SPEC_VOCAB_CHECK_START_TOKEN_ID; i < std::min(n_vocab_tgt, n_vocab_dft); ++i) { + const char * token_text_tgt = llama_token_get_text(model_tgt, i); + const char * token_text_dft = llama_token_get_text(model_dft, i); + if (std::strcmp(token_text_tgt, token_text_dft) != 0) { + LOG_ERR("%s: draft model vocab must match target model to use speculation but " + "token %d content differs - target '%s', draft '%s'\n", __func__, i, + llama_token_to_piece(ctx_tgt, i).c_str(), + llama_token_to_piece(ctx_dft, i).c_str()); + return false; + } + } + } + + return true; +} + +llama_tokens common_speculative_gen_draft( + struct common_speculative * spec, + struct common_speculative_params params, + const llama_tokens & prompt_tgt, + llama_token id_last) { + auto & batch = spec->batch; + auto & ctx = spec->ctx; + auto & smpl = spec->smpl; + auto & prompt = spec->prompt; + + int reuse_i = 0; + int reuse_n = 0; + + const int n_ctx = llama_n_ctx(ctx) - params.n_draft; + + const int i_start = std::max(0, (int) prompt_tgt.size() - n_ctx); + + // reuse as much as possible from the old draft context + // ideally, the draft context should be as big as the target context and we will always reuse the entire prompt + for (int i = 0; i < (int) prompt.size(); ++i) { + int cur = 0; + while (i_start + cur < (int) prompt_tgt.size() && + i + cur < (int) prompt.size() && + prompt_tgt[i_start + cur] == prompt[i + cur]) { + cur++; + } + + if ((cur >= params.n_reuse || n_ctx >= (int) prompt_tgt.size()) && cur > reuse_n) { + reuse_i = i; + reuse_n = cur; + } + } + + LOG_DBG("%s: reuse_i = %d, reuse_n = %d, prompt = %d\n", __func__, reuse_i, reuse_n, (int) prompt.size()); + + llama_tokens result; + result.reserve(params.n_draft); + + if (reuse_n == 0) { + llama_kv_cache_clear(ctx); + + prompt.clear(); + } else { + // this happens when a previous draft has been discarded (for example, due to being too small), but the + // target model agreed with it. in this case, we simply pass back the previous results to save compute + if (reuse_i + reuse_n < (int) prompt.size() && prompt[reuse_i + reuse_n] == id_last) { + for (int i = reuse_i + reuse_n + 1; i < (int) prompt.size(); ++i) { + result.push_back(prompt[i]); + + if (params.n_draft <= (int) result.size()) { + break; + } + } + + return result; + } + + if (reuse_i > 0) { + llama_kv_cache_seq_rm (ctx, 0, 0, reuse_i); + llama_kv_cache_seq_add(ctx, 0, reuse_i, -1, -reuse_i); + + prompt.erase(prompt.begin(), prompt.begin() + reuse_i); + } + + if (reuse_n < (int) prompt.size()) { + llama_kv_cache_seq_rm (ctx, 0, reuse_n, -1); + + prompt.erase(prompt.begin() + reuse_n, prompt.end()); + } + } + + // prepare a batch to evaluate any new tokens in the prompt + llama_batch_clear(batch); + + for (size_t i = i_start + reuse_n; i < prompt_tgt.size(); ++i) { + //LOG_DBG("i = %d, i_start = %d, reuse_n = %d, i - i_start = %d, id = %6d\n", i, i_start, reuse_n, i - i_start, prompt_tgt[i]); + llama_batch_add(batch, prompt_tgt[i], i - i_start, { 0 }, false); + + prompt.push_back(prompt_tgt[i]); + } + + // we should rarely end-up here during normal decoding + if (batch.n_tokens > 0) { + //LOG_DBG("%s: draft prompt batch: %s\n", __func__, string_from(ctx, batch).c_str()); + + llama_decode(ctx, batch); + } + + const llama_pos n_past = prompt.size(); + + LOG_DBG("%s: n_past = %d\n", __func__, n_past); + + llama_batch_clear(batch); + llama_batch_add (batch, id_last, n_past, { 0 }, true); + + prompt.push_back(id_last); + + //LOG_DBG("%s: draft prompt: %s\n", __func__, string_from(ctx, prompt).c_str()); + + llama_decode(ctx, batch); + + gpt_sampler_reset(smpl); + + // sample n_draft tokens from the draft model + for (int i = 0; i < params.n_draft; ++i) { + llama_batch_clear(batch); + + gpt_sampler_sample(smpl, ctx, 0, true); + + const auto * cur_p = gpt_sampler_get_candidates(smpl); + + for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { + LOG_DBG(" - draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", + k, i, cur_p->data[k].id, cur_p->data[k].p, llama_token_to_piece(ctx, cur_p->data[k].id).c_str()); + } + + // add drafted token for each sequence + const llama_token id = cur_p->data[0].id; + + // only collect very high-confidence draft tokens + if (cur_p->data[0].p < params.p_min) { + break; + } + + gpt_sampler_accept(smpl, id, true); + + result.push_back(id); + + if (params.n_draft <= (int) result.size()) { + break; + } + + llama_batch_add(batch, id, n_past + i + 1, { 0 }, true); + + // evaluate the drafted tokens on the draft model + llama_decode(ctx, batch); + + prompt.push_back(id); + } + + return result; +} \ No newline at end of file diff --git a/common/speculative.h b/common/speculative.h new file mode 100644 index 00000000..0af10ea4 --- /dev/null +++ b/common/speculative.h @@ -0,0 +1,28 @@ +#pragma once + +#include "llama.h" +#include "common.h" + +struct common_speculative; + +struct common_speculative_params { + int n_draft = 16; // max drafted tokens + int n_reuse = 256; + + float p_min = 0.9f; // min probabiliy required to accept a token in the draft +}; + +struct common_speculative * common_speculative_init(struct llama_context * ctx_dft); + +void common_speculative_free(struct common_speculative * spec); + +bool common_speculative_are_compatible( + const struct llama_context * ctx_tgt, + const struct llama_context * ctx_dft); + +// sample up to n_draft tokens and add them to the batch using the draft model +llama_tokens common_speculative_gen_draft( + struct common_speculative * spec, + struct common_speculative_params params, + const llama_tokens & prompt, + llama_token id_last); \ No newline at end of file diff --git a/examples/server/server.cpp b/examples/server/server.cpp index f67ba107..2a899ed5 100644 --- a/examples/server/server.cpp +++ b/examples/server/server.cpp @@ -6,6 +6,7 @@ #include "sampling.h" #include "json-schema-to-grammar.h" #include "llama.h" +#include "speculative.h" // Change JSON_ASSERT from assert() to GGML_ASSERT: #define JSON_ASSERT GGML_ASSERT @@ -126,13 +127,16 @@ struct server_task_result { struct slot_params { bool stream = true; - bool cache_prompt = false; // remember the prompt to avoid reprocessing all prompt + bool cache_prompt = true; // remember the prompt to avoid reprocessing all prompt int32_t n_keep = 0; // number of tokens to keep from initial prompt int32_t n_discard = 0; // number of tokens after n_keep that may be discarded when shifting context, 0 defaults to half int32_t n_predict = -1; // new tokens to predict std::vector antiprompt; + + struct gpt_sampler_params sampling; + struct common_params_speculative speculative; json input_prefix; json input_suffix; @@ -142,6 +146,12 @@ struct server_slot { int id; int id_task = -1; + llama_batch batch_spec; + + llama_context * ctx_dft = nullptr; + + common_speculative * spec = nullptr; + // the index relative to completion multi-task request size_t index = 0; @@ -231,7 +241,7 @@ struct server_slot { generated_token_probs.clear(); } - bool has_budget(gpt_params &global_params) { + bool has_budget(const gpt_params &global_params) { if (params.n_predict == -1 && global_params.n_predict == -1) { return true; // limitless } @@ -251,6 +261,10 @@ struct server_slot { return state != SLOT_STATE_IDLE; } + bool can_speculate() const { + return ctx_dft && params.speculative.n_max > 0 && params.cache_prompt; + } + void add_token(const completion_token_output & token) { if (!is_processing()) { SLT_WRN(*this, "%s", "slot is not processing\n"); @@ -615,6 +629,9 @@ struct server_context { gpt_params params; + llama_model * model_dft = nullptr; + llama_context_params cparams_dft; + llama_batch batch = {}; bool clean_kv_cache = true; @@ -652,21 +669,71 @@ struct server_context { model = nullptr; } + if (model_dft) { + llama_free_model(model_dft); + model_dft = nullptr; + } + // Clear any sampling context for (server_slot & slot : slots) { if (slot.smpl != nullptr) { gpt_sampler_free(slot.smpl); } + slot.smpl = nullptr; + + llama_free(slot.ctx_dft); + slot.ctx_dft = nullptr; + + common_speculative_free(slot.spec); + slot.spec = nullptr; + + llama_batch_free(slot.batch_spec); } llama_batch_free(batch); } bool load_model(const gpt_params & params_) { + SRV_INF("loading model '%s'\n", params.model.c_str()); + params = params_; // dedicate one sequence to the system prompt params.n_parallel += 1; + + // load draft model first + llama_init_result llama_init_dft; + if (!params.speculative.model.empty()) { + SRV_INF("loading draft model '%s'\n", params.speculative.model.c_str()); + + auto params_dft = params; + + params_dft.model = params.speculative.model; + params_dft.n_ctx = params.speculative.n_ctx; + params_dft.n_gpu_layers = params.speculative.n_gpu_layers; + params_dft.use_mlock = true; + params_dft.n_world = 1; // do not split the draft model across devicesAdd commentMore actions + params_dft.rank = 0; // always load the draft model on the head device + + std::fill_n(params_dft.n_layer_window, params.n_world, 0); + + llama_init_dft = llama_init_from_gpt_params(params_dft); + + model_dft = llama_init_dft.model; + + if (model_dft == nullptr) { + SRV_ERR("failed to load draft model, '%s'\n", params.speculative.model.c_str()); + return false; + } + + cparams_dft = llama_context_params_from_gpt_params(params); + cparams_dft.n_batch = llama_n_ctx(llama_init_dft.context); + cparams_dft.n_world = 1; + cparams_dft.rank = 0; + std::fill_n(cparams_dft.n_layer_window, 32, 0); + cparams_dft.n_layer_window[0] = llama_n_layer(model_dft); + cparams_dft.n_gpu_layers = params.speculative.n_gpu_layers; + } llama_init_result llama_init = llama_init_from_gpt_params(params); @@ -685,6 +752,22 @@ struct server_context { add_bos_token = llama_add_bos_token(model); has_eos_token = !llama_add_eos_token(model); + + if (!params.speculative.model.empty()){ + if (!common_speculative_are_compatible(ctx, llama_init_dft.context)) { + SRV_ERR("the draft model '%s' is not compatible with the target model '%s'\n", params.speculative.model.c_str(), params.model.c_str()); + + llama_free (llama_init_dft.context); + llama_free_model(llama_init_dft.model); + + model_dft = nullptr; + + return false; + } + + // the context is not needed - we will create one for each slot + llama_free(llama_init_dft.context); + } return true; } @@ -708,6 +791,30 @@ struct server_context { slot.id = i; slot.n_ctx = n_ctx_slot; slot.n_predict = params.n_predict; + + if (model_dft) { + slot.batch_spec = llama_batch_init(params.speculative.n_max + 1, 0, 1); + + slot.ctx_dft = llama_new_context_with_model(model_dft, cparams_dft); + + if (llama_context_setup_backend(model_dft, cparams_dft, slot.ctx_dft) == nullptr) { + SRV_ERR("%s: failed to setup context with model '%s'\n", __func__, params.speculative.model.c_str()); + llama_free(slot.ctx_dft); + llama_free_model(model_dft); + return; + } + + if (slot.ctx_dft == nullptr) { + SRV_ERR("%s", "failed to create draft context\n"); + return; + } + + slot.spec = common_speculative_init(slot.ctx_dft); + if (slot.spec == nullptr) { + SRV_ERR("%s", "failed to create speculator\n"); + return; + } + } SLT_INF(slot, "new slot n_ctx_slot = %d\n", slot.n_ctx); @@ -875,6 +982,8 @@ struct server_context { slot_params default_params; // Sampling parameter defaults are loaded from the global server context (but individual requests can still override them) auto default_sparams = params.sparams; + default_params.speculative = params.speculative; + const auto & data = task.data; if (data.count("__oaicompat") != 0) { @@ -886,7 +995,7 @@ struct server_context { } slot.params.stream = json_value(data, "stream", false); - slot.params.cache_prompt = json_value(data, "cache_prompt", false); + slot.params.cache_prompt = json_value(data, "cache_prompt", true); slot.params.n_predict = json_value(data, "n_predict", json_value(data, "max_tokens", default_params.n_predict)); slot.sparams.top_k = json_value(data, "top_k", default_sparams.top_k); slot.sparams.top_p = json_value(data, "top_p", default_sparams.top_p); @@ -904,11 +1013,17 @@ struct server_context { slot.sparams.mirostat_tau = json_value(data, "mirostat_tau", default_sparams.mirostat_tau); slot.sparams.mirostat_eta = json_value(data, "mirostat_eta", default_sparams.mirostat_eta); slot.sparams.penalize_nl = json_value(data, "penalize_nl", default_sparams.penalize_nl); - slot.params.n_keep = json_value(data, "n_keep", slot.params.n_keep); + slot.params.n_keep = json_value(data, "n_keep", params.n_keep); slot.params.n_discard = json_value(data, "n_discard", default_params.n_discard); slot.sparams.seed = json_value(data, "seed", default_sparams.seed); slot.sparams.n_probs = json_value(data, "n_probs", default_sparams.n_probs); slot.sparams.min_keep = json_value(data, "min_keep", default_sparams.min_keep); + + slot.params.speculative.n_min = json_value(data, "speculative.n_min", default_params.speculative.n_min); + slot.params.speculative.n_max = json_value(data, "speculative.n_max", default_params.speculative.n_max); + slot.params.speculative.p_min = json_value(data, "speculative.p_min", default_params.speculative.p_min); + + slot.params.speculative.n_min = std::min(slot.params.speculative.n_max, slot.params.speculative.n_min); // process "json_schema" and "grammar" if (data.contains("json_schema") && !data.at("json_schema").is_null() && data.contains("grammar") && !data.at("grammar").is_null()) { @@ -1049,6 +1164,12 @@ struct server_context { return false; } } + + if (slot.ctx_dft) { + llama_batch_free(slot.batch_spec); + + slot.batch_spec = llama_batch_init(slot.params.speculative.n_max + 1, 0, 1); + } slot.state = SLOT_STATE_PROCESSING_PROMPT; slot.prompt_tokens.clear(); @@ -1059,13 +1180,9 @@ struct server_context { } void kv_cache_clear() { - SRV_DBG("%s", "clearing KV cache\n"); - - // clear the entire KV cache + SRV_DBG("%s", "clearing all KV cache\n"); llama_kv_cache_clear(ctx); - llama_send_kv_cache_clear(ctx); - clean_kv_cache = false; } @@ -1090,7 +1207,7 @@ struct server_context { llama_batch_add(batch, system_tokens[i + j], i + j, { 0 }, false); } - if (llama_decode(ctx, batch) != 0) { + if (llama_decode(ctx, batch, true) != 0) { SRV_ERR("%s", "llama_decode() failed\n"); return; } @@ -1098,7 +1215,8 @@ struct server_context { // assign the system KV cache to all parallel sequences for (int32_t i = 1; i <= params.n_parallel; ++i) { - llama_kv_cache_seq_cp(ctx, 0, i, -1, -1); + llama_kv_cache_seq_cp (ctx, 0, i, -1, -1); + llama_send_kv_cache_seq_cp(ctx, 0, i - 1, -1, -1); } } @@ -1912,7 +2030,6 @@ struct server_context { } // apply context-shift if needed - // TODO: simplify and improve for (server_slot & slot : slots) { if (slot.ga_n == 1) { if (slot.is_processing() && (int) system_tokens.size() + slot.n_past >= slot.n_ctx - 1) { @@ -2087,14 +2204,14 @@ struct server_context { } else { if (!params.ctx_shift) { // if context shift is disabled, we make sure prompt size is smaller than KV size - if ((int) system_tokens.size() + slot.n_prompt_tokens >= slot.n_ctx) { + if ((int)system_tokens.size() + slot.n_prompt_tokens >= slot.n_ctx) { slot.release(); send_error(slot, "the request exceeds the available context size. try increasing the context size or enable context shift", ERROR_TYPE_INVALID_REQUEST); continue; } } if (slot.params.n_keep < 0) { - slot.params.n_keep = slot.n_prompt_tokens; + slot.params.n_keep = (int)system_tokens.size() + slot.n_prompt_tokens + 3; // +3 for tag } slot.params.n_keep = std::min(slot.n_ctx - 4, slot.params.n_keep); @@ -2311,7 +2428,7 @@ struct server_context { 0, 0, 0, // unused }; - const int ret = llama_decode(ctx, batch_view); + const int ret = llama_decode(ctx, batch_view, true); metrics.on_decoded(slots); if (ret != 0) { @@ -2361,38 +2478,101 @@ struct server_context { continue; // continue loop of slots } - completion_token_output result; - const llama_token id = gpt_sampler_sample(slot.smpl, ctx, slot.i_batch - i); + llama_token id; - gpt_sampler_accept(slot.smpl, id, true); + { + completion_token_output result; - slot.n_decoded += 1; - if (slot.n_decoded == 1) { - slot.t_start_generation = ggml_time_us(); - slot.t_prompt_processing = (slot.t_start_generation - slot.t_start_process_prompt) / 1e3; - metrics.on_prompt_eval(slot); + id = gpt_sampler_sample(slot.smpl, ctx, slot.i_batch - i); + + slot.i_batch = -1; + + gpt_sampler_accept(slot.smpl, id, true); + + slot.n_decoded += 1; + if (slot.n_decoded == 1) { + slot.t_start_generation = ggml_time_us(); + slot.t_prompt_processing = (slot.t_start_generation - slot.t_start_process_prompt) / 1e3; + metrics.on_prompt_eval(slot); + } + + result.tok = id; + + const auto * cur_p = gpt_sampler_get_candidates(slot.smpl); + + for (size_t i = 0; i < (size_t) slot.params.sampling.n_probs; ++i) { + result.probs.push_back({ + cur_p->data[i].id, + i >= cur_p->size ? 0.0f : cur_p->data[i].p, + }); + } + + if (!process_token(result, slot)) { + // release slot because of stop condition + slot.release(); + slot.print_timings(); + send_final_response(slot); + metrics.on_prediction(slot); + continue; + } } - result.tok = id; - - const auto * cur_p = gpt_sampler_get_candidates(slot.smpl); - - for (size_t i = 0; i < (size_t) slot.sparams.n_probs; ++i) { - result.probs.push_back({ - cur_p->data[i].id, - i >= cur_p->size ? 0.0f : cur_p->data[i].p, - }); + // check if the slot supports speculative decoding + if (!slot.can_speculate()) { + continue; } - if (!process_token(result, slot)) { - // release slot because of stop condition - slot.release(); - slot.print_timings(); - send_final_response(slot); - metrics.on_prediction(slot); + struct common_speculative_params params_spec; + params_spec.n_draft = slot.params.speculative.n_max; + params_spec.n_reuse = llama_n_ctx(slot.ctx_dft) - slot.params.speculative.n_max; + params_spec.p_min = slot.params.speculative.p_min; + + llama_tokens draft = common_speculative_gen_draft(slot.spec, params_spec, slot.cache_tokens, id); + + // ignore small drafts + if (slot.params.speculative.n_min > (int) draft.size()) { + continue; } - slot.i_batch = -1; + // construct the speculation batch + llama_batch_clear(slot.batch_spec); + llama_batch_add (slot.batch_spec, id, slot.n_past, { slot.id + 1 }, true); + + for (size_t i = 0; i < draft.size(); ++i) { + llama_batch_add(slot.batch_spec, draft[i], slot.n_past + 1 + i, { slot.id + 1 }, true); + } + + llama_decode(ctx, slot.batch_spec, true); + + // the accepted tokens from the speculation + const auto ids = gpt_sampler_sample_and_accept_n(slot.smpl, ctx, draft); + + slot.n_past += ids.size(); + slot.n_decoded += ids.size(); + + slot.cache_tokens.push_back(id); + slot.cache_tokens.insert(slot.cache_tokens.end(), ids.begin(), ids.end() - 1); + + llama_kv_cache_seq_rm (ctx, slot.id + 1, slot.n_past, -1); + llama_send_kv_cache_seq_rm(ctx, slot.id , slot.n_past, -1); + + for (size_t i = 0; i < ids.size(); ++i) { + completion_token_output result; + + result.tok = ids[i]; + + if (!process_token(result, slot)) { + // release slot because of stop condition + slot.release(); + slot.print_timings(); + send_final_response(slot); + metrics.on_prediction(slot); + break; + } + } + + SRV_DBG("accepted %d/%d draft tokens\n", (int) ids.size() - 1, (int) draft.size()); + } } @@ -3388,6 +3568,8 @@ int main(int argc, char ** argv) { LOG_INF("%s: loading model\n", __func__); if (!ctx_server.load_model(params)) { + char * stop_signal = nullptr; + llama_free_sockets(ctx_server.ctx, &stop_signal); clean_up(); t.join(); LOG_ERR("%s: exiting due to model loading error\n", __func__); @@ -3408,7 +3590,7 @@ int main(int argc, char ** argv) { } // print sample chat example to make it clear which template is used - LOG_INF("%s: chat template, built_in: %d, chat_example: '%s'\n", __func__, params.chat_template.empty(), llama_chat_format_example(ctx_server.model, params.chat_template).c_str()); + // LOG_INF("%s: chat template, built_in: %d, chat_example: '%s'\n", __func__, params.chat_template.empty(), llama_chat_format_example(ctx_server.model, params.chat_template).c_str()); ctx_server.queue_tasks.on_new_task(std::bind( &server_context::process_single_task, &ctx_server, std::placeholders::_1)); diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp index adf6255e..ad610ac8 100644 --- a/examples/speculative/speculative.cpp +++ b/examples/speculative/speculative.cpp @@ -12,7 +12,7 @@ #include #include -#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 100 +#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 128 #define SPEC_VOCAB_CHECK_START_TOKEN_ID 5 struct seq_draft { @@ -41,7 +41,7 @@ int main(int argc, char ** argv) { gpt_init(); - if (params.model_draft.empty()) { + if (params.speculative.model.empty()) { LOG_ERR("%s: --model-draft is required\n", __func__); return 1; } @@ -65,23 +65,30 @@ int main(int argc, char ** argv) { llama_context * ctx_tgt = NULL; llama_context * ctx_dft = NULL; + // load the draft model + // make a hard copy of params to use for the draft model + gpt_params params_draft = params; + params_draft.model = params_draft.speculative.model; + params_draft.n_gpu_layers = params_draft.n_gpu_layers_draft; + params_draft.n_world = 1; // do not split the draft model across devices + params_draft.rank = 0; // always load the draft model on the head device + params_draft.use_mlock = true; // always use mlock for the draft model + std::fill_n(params_draft.n_layer_window, params.n_world, 0); + + if (params_draft.draft_cpuparams.n_threads > 0) { + params_draft.cpuparams.n_threads = params_draft.draft_cpuparams.n_threads; + } + + params_draft.cpuparams_batch.n_threads = params_draft.draft_cpuparams_batch.n_threads; + llama_init_result llama_init_dft = llama_init_from_gpt_params(params_draft); + model_dft = llama_init_dft.model; + ctx_dft = llama_init_dft.context; + // load the target model llama_init_result llama_init_tgt = llama_init_from_gpt_params(params); model_tgt = llama_init_tgt.model; ctx_tgt = llama_init_tgt.context; - // load the draft model - params.model = params.model_draft; - params.n_gpu_layers = params.n_gpu_layers_draft; - if (params.draft_cpuparams.n_threads > 0) { - params.cpuparams.n_threads = params.draft_cpuparams.n_threads; - } - - params.cpuparams_batch.n_threads = params.draft_cpuparams_batch.n_threads; - llama_init_result llama_init_dft = llama_init_from_gpt_params(params); - model_dft = llama_init_dft.model; - ctx_dft = llama_init_dft.context; - const bool vocab_type_tgt = llama_vocab_type(model_tgt); LOG_DBG("vocab_type tgt: %d\n", vocab_type_tgt); @@ -161,11 +168,8 @@ int main(int argc, char ** argv) { const auto t_enc_end = ggml_time_us(); - // the 2 models should have the same vocab - //GGML_ASSERT(n_vocab == llama_n_vocab(model_dft)); - // how many tokens to draft each time - int n_draft = params.n_draft; + int n_draft = params.speculative.n_max; int n_predict = 0; int n_drafted = 0; @@ -180,8 +184,6 @@ int main(int argc, char ** argv) { // target model sampling context (reuse the llama_context's sampling instance) struct gpt_sampler * smpl = gpt_sampler_init(model_tgt, params.sparams); - struct llama_sampler * softmax = llama_sampler_init_softmax(); - // draft sequence data std::vector drafts(n_seq_dft); @@ -258,10 +260,13 @@ int main(int argc, char ** argv) { float r = u_dist(rng); llama_token_data_array dist_dft = { drafts[s].dists[i_dft].data() , drafts[s].dists[i_dft].size(), LLAMA_TOKEN_NULL, true }; - //GGML_ASSERT(dist_tgt.size <= dist_dft.size); + // if (dist_tgt.size > dist_dft.size) { + // LOG_ERR("dist_tgt.size (%zu) must be less than or equal to dist_dft.size (%zu)\n", dist_tgt.size, dist_dft.size); + // GGML_ASSERT(dist_tgt.size <= dist_dft.size); + // } // acquire the token probabilities assigned by the draft and target models - for (size_t i = 0; i < dist_tgt.size; i++) { + for (size_t i = 0; i < dist_tgt.size && i < dist_dft.size; i++) { if (dist_tgt.data[i].id == drafts[s].tokens[i_dft]) { p_tgt = dist_tgt.data[i].p; } @@ -406,7 +411,6 @@ int main(int argc, char ** argv) { { LOG_DBG("the sampled target token (%d, '%s') did not match, or we ran out of drafted tokens\n", token_id, token_str.c_str()); - // TODO: simplify { LOG_DBG("keeping sequence %d, n_past_tgt = %d, n_past_dft = %d\n", s_keep, n_past_tgt, n_past_dft); @@ -418,6 +422,12 @@ int main(int argc, char ** argv) { llama_kv_cache_seq_keep(ctx_tgt, s_keep); llama_kv_cache_seq_cp (ctx_tgt, s_keep, 0, -1, -1); llama_kv_cache_seq_keep(ctx_tgt, 0); + + // notify other devices to manage the KV cache in the same way + llama_send_kv_cache_seq_rm (ctx_tgt, s_keep, n_past_tgt, -1); + llama_send_kv_cache_seq_keep(ctx_tgt, s_keep); + llama_send_kv_cache_seq_cp (ctx_tgt, s_keep, 0, -1, -1); + llama_send_kv_cache_seq_keep(ctx_tgt, 0); } for (int s = 0; s < n_seq_dft; ++s) { @@ -435,7 +445,6 @@ int main(int argc, char ** argv) { llama_batch_add (batch_dft, token_id, n_past_dft, { 0 }, true); llama_kv_cache_seq_rm(ctx_dft, 0, n_past_dft, -1); - // LOG_DBG("dft batch: %s\n", LOG_BATCH_TOSTR_PRETTY(ctx_dft, batch_dft).c_str()); llama_decode(ctx_dft, batch_dft); ++n_past_dft; @@ -575,12 +584,13 @@ int main(int argc, char ** argv) { // evaluate the target model on the drafted tokens { - llama_kv_cache_seq_keep(ctx_tgt, 0); + llama_kv_cache_seq_keep (ctx_tgt, 0); + llama_send_kv_cache_seq_keep(ctx_tgt, 0); for (int s = 1; s < n_seq_dft; ++s) { - llama_kv_cache_seq_cp(ctx_tgt, 0, s, -1, -1); + llama_kv_cache_seq_cp (ctx_tgt, 0, s, -1, -1); + llama_send_kv_cache_seq_cp(ctx_tgt, 0, s, -1, -1); } - // LOG_DBG("target batch: %s\n", LOG_BATCH_TOSTR_PRETTY(ctx_tgt, batch_tgt).c_str()); llama_decode(ctx_tgt, batch_tgt); ++n_past_tgt; } @@ -612,19 +622,21 @@ int main(int argc, char ** argv) { LOG_INF("\n"); LOG_INF("draft:\n\n"); - // TODO: print sampling/grammar timings for all drafts + llama_perf_context_print(ctx_dft); LOG_INF("\n"); LOG_INF("target:\n\n"); gpt_perf_print(ctx_tgt, smpl); + char * stop_signal = nullptr; + llama_free_sockets(ctx_tgt, &stop_signal); + gpt_sampler_free(smpl); for (int s = 0; s < n_seq_dft; ++s) { gpt_sampler_free(drafts[s].smpl); } - llama_sampler_free(softmax); llama_batch_free(batch_dft); llama_free(ctx_tgt); diff --git a/include/llama.h b/include/llama.h index 45a2bad3..5795208b 100644 --- a/include/llama.h +++ b/include/llama.h @@ -67,6 +67,16 @@ extern "C" { typedef int32_t llama_token; typedef int32_t llama_seq_id; + enum backend_type { + BACKEND_CPU = 0, + BACKEND_CUDA = 1, + BACKEND_METAL = 2, + BACKEND_VULKAN = 3, + BACKEND_KOMPUTE = 4, + BACKEND_GPUBLAS = 5, + BACKEND_SYCL = 6 + }; + enum llama_vocab_type { LLAMA_VOCAB_TYPE_NONE = 0, // For models without vocab LLAMA_VOCAB_TYPE_SPM = 1, // LLaMA tokenizer based on byte-level BPE with byte fallback @@ -317,6 +327,7 @@ extern "C" { bool use_mlock; // force system to keep model in RAM bool check_tensors; // validate model tensor data bool keep_out_in_metal; // whether to keep output weights in metal memory + bool keep_out_in_cuda; // whether to run the output layer on CUDA }; // NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations @@ -329,7 +340,9 @@ extern "C" { uint32_t n_cycles; // number of cycles to output one token bool prefetch; // whether to prefetch layer weights bool force; // force to start prefetching after computation + float master_priority; // priority to assign workload to the master (set 1.01 to use master first, and 0.99 to offload to other devices) bool keep_out_in_metal; // whether to keep output weights in metal memory + bool keep_out_in_cuda; // whether to run the output layer on CUDA char * master_ip; // ip address of the master node char * next_node_ip; // ip address of the next node uint32_t data_port; // base data port for this node @@ -461,6 +474,12 @@ extern "C" { struct llama_model_params params); LLAMA_API void llama_free_model(struct llama_model * model); + + enum NodeType{ + NODE_TYPE_WORKER, + NODE_TYPE_FORWARDER, + NODE_TYPE_EXIT, + }; 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); @@ -468,7 +487,12 @@ extern "C" { LLAMA_API int llama_send_device_info (struct llama_context * ctx, struct device_info * dev_info); LLAMA_API int llama_bcast_startup_args(struct llama_context * ctx, uint32_t rank, struct startup_args * args); LLAMA_API int llama_bcast_layer_setup (struct llama_context * ctx, uint32_t * n_layer_window, uint32_t * n_gpu_layers); - LLAMA_API int llama_rebuild_topo (struct llama_context * ctx, uint32_t * n_layer_window, struct device_info * dev_info_set); + LLAMA_API int llama_rebuild_topo (struct llama_context * ctx, + uint32_t * n_layer_window, + struct device_info * desv_info_set, + NodeType* node_type, + char * is_forwarder); + LLAMA_API int llama_forward_messages (struct llama_context * ctx); LLAMA_API int llama_recv_layer_setup (struct llama_context * ctx, uint32_t * n_layer_window, uint32_t * n_gpu_layers); LLAMA_API int llm_load_tensors( @@ -479,7 +503,9 @@ extern "C" { LLAMA_API void llama_update_context_with_rankworld( struct llama_context * ctx, uint32_t rank, - uint32_t n_world); + uint32_t n_world, + uint32_t worker_rank, + uint32_t n_worker); LLAMA_API struct llama_context * llama_new_context_with_model( struct llama_model * model, @@ -571,10 +597,11 @@ extern "C" { int64_t * gpu_buf, const struct llama_model * model, const struct llama_context_params cparams, - bool use_gpu, - bool is_master, - int n_layers, - int n_gpu_layers); + enum backend_type backend, + int my_rank, + struct model_bytes n_bytes, + bool offload, + bool has_gpu_layers); // Return the size of KV cache in the model LLAMA_API void llama_total_kv_size( @@ -772,6 +799,11 @@ extern "C" { LLAMA_API void llama_kv_cache_seq_keep( struct llama_context * ctx, llama_seq_id seq_id); + + // Notify other nodes to keep only the specified sequence in their KV cache + LLAMA_API void llama_send_kv_cache_seq_keep( + struct llama_context * ctx, + llama_seq_id seq_id); // Adds relative position "delta" to all tokens that belong to the specified sequence and have positions in [p0, p1) // If the KV cache is RoPEd, the KV data is updated accordingly: @@ -970,7 +1002,8 @@ extern "C" { // < 0 - error LLAMA_API int32_t llama_decode( struct llama_context * ctx, - struct llama_batch batch); + struct llama_batch batch, + bool server_mode = false); // Set the number of threads used for decoding // n_threads is the number of threads used for generation (single token) diff --git a/src/llama.cpp b/src/llama.cpp index 95e6796e..6360a98a 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -11,6 +11,7 @@ #include "ggml-backend.h" #include "profiler.h" +#include "network-utils.h" #ifdef GGML_USE_RPC # include "ggml-rpc.h" @@ -246,12 +247,12 @@ static void zeros(std::ofstream & file, size_t n) { } // zmq helpers -static std::vector dev_infos_to_messages(const device_info* infos, - uint32_t n_world){ +static std::vector topohelper_to_messages(const TopoRebuildHelperInfo* infos, + uint32_t n_world){ std::vector res; for (uint32_t i = 0; i < n_world; ++i) { char * buffer = nullptr; - size_t buffer_size = serialize(&infos[i], &buffer); + size_t buffer_size = infos[i].serialize(&buffer); res.emplace_back(buffer, buffer_size); free(buffer); } @@ -2669,6 +2670,9 @@ static_assert(std::is_trivially_copyable::value, "llama_hparams m struct llama_cparams { uint32_t n_world; uint32_t rank; + NodeType node_type; + uint32_t n_worker; + uint32_t worker_rank; uint32_t original_next_rank; // original rank of the next node uint32_t n_layer_window[32]; bool prefetch; @@ -3510,8 +3514,8 @@ struct llama_context { 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"; + std::string master_ip = "127.0.0.1"; + std::string next_node_ip = "127.0.0.1"; uint32_t next_node_data_port = 9000; uint32_t next_node_signal_port = 10000; uint32_t master_data_port = 9000; @@ -3708,6 +3712,10 @@ void llama_profile_device( dev_info->memory.total_physical = round(device_physical_memory(false) / (double)(1 << 30) * 100) / 100; dev_info->memory.available_physical = round(device_physical_memory(true) / (double)(1 << 30) * 100) / 100; + + GGML_ASSERT(dev_info->memory.total_physical > 0 && "Failed to parse total physical memory\n"); + GGML_ASSERT(dev_info->memory.available_physical > 0 && "Failed to parse available physical memory\n"); + dev_info->memory.used_can_swap = round(device_swappable_memory() / (double)(1 << 30) * 100) / 100; dev_info->memory.total_swap = round(device_swap_memory(false) / (double)(1 << 30) * 100) / 100; dev_info->memory.available_swap = round(device_swap_memory(true) / (double)(1 << 30) * 100) / 100; @@ -3744,16 +3752,15 @@ void llama_profile_device( #if defined(GGML_USE_METAL) || defined(GGML_USE_CUDA) dev_info->gpu_props.name = gpu_props.name; dev_info->gpu_props.description = gpu_props.description; - - // reserved/limit memory to avoid potential OOM, default to 300 MiB - dev_info->gpu_props.memory_free = round(gpu_props.memory_free / (double)(1 << 30) * 100) / 100; - dev_info->gpu_props.memory_free = std::min((float)gpu_mem, dev_info->gpu_props.memory_free) - 0.3; - + dev_info->gpu_props.memory_free = round(gpu_props.memory_free / (double)(1 << 30) * 100) / 100; + dev_info->gpu_props.memory_free = std::min((float)gpu_mem, dev_info->gpu_props.memory_free); dev_info->gpu_props.memory_total = round(gpu_props.memory_total / (double)(1 << 30) * 100) / 100; dev_info->gpu_props.metal_read_vram_bw = device_metal_read_vram_bw(); dev_info->gpu_props.cuda_read_vram_bw = device_cuda_read_vram_bw(); dev_info->gpu_props.metal_mem_cpy_delay = device_metal_mem_copy(model); dev_info->gpu_props.cuda_mem_cpy_delay = device_cuda_mem_copy(model); +#else + (void)gpu_mem; #endif if (is_dtype_exist(n_params, GGML_TYPE_F32)) { @@ -7570,10 +7577,7 @@ static void llm_load_qwen2_tensors( uint32_t n_world, uint32_t my_rank, const uint32_t * n_layer_window, - bool * use_mmap_buffer, bool set_needed) { - (void)use_mmap_buffer; // unused in this function - const auto tn = LLM_TN(model.arch); ggml_context * ctx_input = nullptr; @@ -7656,6 +7660,7 @@ static bool llm_load_tensors_impl( int main_gpu, bool use_mlock, bool keep_out_in_metal, + bool keep_out_in_cuda, llama_progress_callback progress_callback, void * progress_callback_user_data) { auto & hparams = model.hparams; @@ -7688,10 +7693,10 @@ static bool llm_load_tensors_impl( GGML_ASSERT(local_i != -1); if (local_i % window_size >= window_size - n_gpu_layers) { - // LLAMA_LOG_INFO("Layer %i assigned to gpu (cache index %i)\n", i, local_i); + LLAMA_LOG_DEBUG("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); + LLAMA_LOG_DEBUG("Layer %i assigned to cpu (cache index %i)\n", i, local_i); model.buft_layer[local_i] = llama_default_buffer_type_cpu(model, true); } } @@ -7700,9 +7705,15 @@ static bool llm_load_tensors_impl( // assign the input and output layers on CPU by default if (my_rank == 0) { model.buft_input = llama_default_buffer_type_cpu(model, true); - model.buft_output = llama_default_buffer_type_cpu(model, true); - // LLAMA_LOG_INFO("Layer input assigned to cpu\n"); - // LLAMA_LOG_INFO("Layer output assigned to cpu\n"); + LLAMA_LOG_DEBUG("Layer input assigned to cpu\n"); + + if (keep_out_in_cuda) { + model.buft_output = llama_default_buffer_type_offload(model, main_gpu); + LLAMA_LOG_DEBUG("Layer output assigned to gpu\n"); + } else { + model.buft_output = llama_default_buffer_type_cpu(model, true); + LLAMA_LOG_DEBUG("Layer output assigned to cpu\n"); + } } // count used buffer types @@ -8310,7 +8321,7 @@ static bool llm_load_tensors_impl( } } break; case LLM_ARCH_QWEN2: - llm_load_qwen2_tensors(ml, model, ctx_map, n_world, my_rank, n_layer_window, &use_mmap_buffer, true); + llm_load_qwen2_tensors(ml, model, ctx_map, n_world, my_rank, n_layer_window, true); break; case LLM_ARCH_QWEN2MOE: { @@ -9629,7 +9640,8 @@ int llm_load_tensors( try { if (!llm_load_tensors_impl( *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.keep_out_in_metal, params.progress_callback, params.progress_callback_user_data + params.main_gpu, params.use_mlock, params.keep_out_in_metal, params.keep_out_in_cuda, params.progress_callback, + params.progress_callback_user_data )) { return -2; } @@ -11281,8 +11293,6 @@ struct llm_build_context { cur = llm_build_out_embd(ctx0, lctx, hparams, cb); - // 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); @@ -12823,18 +12833,19 @@ struct llm_build_context { } std::vector build_qwen2() { + // 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); // create a vector to hold the subgraphs std::vector sub_gfs; struct ggml_cgraph * sub_gf = nullptr; - struct ggml_tensor * cur = nullptr; - struct ggml_tensor * inpL = nullptr; - struct ggml_tensor * inpB = nullptr; - + struct ggml_tensor * cur = nullptr; + struct ggml_tensor * inpL = nullptr; + struct ggml_tensor * inpB = nullptr; 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; @@ -12850,7 +12861,7 @@ struct llm_build_context { sub_gfs.push_back(sub_gf); sub_gf = nullptr; - inpL = nullptr; + inpL = nullptr; } // inpB - contains the output embedding from other nodes @@ -12862,6 +12873,7 @@ struct llm_build_context { // KQ_mask (mask for 1 head, it will be broadcasted to all heads) struct ggml_tensor * KQ_mask = build_inp_KQ_mask(); + 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) { 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 @@ -12870,7 +12882,6 @@ struct llm_build_context { sub_gfs.push_back(sub_gf); sub_gf = nullptr; } - // synchronous input tensor if (inpL != inpB) { inpL = inpB; } @@ -12900,21 +12911,27 @@ struct llm_build_context { // compute Q and K and RoPE them struct ggml_tensor * Qcur = llm_build_lora_mm(lctx, ctx0, model.layers[local_il].wq, cur); cb(Qcur, "Qcur", il); - Qcur = ggml_add(ctx0, Qcur, model.layers[local_il].bq); - cb(Qcur, "Qcur", il); + 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[local_il].wk, cur); cb(Kcur, "Kcur", il); - Kcur = ggml_add(ctx0, Kcur, model.layers[local_il].bk); - cb(Kcur, "Kcur", il); + 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[local_il].wv, cur); cb(Vcur, "Vcur", il); - Vcur = ggml_add(ctx0, Vcur, model.layers[local_il].bv); - cb(Vcur, "Vcur", il); + if (model.layers[local_il].bv) { + Vcur = ggml_add(ctx0, Vcur, model.layers[local_il].bv); + cb(Vcur, "Vcur", il); + } Qcur = ggml_rope_ext( - ctx0, ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens), inp_pos, nullptr, + ctx0, ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens), inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow ); @@ -12929,7 +12946,7 @@ struct llm_build_context { 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, 1.0f/sqrtf(float(n_embd_head)), cb, il); + Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, kq_scale, cb, il); } if (il == n_layer - 1) { @@ -12939,7 +12956,7 @@ struct llm_build_context { inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); } - 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 @@ -12957,6 +12974,8 @@ struct llm_build_context { cb(cur, "ffn_out", il); cur = ggml_add(ctx0, cur, ffn_inp); // shortcut + cb(cur, "ffn_out", il); + cur = lctx.cvec.apply_to(ctx0, cur, il); cb(cur, "l_out", il); @@ -17133,7 +17152,6 @@ static std::vector llama_build_graph( } break; case LLM_ARCH_QWEN2: { - // result.push_back(llm.build_qwen2()); result = llm.build_qwen2(); } break; case LLM_ARCH_QWEN2MOE: @@ -17937,6 +17955,9 @@ struct sync_meta { llama_pos cp_p0 = 0; llama_pos cp_p1 = 0; + bool kv_seq_keep = false; + llama_seq_id keep_seq_id = 0; + // signal to divide the kv cache range bool kv_seq_div = false; llama_seq_id div_seq_id = 0; @@ -17945,7 +17966,7 @@ struct sync_meta { int div_factor = 1; }; -static void llama_send_meta(zmq::socket_t & socket, struct sync_meta * meta) { +static void llama_send_meta(zmq::socket_t & socket, struct sync_meta * meta, bool align_seq_ids = false) { GGML_ASSERT(meta != nullptr); try { std::vector send_msgs; @@ -17956,23 +17977,24 @@ static void llama_send_meta(zmq::socket_t & socket, struct sync_meta * meta) { if (meta->pos != nullptr) { send_msgs.emplace_back("pos", strlen("pos")); - send_msgs.emplace_back(meta->pos, meta->n_ctx * sizeof(llama_pos)); + send_msgs.emplace_back(meta->pos, meta->n_tokens * sizeof(llama_pos)); } if (meta->n_seq_id != nullptr) { - GGML_ASSERT(meta->n_ctx > 0); + GGML_ASSERT(meta->n_tokens > 0); send_msgs.emplace_back("n_seq_id", strlen("n_seq_id")); - send_msgs.emplace_back(meta->n_seq_id, meta->n_ctx * sizeof(int32_t)); + send_msgs.emplace_back(meta->n_seq_id, meta->n_tokens * sizeof(int32_t)); // here we assume only a single seq_id per token is needed // pack all single seq_id values into a contiguous array - llama_seq_id * all_seq_ids = (llama_seq_id *) malloc(meta->n_ctx * sizeof(llama_seq_id)); - for (uint32_t i = 0; i < meta->n_ctx; ++i) { - all_seq_ids[i] = meta->seq_id[i][0]; + llama_seq_id * all_seq_ids = (llama_seq_id *) malloc(meta->n_tokens * sizeof(llama_seq_id)); + int seq_id_offset = align_seq_ids ? 1 : 0; + for (int32_t i = 0; i < meta->n_tokens; ++i) { + all_seq_ids[i] = meta->seq_id[i][0] - seq_id_offset; } send_msgs.emplace_back("seq_id", strlen("seq_id")); - send_msgs.emplace_back(all_seq_ids, meta->n_ctx * sizeof(llama_seq_id)); + send_msgs.emplace_back(all_seq_ids, meta->n_tokens * sizeof(llama_seq_id)); free(all_seq_ids); } @@ -18038,8 +18060,14 @@ static int llama_recv_meta(zmq::socket_t & socket, struct sync_meta * meta) { return 0; } + if (cmd == "kv_seq_keep" && recv_msgs.size() == 2) { + meta->kv_seq_keep = true; + std::memcpy(&meta->keep_seq_id, recv_msgs[idx++].data(), sizeof(meta->keep_seq_id)); + return 0; + } + if (cmd == "kv_seq_div" && recv_msgs.size() == 5) { - meta->kv_seq_div = true; + meta->kv_seq_div = true; std::memcpy(&meta->div_seq_id, recv_msgs[idx++].data(), sizeof(meta->div_seq_id)); std::memcpy(&meta->div_p0, recv_msgs[idx++].data(), sizeof(meta->div_p0)); std::memcpy(&meta->div_p1, recv_msgs[idx++].data(), sizeof(meta->div_p1)); @@ -18057,23 +18085,23 @@ static int llama_recv_meta(zmq::socket_t & socket, struct sync_meta * meta) { } if (key == "pos") { - meta->pos = (llama_pos *) malloc(meta->n_ctx * sizeof(llama_pos)); - std::memcpy(meta->pos, data_msg.data(), meta->n_ctx * sizeof(llama_pos)); + meta->pos = (llama_pos *) malloc(meta->n_tokens * sizeof(llama_pos)); + std::memcpy(meta->pos, data_msg.data(), meta->n_tokens * sizeof(llama_pos)); } if (key == "n_seq_id") { - GGML_ASSERT(meta->n_ctx > 0); - GGML_ASSERT(data_msg.size() == meta->n_ctx * sizeof(int32_t)); - meta->n_seq_id = (int32_t *) malloc(meta->n_ctx * sizeof(int32_t)); - std::memcpy(meta->n_seq_id, data_msg.data(), meta->n_ctx * sizeof(int32_t)); + GGML_ASSERT(meta->n_tokens > 0); + GGML_ASSERT(data_msg.size() == meta->n_tokens * sizeof(int32_t)); + meta->n_seq_id = (int32_t *) malloc(meta->n_tokens * sizeof(int32_t)); + std::memcpy(meta->n_seq_id, data_msg.data(), meta->n_tokens * sizeof(int32_t)); } if (key == "seq_id") { - GGML_ASSERT(meta->n_ctx > 0); - GGML_ASSERT(data_msg.size() == meta->n_ctx * sizeof(llama_seq_id)); + GGML_ASSERT(meta->n_tokens > 0); + GGML_ASSERT(data_msg.size() == meta->n_tokens * sizeof(llama_seq_id)); const llama_seq_id * all_seq_ids = (llama_seq_id *) data_msg.data(); - meta->seq_id = (llama_seq_id **) malloc(meta->n_ctx * sizeof(llama_seq_id *)); - for (uint32_t i = 0; i < meta->n_ctx; ++i) { + meta->seq_id = (llama_seq_id **) malloc(meta->n_tokens * sizeof(llama_seq_id *)); + for (int32_t i = 0; i < meta->n_tokens; ++i) { meta->seq_id[i] = (llama_seq_id *) malloc(sizeof(llama_seq_id)); meta->seq_id[i][0] = all_seq_ids[i]; } @@ -18317,13 +18345,17 @@ static void manage_graph_tensors(struct ggml_cgraph * cgraph, int advice, bool f // static int llama_decode_internal( llama_context & lctx, - llama_batch batch_all) { // TODO: rename back to batch + llama_batch batch_all, + bool server_mode) { 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; + const uint32_t n_worker = cparams.n_worker; + const uint32_t worker_rank = cparams.worker_rank; + lctx.is_encoding = false; const uint32_t n_tokens_all = batch_all.n_tokens; if (my_rank != 0) { @@ -18379,7 +18411,7 @@ static int llama_decode_internal( sync_meta meta; meta.n_ctx = cparams.n_ctx; - bool is_last_dev = (my_rank == n_world - 1); + bool is_last_dev = (worker_rank == n_worker - 1); if (my_rank != 0) { if (llama_recv_meta(*lctx.recv_socket, &meta) == -1) { @@ -18389,16 +18421,16 @@ static int llama_decode_internal( if (meta.n_tokens > 0) { batch_all.n_tokens = meta.n_tokens; if (meta.pos != nullptr) { - batch_all.pos = (llama_pos *) malloc(cparams.n_ctx * sizeof(llama_pos)); - std::memcpy(batch_all.pos, meta.pos, cparams.n_ctx * sizeof(llama_pos)); + batch_all.pos = (llama_pos *) malloc(meta.n_tokens * sizeof(llama_pos)); + std::memcpy(batch_all.pos, meta.pos, meta.n_tokens * sizeof(llama_pos)); } if (meta.n_seq_id != nullptr) { - batch_all.n_seq_id = (int32_t *) malloc(cparams.n_ctx * sizeof(int32_t)); - std::memcpy(batch_all.n_seq_id, meta.n_seq_id, cparams.n_ctx * sizeof(int32_t)); + batch_all.n_seq_id = (int32_t *) malloc(meta.n_tokens * sizeof(int32_t)); + std::memcpy(batch_all.n_seq_id, meta.n_seq_id, meta.n_tokens * sizeof(int32_t)); } if (meta.seq_id != nullptr) { - batch_all.seq_id = (llama_seq_id **) malloc(cparams.n_ctx * sizeof(llama_seq_id *)); - for (size_t i = 0; i < cparams.n_ctx; ++i) { + batch_all.seq_id = (llama_seq_id **) malloc(meta.n_tokens * sizeof(llama_seq_id *)); + for (int32_t i = 0; i < meta.n_tokens; ++i) { batch_all.seq_id[i] = (llama_seq_id *) malloc(sizeof(llama_seq_id)); batch_all.seq_id[i][0] = meta.seq_id[i][0]; } @@ -18443,6 +18475,14 @@ static int llama_decode_internal( return -1; } + if (kv_cache_op(meta.kv_seq_keep, + [&]{ llama_kv_cache_seq_keep (&lctx, meta.keep_seq_id); }, + [&]{ llama_send_kv_cache_seq_keep(&lctx, meta.keep_seq_id); }, + is_last_dev)) { + LLAMA_LOG_DEBUG("%s: received signal kv_cache_seq_keep\n", __func__); + return -1; + } + if (kv_cache_op(meta.kv_seq_div, [&]{ llama_kv_cache_seq_div (&lctx, meta.div_seq_id, meta.div_p0, meta.div_p1, meta.div_factor); }, [&]{ llama_send_kv_cache_seq_div(&lctx, meta.div_seq_id, meta.div_p0, meta.div_p1, meta.div_factor); }, @@ -18460,7 +18500,7 @@ static int llama_decode_internal( meta.logits = batch_all.logits; meta.all_pos_0 = batch_all.all_pos_0; meta.all_pos_1 = batch_all.all_pos_1; - llama_send_meta(*lctx.send_socket, &meta); + llama_send_meta(*lctx.send_socket, &meta, server_mode); } lctx.sbatch.from_batch(batch_all, n_embd, @@ -19217,19 +19257,19 @@ static void llama_kv_cache_update_internal(struct llama_context & lctx) { GGML_ABORT("Deepseek2 does not support K-shift"); } - for (size_t i = 0; i < lctx.sched.size(); ++i) { - ggml_backend_sched_reset(lctx.sched[i]); + auto * sched = lctx.sched.at(0); - ggml_cgraph * gf = llama_build_graph_k_shift(lctx); + ggml_backend_sched_reset(sched); - ggml_backend_sched_alloc_graph(lctx.sched[i], gf); + ggml_cgraph * gf = llama_build_graph_k_shift(lctx); - llama_set_k_shift(lctx); + ggml_backend_sched_alloc_graph(sched, gf); - llama_graph_compute(lctx, gf, lctx.sched[i], lctx.cparams.n_threads, lctx.threadpool); + llama_set_k_shift(lctx); - need_reserve = true; - } + llama_graph_compute(lctx, gf, sched, lctx.cparams.n_threads, lctx.threadpool); + + need_reserve = true; { auto & kv_self = lctx.kv_self; @@ -20365,6 +20405,7 @@ struct llama_model_params llama_model_default_params() { /*.use_mlock =*/ false, /*.check_tensors =*/ false, /*.keep_out_in_metal =*/ true, + /*.keep_out_in_cuda =*/ false, }; #ifdef GGML_USE_METAL @@ -20384,7 +20425,9 @@ struct llama_context_params llama_context_default_params() { /*.n_cycles =*/ 0, /*.prefetch =*/ false, /*.force =*/ false, + /*.master_priority =*/ 1.01, /*.keep_out_in_metal =*/ true, + /*.keep_out_in_cuda =*/ false, /*.master_ip =*/ nullptr, /*.next_node_ip =*/ nullptr, /*.data_port =*/ 9000, @@ -20580,6 +20623,37 @@ static uint32_t map_rank_to_port(uint32_t rank, uint32_t data_port) { return data_port + rank; } +static std::string try_connect(llama_context * ctx, uint32_t rank, TopoRebuildHelperInfo * infos, uint32_t n_world, zmq::socket_t ** socket){ + LLAMA_LOG_ERROR( + "ERROR: `try_connect()` is not supported because " + "the data_port, signal_port, next_data_port, etc. are specified, " + "but the current TopoRebuildHelper only handles scenarios where " + "data ports are fixed (e.g., rank0 data port = 9000, rank1 data port = 9001, " + "rank2 data port = 9002). In those cases the port information is not required " + "and only the IP address is used. To support dynamic port specifications, " + "the rebuild helper and device info must be enhanced to handle port " + "information as well. (plan to fix in the future)" + ); + auto prev_rank = (rank - 1 + n_world) % n_world; + std::string ip = infos[prev_rank].dev_info.next_ip; + auto port = map_rank_to_port(rank, ctx->data_port); + + if (!is_port_open(ip, port)) { + *socket = nullptr; + return ""; + } + std::string send_endp = "tcp://" + ip + ":" + std::to_string(port); + *socket = new zmq::socket_t(*ctx->sock_context, zmq::socket_type::push); + try { + (*socket)->connect(send_endp); + } catch (const zmq::error_t& e) { + delete *socket; + *socket = nullptr; + return ""; + } + return ip; +} + void llama_init_sockets(struct llama_context * ctx, uint32_t n_world, uint32_t my_rank) { if (n_world == 1) { return; @@ -20618,12 +20692,11 @@ void llama_init_sockets(struct llama_context * ctx, uint32_t n_world, uint32_t m int llama_gather_device_info(struct llama_context * ctx, struct device_info * dev_info_set) { uint32_t n_world = ctx->cparams.n_world; - if (n_world == 1) { - return 0; - } - + if (n_world == 1) return 0; GGML_ASSERT(dev_info_set != nullptr); GGML_ASSERT(ctx != nullptr && ctx->send_socket != nullptr); + + // master rank sends its device info to its next rank try { char * buffer = nullptr; size_t buffer_size = serialize(&dev_info_set[0], &buffer); @@ -20638,6 +20711,7 @@ int llama_gather_device_info(struct llama_context * ctx, struct device_info * de return -1; } + // master rank receives aggregated device info from its previous rank (a barrier op) std::vector recv_msgs; if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(recv_msgs))) { return -1; @@ -20651,6 +20725,7 @@ int llama_gather_device_info(struct llama_context * ctx, struct device_info * de } int llama_send_device_info(struct llama_context * ctx, struct device_info * dev_info) { + // non-master ranks receive device info from their previous rank std::vector recv_msgs; if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(recv_msgs))) { return -1; @@ -20659,6 +20734,7 @@ int llama_send_device_info(struct llama_context * ctx, struct device_info * dev_ GGML_ASSERT(dev_info != nullptr); GGML_ASSERT(ctx != nullptr && ctx->send_socket != nullptr); + // non-master ranks aggregate their device info and forward it to their next rank try { char * buffer = nullptr; size_t buffer_size = serialize(dev_info, &buffer); @@ -20677,171 +20753,262 @@ int llama_send_device_info(struct llama_context * ctx, struct device_info * dev_ int llama_bcast_startup_args(llama_context * ctx, uint32_t rank, startup_args * args) { int32_t n_world = ctx->cparams.n_world; - GGML_ASSERT(n_world > 0); + if (n_world == 1) return 0; GGML_ASSERT(ctx != nullptr && ctx->send_socket != nullptr); + std::vector msgs; + if (rank == 0){ - // send + // master rank sends its startup args to its next rank try { - std::vector send_msgs; - - send_msgs.emplace_back("should_profile", strlen("should_profile")); - send_msgs.emplace_back(&args->should_profile, sizeof(args->should_profile)); - - send_msgs.emplace_back("n_ctx", strlen("n_ctx")); - send_msgs.emplace_back(&args->n_ctx, sizeof(args->n_ctx)); - - zmq::send_multipart(*ctx->send_socket, send_msgs); + msgs.emplace_back("should_profile", strlen("should_profile")); + msgs.emplace_back(&args->should_profile, sizeof(args->should_profile)); + msgs.emplace_back("n_ctx", strlen("n_ctx")); + msgs.emplace_back(&args->n_ctx, sizeof(args->n_ctx)); + zmq::send_multipart(*ctx->send_socket, msgs); } catch (const zmq::error_t& e) { LLAMA_LOG_INFO("Failed to send data: %s\n", e.what()); return -1; } - } else { - // receive - std::vector recv_msgs; - if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(recv_msgs))) { + + // master rank receives ack from its previous rank (a barrier op) + msgs.clear(); + if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) { return -1; } - GGML_ASSERT(recv_msgs[0].to_string() == "should_profile"); - GGML_ASSERT(recv_msgs[1].size() == sizeof(bool)); - bool should_profile = *static_cast(recv_msgs[1].data()); - args->should_profile = should_profile; + GGML_ASSERT(msgs.size() == 1); + GGML_ASSERT(msgs[0].to_string() == "ACK"); + } else { + // non-master ranks receive startup args from their previous rank + if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) { + return -1; + } - GGML_ASSERT(recv_msgs[2].to_string() == "n_ctx"); - GGML_ASSERT(recv_msgs[3].size() == sizeof(uint32_t)); - uint32_t n_ctx = *static_cast(recv_msgs[3].data()); - args->n_ctx = n_ctx; + GGML_ASSERT(msgs[0].to_string() == "should_profile"); + GGML_ASSERT(msgs[1].size() == sizeof(bool)); + GGML_ASSERT(msgs[2].to_string() == "n_ctx"); + GGML_ASSERT(msgs[3].size() == sizeof(uint32_t)); - if ((int)rank != (int)n_world - 1){ - // send + args->should_profile = *static_cast(msgs[1].data()); + args->n_ctx = *static_cast(msgs[3].data()); + + // non-master ranks forward the startup args to their next rank + try { + // last rank just sends an ack to pass the barrier + if ((int)rank == (int)n_world - 1) { + msgs.clear(); + msgs.emplace_back("ACK", strlen("ACK")); + } + zmq::send_multipart(*ctx->send_socket, msgs); + } catch (const zmq::error_t & e) { + LLAMA_LOG_INFO("Failed to send data: %s\n", e.what()); + return -1; + } + } + + return 0; +} + +int llama_bcast_layer_setup(struct llama_context * ctx, uint32_t * n_layer_window, uint32_t * n_gpu_layers) { + uint32_t n_world = ctx->cparams.n_world; + if (n_world == 1) return 0; + GGML_ASSERT(ctx != nullptr && ctx->send_socket != nullptr); + + std::vector msgs; + + // master rank sends its layer setup to its next rank + try { + msgs.emplace_back("n_layer_window", strlen("n_layer_window")); + msgs.emplace_back(n_layer_window, sizeof(uint32_t) * 32); + + if (n_gpu_layers != nullptr) { + msgs.emplace_back("n_gpu_layers", strlen("n_gpu_layers")); + msgs.emplace_back(n_gpu_layers, sizeof(uint32_t) * 32); + } + + zmq::send_multipart(*ctx->send_socket, msgs); + } catch (const zmq::error_t & e) { + LLAMA_LOG_INFO("Failed to send data: %s\n", e.what()); + return -1; + } + + // master rank receives ack from its previous rank (a barrier op) + msgs.clear(); + if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) { + return -1; + } + + GGML_ASSERT(msgs.size() == 1); + GGML_ASSERT(msgs[0].to_string() == "ACK"); + + return 0; +} + +int llama_rebuild_topo(llama_context * ctx, + uint32_t * n_layer_window, + device_info * dev_info_set, + NodeType * node_type, + char * is_forwarder) { + uint32_t n_world = ctx->cparams.n_world; + uint32_t my_rank = ctx->cparams.rank; + TopoRebuildHelperInfo* topo_helper = new TopoRebuildHelperInfo[n_world]; + + if (dev_info_set == nullptr) { + // for non-master ranks, receive devices info from their previous rank + std::vector msgs; + if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) { + return -1; + } + for (size_t i = 0; i < msgs.size(); i++) { + topo_helper[i].deserialize((char *)msgs[i].data()); + } + GGML_ASSERT(msgs.size() == n_world); + } else { + for (size_t i = 0; i < n_world; i++) { + topo_helper[i].dev_info = dev_info_set[i]; + topo_helper[i].is_forwarder = 0; + } + } + + GGML_ASSERT(ctx != nullptr && ctx->send_socket != nullptr); + + auto next_rank = (my_rank + 1) % n_world; + auto next_connect_rank = (my_rank + 1) % n_world; + zmq::socket_t * socket_to_close = nullptr; + bool is_not_exit = n_layer_window[my_rank] > 0 || topo_helper[my_rank].is_forwarder == 1; + + if (is_not_exit) { + // reconstruct socket to the next valid rank + auto current_rank = my_rank; + std::vector nodes; + auto next_rank_ = next_rank; + + while (next_rank_ != my_rank) { + nodes.push_back(next_rank_); + if (n_layer_window[next_rank_] > 0) { + break; + } + next_rank_ = (next_rank_ + 1) % n_world; + current_rank = (current_rank + 1) % n_world; + } + + if (next_rank_ == my_rank) { + // only one node + ctx->next_node_ip = ""; + socket_to_close = ctx->send_socket; + ctx->send_socket = nullptr; + } else { + // iterate node reverse + zmq::socket_t * socket = nullptr; + std::string ip; + for (int i = nodes.size() - 1; i > 0; --i) { + auto rank = nodes[i]; + ip = try_connect(ctx, rank, topo_helper, n_world, &socket); + if (!ip.empty()) { + next_connect_rank = rank; + break; + } + } + topo_helper[next_connect_rank].is_forwarder = 1; + if (next_connect_rank != next_rank) { + // reset socket + GGML_ASSERT(socket != nullptr); + GGML_ASSERT(!ip.empty()); + socket_to_close = ctx->send_socket; + ctx->send_socket = socket; + ctx->next_node_ip = ip; + ctx->cparams.original_next_rank = next_connect_rank; + } + } + } else if (n_layer_window[next_rank] <= 0 && topo_helper[next_rank].is_forwarder == 0) { + socket_to_close = ctx->send_socket; + } + + // notify next exiting node + if (socket_to_close != nullptr) { + GGML_ASSERT(n_layer_window[next_rank] <= 0 && topo_helper[next_rank].is_forwarder == 0); + try { + auto msgs = topohelper_to_messages(topo_helper, n_world); + socket_to_close->set(zmq::sockopt::linger, 3500); + zmq::send_multipart(*socket_to_close, msgs); + } catch (const zmq::error_t & e) { + LLAMA_LOG_INFO("Failed to send data: %s\n", e.what()); + return -1; + } + } + + // notify next connect node + if (!ctx->next_node_ip.empty() && is_not_exit) { + GGML_ASSERT(ctx->send_socket != nullptr); + try { + auto msgs = topohelper_to_messages(topo_helper, n_world); + zmq::send_multipart(*ctx->send_socket, msgs); + } catch (const zmq::error_t & e) { + LLAMA_LOG_INFO("Failed to send data: %s\n", e.what()); + return -1; + } + } + + if (n_layer_window[my_rank] > 0) { + *node_type = NodeType::NODE_TYPE_WORKER; + } else if (topo_helper[my_rank].is_forwarder == 1) { + *node_type = NodeType::NODE_TYPE_FORWARDER; + } else { + *node_type = NodeType::NODE_TYPE_EXIT; + } + + if (ctx->send_socket != nullptr && *node_type != NodeType::NODE_TYPE_EXIT) { + // recv the whole view of all nodes + std::vector msgs; + if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) { + return -1; + } + GGML_ASSERT(msgs.size() == n_world); + for (size_t i = 0; i < msgs.size(); i++) { + topo_helper[i].deserialize((char *)msgs[i].data()); + } + // broadcast the whole view + if (next_connect_rank != 0) { try { - zmq::send_multipart(*ctx->send_socket, recv_msgs); + zmq::send_multipart(*ctx->send_socket, msgs); } catch (const zmq::error_t & e) { LLAMA_LOG_INFO("Failed to send data: %s\n", e.what()); return -1; } } } - return 0; -} -int llama_bcast_layer_setup(struct llama_context * ctx, uint32_t * n_layer_window, uint32_t * n_gpu_layers) { - uint32_t n_world = ctx->cparams.n_world; - if (n_world == 1) { - return 0; + for (size_t i = 0; i < n_world; i++) { + is_forwarder[i] = topo_helper[i].is_forwarder; } - GGML_ASSERT(ctx != nullptr && ctx->send_socket != nullptr); - try { - std::vector send_msgs; + ctx->cparams.node_type = *node_type; - send_msgs.emplace_back("n_layer_window", strlen("n_layer_window")); - send_msgs.emplace_back(n_layer_window, sizeof(uint32_t) * 32); - - if (n_gpu_layers != nullptr) { - send_msgs.emplace_back("n_gpu_layers", strlen("n_gpu_layers")); - send_msgs.emplace_back(n_gpu_layers, sizeof(uint32_t) * 32); - } - - zmq::send_multipart(*ctx->send_socket, send_msgs); - } catch (const zmq::error_t& e) { - LLAMA_LOG_INFO("Failed to send data: %s\n", e.what()); - return -1; - } - - return 0; -} - -int llama_rebuild_topo(llama_context * ctx, uint32_t * n_layer_window, device_info * dev_info_set) { - uint32_t n_world = ctx->cparams.n_world; - uint32_t my_rank = ctx->cparams.rank; - device_info * dev_info_ptr = nullptr; - - if (dev_info_set == nullptr) { - std::vector msgs; - if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) { - return -1; - } - dev_info_ptr = new device_info[n_world]; - for (size_t i = 0; i < msgs.size(); i++) { - deserialize((const char *)msgs[i].data(), &dev_info_ptr[i]); - } - GGML_ASSERT(msgs.size() == n_world); - } else { - dev_info_ptr = dev_info_set; - } - - GGML_ASSERT(ctx != nullptr && ctx->send_socket != nullptr); - - // notify next rank - auto next_rank = (my_rank + 1) % n_world; - if (n_layer_window[next_rank] <= 0 && next_rank != 0) { - try { - auto msgs = dev_infos_to_messages(dev_info_ptr, n_world); - ctx->send_socket->set(zmq::sockopt::linger, 3500); - zmq::send_multipart(*ctx->send_socket, msgs); - } catch (const zmq::error_t& e) { - LLAMA_LOG_INFO("Failed to send data: %s\n", e.what()); - if(!dev_info_set){ - delete[] dev_info_ptr; - } - return -1; - } - } - - zmq::socket_t * socket_to_close = nullptr; - if (n_layer_window[my_rank] > 0) { - // reconstruct socket to the next valid rank - std::string next_ip; - auto current_rank = my_rank; - - while (next_rank != my_rank) { - if (n_layer_window[next_rank] > 0) { - next_ip = dev_info_ptr[current_rank].next_ip; - break; - } - next_rank = (next_rank + 1) % n_world; - current_rank = (current_rank + 1) % n_world; - } - - if (!next_ip.empty()) { - if ((my_rank + 1) % n_world != next_rank) { - socket_to_close = ctx->send_socket; - ctx->send_socket = new zmq::socket_t(*ctx->sock_context, zmq::socket_type::push); - std::string send_endp = "tcp://" + next_ip + ":" + std::to_string(ctx->next_node_data_port); - ctx->send_socket->connect(send_endp); - ctx->next_node_ip = next_ip; - ctx->cparams.original_next_rank = next_rank; - } - - if (next_rank != 0) { - try { - auto msgs = dev_infos_to_messages(dev_info_ptr, n_world); - zmq::send_multipart(*ctx->send_socket, msgs); - } catch (const zmq::error_t &e) { - LLAMA_LOG_INFO("Error binding/connecting recv socket to endpoint: %s", e.what()); - if(!dev_info_set){ - delete[] dev_info_ptr; - } - return -1; - } - } - } else { - // only one node - ctx->next_node_ip = ""; - } - } - - if (!dev_info_set) { - delete[] dev_info_ptr; - } - - if(socket_to_close != nullptr){ + if (socket_to_close != nullptr) { socket_to_close->close(); delete socket_to_close; } + delete [] topo_helper; + return 0; +} +int llama_forward_messages(llama_context *ctx) { + zmq::message_t message; + int more = true; + int timeout_ms = 10; + ctx->recv_socket->setsockopt(ZMQ_RCVTIMEO, &timeout_ms, sizeof(timeout_ms)); + while (more) { + auto recv_result = ctx->recv_socket->recv(message, zmq::recv_flags::none); + if (!recv_result) { + return -1; + } + size_t more_size = sizeof(more); + ctx->recv_socket->getsockopt(ZMQ_RCVMORE, &more, &more_size); + + ctx->send_socket->send(message, + more ? zmq::send_flags::sndmore : zmq::send_flags::none); + } return 0; } @@ -20849,28 +21016,33 @@ int llama_recv_layer_setup(struct llama_context * ctx, uint32_t * n_layer_window uint32_t n_world = ctx->cparams.n_world; uint32_t my_rank = ctx->cparams.rank; - std::vector recv_msgs; - if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(recv_msgs))) { + // non-master ranks receive data from their previous rank + std::vector msgs; + if (!zmq::recv_multipart(*ctx->recv_socket, std::back_inserter(msgs))) { return -1; } - GGML_ASSERT(recv_msgs[0].to_string() == "n_layer_window"); - GGML_ASSERT(recv_msgs[1].size() == sizeof(uint32_t) * 32); - memcpy(n_layer_window, recv_msgs[1].data(), sizeof(uint32_t) * 32); + GGML_ASSERT(msgs[0].to_string() == "n_layer_window"); + GGML_ASSERT(msgs[1].size() == sizeof(uint32_t) * 32); + memcpy(n_layer_window, msgs[1].data(), sizeof(uint32_t) * 32); - if (recv_msgs.size() > 2) { - GGML_ASSERT(recv_msgs[2].to_string() == "n_gpu_layers"); - GGML_ASSERT(recv_msgs[3].size() == sizeof(uint32_t) * 32); - memcpy(n_gpu_layers, recv_msgs[3].data(), sizeof(uint32_t) * 32); + if (msgs.size() > 2) { + GGML_ASSERT(msgs[2].to_string() == "n_gpu_layers"); + GGML_ASSERT(msgs[3].size() == sizeof(uint32_t) * 32); + memcpy(n_gpu_layers, msgs[3].data(), sizeof(uint32_t) * 32); } - if (my_rank != n_world - 1) { - try { - zmq::send_multipart(*ctx->send_socket, recv_msgs); - } catch (const zmq::error_t& e) { - LLAMA_LOG_INFO("Failed to send data: %s\n", e.what()); - return -1; + // non-master ranks forward the received message to their next rank + try { + if (my_rank == n_world - 1) { + // last rank just sends an ack to pass the barrier + msgs.clear(); + msgs.emplace_back("ACK", strlen("ACK")); } + zmq::send_multipart(*ctx->send_socket, msgs); + } catch (const zmq::error_t& e) { + LLAMA_LOG_INFO("Failed to send data: %s\n", e.what()); + return -1; } return 0; @@ -20880,9 +21052,7 @@ 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; - if (n_world == 1) { - return; - } + 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(ctx->next_node_signal_port); @@ -20904,10 +21074,16 @@ void llama_free_sockets(struct llama_context * ctx, char ** msg) { } } -void llama_update_context_with_rankworld(struct llama_context * ctx, uint32_t rank, uint32_t n_world) { +void llama_update_context_with_rankworld(struct llama_context * ctx, + uint32_t rank, + uint32_t n_world, + uint32_t worker_rank, + uint32_t n_worker) { if (ctx) { ctx->cparams.rank = rank; ctx->cparams.n_world = n_world; + ctx->cparams.worker_rank = worker_rank; + ctx->cparams.n_worker = n_worker; } } @@ -20929,13 +21105,18 @@ struct llama_context * llama_new_context_with_model( ctx->master_data_port = params.master_data_port; ctx->data_port = params.data_port; ctx->signal_port = params.signal_port; - ctx->cparams.n_world = params.n_world; ctx->cparams.rank = params.rank; ctx->cparams.force = params.force; ctx->cparams.dump_folder = params.dump_folder; ctx->cparams.enable_comm_compute_log = params.enable_comm_compute_log; ctx->cparams.original_next_rank = (params.rank + 1) % params.n_world; + + auto &hparams = model->hparams; + auto &cparams = ctx->cparams; + cparams.n_ctx = params.n_ctx == 0 ? hparams.n_ctx_train : params.n_ctx; + ctx->logits_all = params.logits_all; + return ctx; } @@ -20978,7 +21159,7 @@ void * llama_context_setup_backend( auto & cparams = ctx->cparams; std::copy(std::begin(params.n_layer_window), std::end(params.n_layer_window), cparams.n_layer_window); - cparams.prefetch = params.prefetch; + cparams.prefetch = params.prefetch; 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; @@ -21352,14 +21533,42 @@ void * llama_context_setup_backend( std::vector gf = llama_build_graph(*ctx, ubatch, true); GGML_ASSERT(gf.size() <= MAX_SCHEDULERS && "Number of subgraphs exceeds the maximum number of schedulers\n"); + for (size_t i = gf.size(); i < ctx->sched.size(); ++i) { + ggml_backend_sched_free(ctx->sched[i]); + } ctx->sched.resize(gf.size()); // initialize scheduler with the worst-case graph bool ok = true; GGML_ASSERT(ctx->sched.size() == gf.size()); for (size_t i = 0; i < gf.size(); ++i) { + +#if defined(GGML_USE_CUDA) + // output layer + if (!params.keep_out_in_cuda && cparams.rank == 0 && i == gf.size() - 1) { + continue; + } + + // input layer + if (cparams.rank == 0 && i == 0) { + continue; + } + + // ignore all backend layers if n_gpu_layers is 0 + if (model->n_gpu_layers == 0) { + continue; + } + + // don't reserve for repeated backend layers + if ((cparams.rank == 0 && i > 1 && i < gf.size() - 1) + || (cparams.rank > 0 && i > 0)) { + continue; + } +#endif + 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); @@ -21890,11 +22099,13 @@ void llama_model_compute_buf_size( int64_t * gpu_buf, const struct llama_model * model, const struct llama_context_params cparams, - bool use_gpu, - bool is_master, - int n_layers, - int n_gpu_layers) { + enum backend_type backend, + int my_rank, + struct model_bytes n_bytes, + bool offload, + bool has_gpu_layers) { const llama_hparams hparams = model->hparams; + bool is_master = my_rank == 0; // input tensors const int64_t n_inp_toks = cparams.n_ubatch; @@ -21904,33 +22115,150 @@ void llama_model_compute_buf_size( const int64_t n_bak_embd = hparams.n_embd * cparams.n_ubatch; const int64_t n_inp_pos = cparams.n_ubatch; const int64_t n_kq_mask = cparams.n_ctx * cparams.n_ubatch; - const int64_t n_inp_out_ids = cparams.n_ubatch; const int64_t n_norm = hparams.n_embd * cparams.n_ubatch; - const int64_t n_qcur = hparams.n_embd * cparams.n_ubatch * 2; + const int64_t n_qcur = hparams.n_embd * cparams.n_ubatch; + const int64_t n_ffn_gate = hparams.n_ff() * cparams.n_ubatch; + const int64_t n_ffn_up = hparams.n_ff() * cparams.n_ubatch; + const int64_t n_ffn_out = hparams.n_embd * cparams.n_ubatch; + const int64_t n_ffn_inp = hparams.n_embd * cparams.n_ubatch; const int64_t n_kq = cparams.n_ctx * cparams.n_ubatch * hparams.n_head(); + const int64_t n_k = cparams.n_ctx * hparams.n_embd_head_k * hparams.n_head_kv(); + const int64_t n_v = cparams.n_ctx * hparams.n_embd_head_v * hparams.n_head_kv(); + const int64_t n_inp_out_ids = cparams.n_ubatch; // outputs const int64_t n_out_embd = hparams.n_embd * cparams.n_ubatch; - const int64_t n_output = hparams.n_vocab * cparams.n_ubatch; + const int64_t n_result = hparams.n_vocab * cparams.n_ubatch; - // compute buffer size for input, each layer, and output - const int64_t n_buf_inp = (n_inp_toks + n_inp_embd) * ggml_type_size(GGML_TYPE_F32); - const int64_t n_buf_act = (n_bak_embd + n_inp_pos + n_kq_mask + - n_inp_out_ids + n_norm + n_qcur + n_kq - ) * ggml_type_size(GGML_TYPE_F32); - const int64_t n_buf_out = (n_out_embd + n_output) * ggml_type_size(GGML_TYPE_F32); + // weights + const int64_t nb_attn_norm_w = n_bytes.nb_attn_norm_w; + const int64_t nb_attn_q_w = n_bytes.nb_attn_q_w; + const int64_t nb_output_w = n_bytes.nb_output_w; + + // format bytes + const int64_t type_size_f32 = ggml_type_size(GGML_TYPE_F32); + const int64_t type_size_f16 = ggml_type_size(GGML_TYPE_F16); - *cpu_buf = 0; + bool use_gpu = backend != BACKEND_CPU && has_gpu_layers; *gpu_buf = 0; - if (is_master) *cpu_buf = n_buf_inp + n_buf_out; + *cpu_buf = 0; + int64_t gpu_host_buf = 0; - if (use_gpu) { - *gpu_buf += n_buf_act; - if (n_layers > n_gpu_layers) { - *cpu_buf += n_buf_act; + // GPU compute buffer + // estimate the GPU compute buffer, here we can only estimate the upper bound of various models, + // but cannot estimate the exact value. + if (backend == BACKEND_CUDA) { + *gpu_buf = (n_bak_embd + n_norm) * type_size_f32; + + if (is_master) { + if (has_gpu_layers) { + if (offload) { + *gpu_buf += std::max({ + (n_qcur + n_inp_pos + n_kq_mask + n_inp_out_ids) * type_size_f32 + nb_attn_norm_w, + (n_qcur + n_inp_pos + n_norm) * type_size_f32 + nb_attn_norm_w, + (n_qcur + n_qcur + n_kq_mask + n_inp_pos) * type_size_f32, + (n_qcur + n_qcur + n_inp_pos) * type_size_f32 + nb_attn_q_w, + n_inp_pos * type_size_f32 + (n_k + n_v) * type_size_f16 + nb_attn_q_w + }); + } else { + *gpu_buf += (n_qcur + n_inp_pos + n_kq_mask + n_inp_out_ids) * type_size_f32; + } + *gpu_buf += (n_qcur + n_kq) * type_size_f32; + } else { + *gpu_buf += (n_qcur + n_kq) * type_size_f32; + *gpu_buf += std::max({ + (n_kq_mask + n_qcur + n_inp_pos) * type_size_f32 + nb_attn_norm_w, + (n_inp_pos + n_kq_mask) * type_size_f32 + n_v * type_size_f16 + nb_attn_norm_w, + }); + } + // we run the output layer on CPU by default + if (cparams.keep_out_in_cuda) { + *gpu_buf += (n_out_embd + n_result) * type_size_f32 + nb_output_w; + } + gpu_host_buf = (n_inp_toks + n_inp_embd + n_bak_embd + n_inp_pos + n_kq_mask + n_inp_out_ids + n_out_embd) * type_size_f32; + } else { + if (has_gpu_layers) { + if (offload) { + *gpu_buf += (n_kq + n_qcur) * type_size_f32; + *gpu_buf += std::max({ + (n_inp_pos + n_norm + n_kq_mask) * type_size_f32 + nb_attn_norm_w, + (n_inp_pos + n_norm + n_qcur) * type_size_f32 + nb_attn_norm_w, + n_inp_pos * type_size_f32 + (n_k + n_v) * type_size_f16 + nb_attn_q_w, + }); + } else { + *gpu_buf += (n_inp_pos + n_kq_mask + n_qcur + n_qcur + n_kq) * type_size_f32; + } + } else { + *gpu_buf += (n_qcur + n_kq + n_kq_mask + n_qcur + n_inp_pos) * type_size_f32 + nb_attn_norm_w; + } + gpu_host_buf = (n_bak_embd + n_inp_pos + n_kq_mask) * type_size_f32; } - } else { - *cpu_buf += n_buf_act; + } + + else if (backend == BACKEND_METAL) { + const int64_t nb_act_buf_base = (n_inp_pos + n_kq_mask) * type_size_f32; + *gpu_buf = nb_act_buf_base; + *cpu_buf = nb_act_buf_base; + + if (use_gpu) { + if (is_master) { + *cpu_buf += (n_inp_toks + n_inp_embd + n_bak_embd + n_inp_out_ids + n_out_embd + n_result) * type_size_f32; + + if (offload) { + *gpu_buf += (n_ffn_out + n_ffn_inp + n_inp_out_ids) * type_size_f32; + *gpu_buf += std::max(n_ffn_up + n_ffn_gate, n_qcur + n_qcur + n_kq) * type_size_f32; + *cpu_buf += n_norm * type_size_f32; + *cpu_buf += std::max(n_ffn_up + n_ffn_gate, n_qcur + n_qcur + n_kq) * type_size_f32; + } else { + *gpu_buf += (n_bak_embd + n_inp_out_ids + n_norm) * type_size_f32; + *gpu_buf += std::max(n_ffn_up + n_ffn_gate, n_qcur + n_qcur + n_kq) * type_size_f32; + } + } else { + *gpu_buf += (n_ffn_out + n_ffn_inp) * type_size_f32; + *gpu_buf += std::max(n_ffn_up + n_ffn_gate, n_qcur + n_qcur + n_kq) * type_size_f32; + + *cpu_buf += n_bak_embd * type_size_f32; + if (offload) { + *cpu_buf += n_norm * type_size_f32; + *cpu_buf += std::max(n_ffn_up + n_ffn_gate, n_qcur + n_qcur + n_kq) * type_size_f32; + } + } + } else { + *gpu_buf = 0; + *cpu_buf = 0; + } + } + + else if (backend != BACKEND_CPU) { + GGML_ASSERT(false && "Unsupported backend type for compute buffer estimation.\n"); + } + + // CPU compute buffer for NUMA system or Metal with ngl=0 + if (*cpu_buf == 0) { + if (is_master) { + *cpu_buf += (n_inp_toks + n_inp_embd + n_inp_out_ids + n_out_embd + n_result) * type_size_f32; + } + if (offload) { + *cpu_buf = (n_inp_pos + n_kq_mask + n_bak_embd + n_norm) * type_size_f32; + *cpu_buf += std::max(n_ffn_gate + n_ffn_up, n_qcur + n_qcur + n_kq) * type_size_f32; + } + *cpu_buf += gpu_host_buf; + } + + LLAMA_LOG_INFO("\n"); + LLAMA_LOG_INFO("%s: here the compute buffer size is a predicted upper bound, not an exact value\n", __func__); + LLAMA_LOG_INFO("%s: (rank %d) compute buffer size = %7.2f MiB (GPU) + %7.2f MiB (CPU & GPU-host buffer)\n", __func__, + my_rank, *gpu_buf / (1024.0 * 1024.0), *cpu_buf / (1024.0 * 1024.0)); + + if (backend == BACKEND_CUDA) { + // context GPU memory usage, i.e. the initial memory cost of creating a CUDA context, + // even before you launch any kernels or allocate your own buffers. + // this value may vary by GPU and CUDA version, but it's lower than 400 MiB in most cases, + // another 300 MiB is used to prevent accidental OOM. + *gpu_buf += 700 * 1024 * 1024; + } else if (backend == BACKEND_METAL) { + // 300 MiB is used to prevent accidental OOM, e.g., automatic quantization conversion. + *gpu_buf += 300 * 1024 * 1024; } } @@ -22058,7 +22386,7 @@ void llama_model_n_flops( llm_load_llama_tensors(*ml, *model, ctx_map, 1, 0, n_layer_window, &use_mmap_buffer, false); break; case LLM_ARCH_QWEN2: - llm_load_qwen2_tensors(*ml, *model, ctx_map, 1, 0, n_layer_window, &use_mmap_buffer, false); + llm_load_qwen2_tensors(*ml, *model, ctx_map, 1, 0, n_layer_window, false); break; default: throw std::runtime_error("unsupported architecture\n"); @@ -22077,9 +22405,11 @@ void llama_model_n_flops( if (blk_suffix == "attn_norm.weight" || blk_suffix == "ffn_norm.weight") { count_n_flops (n_flops, GGML_TYPE_F32, PROFILER_LAYER_BACKEND, 4 * n_embd + 1); // rms norm count_n_flops (n_flops, cur->type, PROFILER_LAYER_BACKEND, n_embd); // norm weights + n_bytes->nb_attn_norm_w = std::max(n_bytes->nb_attn_norm_w, (int64_t)ggml_nbytes(cur)); } else if (blk_suffix == "attn_q.weight") { count_n_flops (n_flops, cur->type, PROFILER_LAYER_BACKEND, 2 * n_embd * n_embd); count_n_flops (n_flops, GGML_TYPE_F32, PROFILER_LAYER_BACKEND, 2.5 * n_embd); // rope + n_bytes->nb_attn_q_w = std::max(n_bytes->nb_attn_q_w, (int64_t)ggml_nbytes(cur)); } else if (blk_suffix == "attn_k.weight") { count_n_flops (n_flops, cur->type, PROFILER_LAYER_BACKEND, 2 * n_embd * n_head_kv * n_embd_head_k); count_n_flops (n_flops, GGML_TYPE_F32, PROFILER_LAYER_BACKEND, 2.5 * n_embd_head_k * n_head_kv); // rope @@ -22094,9 +22424,11 @@ void llama_model_n_flops( } else if (blk_suffix == "ffn_gate.weight") { count_n_flops (n_flops, cur->type, PROFILER_LAYER_BACKEND, 2 * n_embd * n_ff); count_n_flops (n_flops, GGML_TYPE_F32, PROFILER_LAYER_BACKEND, 8 * n_ff); // SiLU + n_bytes->nb_ffn_gate_w = std::max(n_bytes->nb_ffn_gate_w, (int64_t)ggml_nbytes(cur)); } else if (blk_suffix == "ffn_down.weight") { count_n_flops (n_flops, cur->type, PROFILER_LAYER_BACKEND, 2 * n_embd * n_ff); count_n_flops (n_flops, GGML_TYPE_F32, PROFILER_LAYER_BACKEND, n_embd); // shortcut + n_bytes->nb_ffn_down_w = std::max(n_bytes->nb_ffn_down_w, (int64_t)ggml_nbytes(cur)); } else if (blk_suffix == "ffn_up.weight") { count_n_flops (n_flops, cur->type, PROFILER_LAYER_BACKEND, 2 * n_embd * n_ff); count_n_flops (n_flops, GGML_TYPE_F32, PROFILER_LAYER_BACKEND, n_ff); // silu(gate(x)) * up(x) @@ -22129,6 +22461,7 @@ void llama_model_n_flops( count_n_flops (n_flops, GGML_TYPE_F32, PROFILER_LAYER_OUTPUT, 5 * n_vocab); // softmax count_n_params(n_params, cur->type, PROFILER_LAYER_OUTPUT, ggml_nelements(cur)); count_n_bytes (n_bytes, PROFILER_LAYER_OUTPUT, ggml_nbytes(cur)); + n_bytes->nb_output_w = std::max(n_bytes->nb_output_w, (int64_t)ggml_nbytes(cur)); } else if (tensor_name == "rope_freqs.weight") { if (!rope_used) { count_n_params(n_params, cur->type, PROFILER_LAYER_BACKEND, ggml_nelements(cur)); @@ -22506,6 +22839,21 @@ void llama_kv_cache_seq_keep(struct llama_context * ctx, llama_seq_id seq_id) { llama_kv_cache_seq_keep(ctx->kv_self, seq_id); } +void llama_send_kv_cache_seq_keep(struct llama_context * ctx, llama_seq_id seq_id) { + if (ctx->send_socket == nullptr) { + return; + } + + try { + std::vector msgs; + msgs.emplace_back("kv_seq_keep", strlen("kv_seq_keep")); + msgs.emplace_back(&seq_id, sizeof(seq_id)); + zmq::send_multipart(*ctx->send_socket, msgs); + } catch (const zmq::error_t & e) { + LLAMA_LOG_WARN("Failed to send kv_seq_keep: %s\n", e.what()); + } +} + void llama_kv_cache_seq_add(struct llama_context * ctx, llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos delta) { if (delta == 0) { return; @@ -23643,8 +23991,9 @@ int32_t llama_encode( int32_t llama_decode( struct llama_context * ctx, - struct llama_batch batch) { - return llama_decode_internal(*ctx, batch); + struct llama_batch batch, + bool server_mode) { + return llama_decode_internal(*ctx, batch, server_mode); } void llama_synchronize(struct llama_context * ctx) { diff --git a/src/network-utils.cpp b/src/network-utils.cpp new file mode 100644 index 00000000..e7fa5ab1 --- /dev/null +++ b/src/network-utils.cpp @@ -0,0 +1,26 @@ +#include "network-utils.h" + +#include +#include +#include +#include + +bool is_port_open(const std::string& ip, uint32_t port, int timeout_sec) { + int sock = socket(AF_INET, SOCK_STREAM, 0); + if (sock < 0) return false; + + struct timeval tv; + tv.tv_sec = timeout_sec; + tv.tv_usec = 0; + setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)); + setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv)); + + struct sockaddr_in server; + server.sin_addr.s_addr = inet_addr(ip.c_str()); + server.sin_family = AF_INET; + server.sin_port = htons(port); + + int res = connect(sock, (struct sockaddr*)&server, sizeof(server)); + close(sock); + return res == 0; +} \ No newline at end of file diff --git a/src/network-utils.h b/src/network-utils.h new file mode 100644 index 00000000..7a35475a --- /dev/null +++ b/src/network-utils.h @@ -0,0 +1,7 @@ +#pragma once + +#include + +typedef unsigned int uint32_t; + +bool is_port_open(const std::string& ip, uint32_t port, int timeout_sec = 2); \ No newline at end of file