electron-rare/prima-cpp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'upstream/main' into merge/merge_86ca21e_from_upstream

Browse Source
This commit is contained in:
DandinPower
2025-08-07 20:19:18 +08:00
parent 96191d3b6f bdf9d8e74b
commit 7480ffbd30
17 changed files with 1632 additions and 429 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Makefile
+18 -1
View File
@@ -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
README.md
+43 -33
View File
@@ -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.
common/arg.cpp
+53 -8
View File
@@ -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),
common/common.cpp
+118 -36
View File
@@ -28,6 +28,7 @@
#include <unordered_set>
#include <vector>
#include <thread>
#include <atomic>
#if defined(__APPLE__) && defined(__MACH__)
#include <sys/types.h>
@@ -846,6 +847,16 @@ static std::string vec_to_str(const std::vector<T> & 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<float> 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<double>(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<uint32_t>(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<uint32_t>(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<float>::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<float> 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<device_info> 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<uint32_t> 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<bool> 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);
common/common.h
+21 -4
View File
@@ -33,6 +33,8 @@ struct llama_lora_adapter_container : llama_lora_adapter_info {
struct llama_lora_adapter * adapter;
};
using llama_tokens = std::vector<llama_token>;
// 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
common/profiler.cpp
+43 -5
View File
@@ -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<double>(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
}
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;
}
common/profiler.h
+30 -4
View File
@@ -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
common/sampling.cpp
+39
View File
@@ -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<llama_token> gpt_sampler_sample_and_accept_n(struct gpt_sampler * gsmpl, struct llama_context * ctx, const std::vector<int> & idxs, const llama_tokens & draft, bool grammar_first) {
GGML_ASSERT(idxs.size() == draft.size() + 1 && "idxs.size() must be draft.size() + 1");
std::vector<llama_token> 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<llama_token> gpt_sampler_sample_and_accept_n(struct gpt_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first) {
std::vector<int> 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);
}
common/sampling.h
+21
View File
@@ -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<llama_token> gpt_sampler_sample_and_accept_n(struct gpt_sampler * gsmpl, struct llama_context * ctx, const std::vector<int> & idxs, const llama_tokens & draft, bool grammar_first = false);
// assume idxs == [ 0, 1, 2, ..., draft.size() ]
std::vector<llama_token> 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
common/speculative.cpp
+271
View File
@@ -0,0 +1,271 @@
#include "speculative.h"
#include "log.h"
#include "common.h"
#include "sampling.h"
#include <cstring>
#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<int>(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;
}
common/speculative.h
+28
View File
@@ -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);
examples/server/server.cpp
+222 -40
View File
@@ -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<std::string> 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 <think> 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));
examples/speculative/speculative.cpp
+41 -29
View File
@@ -12,7 +12,7 @@
#include <string>
#include <vector>
#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<seq_draft> 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);
include/llama.h
+40 -7
View File
@@ -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)
src/llama.cpp
+611 -262
View File
File diff suppressed because it is too large Load Diff
src/network-utils.cpp
+26
View File
@@ -0,0 +1,26 @@
#include "network-utils.h"
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
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;
}
src/network-utils.h
+7
View File
@@ -0,0 +1,7 @@
#pragma once
#include <string>
typedef unsigned int uint32_t;
bool is_port_open(const std::string& ip, uint32_t port, int timeout_sec = 2);