From 15771072c7aec5b7121f45037509653f3f7ad2b3 Mon Sep 17 00:00:00 2001 From: firecoperana Date: Sun, 30 Nov 2025 11:48:02 -0600 Subject: [PATCH] RPC: support multiple devices including cpu (#1024) * RPC support multiple devices * rpc : update documentation (#16441) Update the README file to match the newly added functionality of exposing multiple devices from a single server. Co-authored-by: Diego Devesa # Conflicts: # examples/rpc/README.md * Remove memory settings * rpc : cache and reuse compute graphs (#15405) Store the last computed graph and reuse it when possible. Also do not return response from GRAPH_COMPUTE and assume it always completes successfully. If this this is not the case, the server closes the connection. This saves us a network round trip to the server. * Add -cpu to include cpu backend --------- Co-authored-by: firecoperana Co-authored-by: Radoslav Gerganov --- common/common.cpp | 39 ++- examples/rpc/README.md | 68 ++-- examples/rpc/rpc-server.cpp | 208 ++++++++--- ggml/include/ggml-rpc.h | 16 +- ggml/src/ggml-backend.cpp | 11 +- ggml/src/ggml-rpc.cpp | 669 ++++++++++++++++++++++++------------ src/llama-model.h | 7 +- src/llama.cpp | 97 +++--- 8 files changed, 734 insertions(+), 381 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index d173c770..294a88e5 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -270,6 +270,28 @@ static std::string parse_device_list(const std::string& value) { return value; } +static std::string add_rpc_devices(std::string& servers) { + std::string rpc_devices; + std::vector rpc_servers = string_split(servers, ","); + if (rpc_servers.empty()) { + throw std::invalid_argument("no RPC servers specified"); + } + for (auto& server : rpc_servers) { + uint32_t dev_count = ggml_backend_rpc_get_device_count(server.c_str()); + uint32_t device = 0; + for (uint32_t i = 0; i < dev_count; ++i) { + const auto buft = ggml_backend_rpc_buffer_type(server.c_str(), device); + if (buft != nullptr) { + rpc_devices = rpc_devices + server + "|" + std::to_string(device) + ","; + ++device; + } + } + } + if (!rpc_devices.empty()) { + rpc_devices = rpc_devices.substr(0, rpc_devices.size() - 1); // remove trailing comma + } + return rpc_devices; +} std::pair> common_remote_get_content(const std::string& url, const common_remote_params&) { if (!url.empty()) { @@ -1296,15 +1318,12 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa if (arg == "--rpc") { CHECK_ARG #ifdef GGML_USE_RPC - params.rpc_servers = argv[i]; - std::string servers(params.rpc_servers); - size_t pos = 0; - while ((pos = servers.find(",")) != std::string::npos) { - std::string server = servers.substr(0, pos); - ggml_backend_rpc_buffer_type(server.c_str()); - servers.erase(0, pos + 1); + std::string servers(argv[i]); + servers = add_rpc_devices(servers); + if (servers.empty()) { + return false; } - ggml_backend_rpc_buffer_type(servers.c_str()); + params.rpc_servers = servers; #endif return true; } @@ -1319,10 +1338,6 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa } if (arg == "--override-tensor" || arg == "-ot") { CHECK_ARG - /*for (auto endpoint : params.rpc_servers.split) - { - - }*/ if (!parse_buft_overrides(std::string{ argv[i] }, params.tensor_buft_overrides)) { fprintf(stderr, "error: Invalid tensor buffer type override: %s\n", argv[i]); invalid_param = true; diff --git a/examples/rpc/README.md b/examples/rpc/README.md index adedc890..3b29d106 100644 --- a/examples/rpc/README.md +++ b/examples/rpc/README.md @@ -4,7 +4,7 @@ > This example and the RPC backend are currently in a proof-of-concept development stage. As such, the functionality is fragile and > insecure. **Never run the RPC server on an open network or in a sensitive environment!** -The `rpc-server` allows running `ggml` backend on a remote host. +The `rpc-server` allows exposing `ggml` devices on a remote host. The RPC backend communicates with one or several instances of `rpc-server` and offloads computations to them. This can be used for distributed LLM inference with `llama.cpp` in the following way: @@ -14,27 +14,34 @@ flowchart TD rpcb---|TCP|srvb rpcb-.-|TCP|srvn subgraph hostn[Host N] - srvn[rpc-server]-.-backend3["Backend (CUDA,Metal,etc.)"] + srvn[rpc-server]<-.->dev4["CUDA0"] + srvn[rpc-server]<-.->dev5["CPU"] end subgraph hostb[Host B] - srvb[rpc-server]---backend2["Backend (CUDA,Metal,etc.)"] + srvb[rpc-server]<-->dev3["Metal"] end subgraph hosta[Host A] - srva[rpc-server]---backend["Backend (CUDA,Metal,etc.)"] + srva[rpc-server]<-->dev["CUDA0"] + srva[rpc-server]<-->dev2["CUDA1"] end subgraph host[Main Host] - ggml[llama.cpp]---rpcb[RPC backend] + local["Local devices"]<-->ggml[llama-cli] + ggml[llama-cli]<-->rpcb[RPC backend] end style hostn stroke:#66,stroke-width:2px,stroke-dasharray: 5 5 + classDef devcls fill:#5B9BD5 + class local,dev,dev2,dev3,dev4,dev5 devcls ``` -Each host can run a different backend, e.g. one with CUDA and another with Metal. -You can also run multiple `rpc-server` instances on the same host, each with a different backend. +By default, `rpc-server` exposes all available accelerator devices on the host. +If there are no accelerators, it exposes a single `CPU` device. ## Usage -On each host, build the corresponding backend with `cmake` and add `-DGGML_RPC=ON` to the build options. -For example, to build the CUDA backend with RPC support: +### Remote hosts + +On each remote host, build the backends for each accelerator by adding `-DGGML_RPC=ON` to the build options. +For example, to build the `rpc-server` with support for CUDA accelerators: ```bash mkdir build-rpc-cuda @@ -43,36 +50,49 @@ cmake .. -DGGML_CUDA=ON -DGGML_RPC=ON cmake --build . --config Release ``` -Then, start the `rpc-server` with the backend: +When started, the `rpc-server` will detect and expose all available `CUDA` devices: ```bash -$ bin/rpc-server -p 50052 -create_backend: using CUDA backend -ggml_cuda_init: GGML_CUDA_FORCE_MMQ: no -ggml_cuda_init: CUDA_USE_TENSOR_CORES: yes +$ bin/rpc-server +ggml_cuda_init: GGML_CUDA_FORCE_MMQ: no +ggml_cuda_init: GGML_CUDA_FORCE_CUBLAS: no ggml_cuda_init: found 1 CUDA devices: - Device 0: NVIDIA T1200 Laptop GPU, compute capability 7.5, VMM: yes -Starting RPC server on 0.0.0.0:50052 + Device 0: NVIDIA GeForce RTX 5090, compute capability 12.0, VMM: yes +Starting RPC server v3.0.0 + endpoint : 127.0.0.1:50052 + local cache : n/a +Devices: + CUDA0: NVIDIA GeForce RTX 5090 (32109 MiB, 31588 MiB free) ``` -When using the CUDA backend, you can specify the device with the `CUDA_VISIBLE_DEVICES` environment variable, e.g.: +You can control the set of exposed CUDA devices with the `CUDA_VISIBLE_DEVICES` environment variable or the `--device` command line option. The following two commands have the same effect: ```bash $ CUDA_VISIBLE_DEVICES=0 bin/rpc-server -p 50052 +$ bin/rpc-server --device CUDA0 -p 50052 ``` -This way you can run multiple `rpc-server` instances on the same host, each with a different CUDA device. +### Main host -On the main host build `llama.cpp` only with `-DGGML_RPC=ON`: +On the main host build `llama.cpp` with the backends for the local devices and add `-DGGML_RPC=ON` to the build options. +Finally, when running `llama-cli` or `llama-server`, use the `--rpc` option to specify the host and port of each `rpc-server`: ```bash -mkdir build-rpc -cd build-rpc -cmake .. -DGGML_RPC=ON -cmake --build . --config Release +$ llama-cli -hf ggml-org/gemma-3-1b-it-GGUF -ngl 99 --rpc 192.168.88.10:50052,192.168.88.11:50052 ``` -Finally, use the `--rpc` option to specify the host and port of each `rpc-server`: +By default, llama.cpp distributes model weights and the KV cache across all available devices -- both local and remote -- in proportion to each device's available memory. +You can override this behavior with the `--tensor-split` option and set custom proportions when splitting tensor data across devices. ```bash $ bin/llama-cli -m ../models/tinyllama-1b/ggml-model-f16.gguf -p "Hello, my name is" --repeat-penalty 1.0 -n 64 --rpc 192.168.88.10:50052,192.168.88.11:50052 -ngl 99 ``` + +By default, the cache is stored in the `$HOME/.cache/llama.cpp/rpc` directory and can be controlled via the `LLAMA_CACHE` environment variable. + +### Troubleshooting + +Use the `GGML_RPC_DEBUG` environment variable to enable debug messages from `rpc-server`: +```bash +$ GGML_RPC_DEBUG=1 bin/rpc-server +``` + diff --git a/examples/rpc/rpc-server.cpp b/examples/rpc/rpc-server.cpp index 5976888b..fc134231 100644 --- a/examples/rpc/rpc-server.cpp +++ b/examples/rpc/rpc-server.cpp @@ -32,6 +32,7 @@ #include #include #include +#include namespace fs = std::filesystem; @@ -145,22 +146,24 @@ static std::string fs_get_cache_directory() { } struct rpc_server_params { - std::string host = "127.0.0.1"; - int port = 50052; - size_t backend_mem = 0; - bool use_cache = false; - int n_threads = std::max(1U, std::thread::hardware_concurrency() / 2); + std::string host = "127.0.0.1"; + int port = 50052; + bool use_cache = false; + bool use_cpu = false; + int n_threads = std::max(1U, std::thread::hardware_concurrency() / 2); + std::vector devices; }; static void print_usage(int /*argc*/, char** argv, rpc_server_params params) { fprintf(stderr, "Usage: %s [options]\n\n", argv[0]); fprintf(stderr, "options:\n"); - fprintf(stderr, " -h, --help show this help message and exit\n"); - fprintf(stderr, " -t, --threads number of threads for the CPU backend (default: %d)\n", params.n_threads); - fprintf(stderr, " -H HOST, --host HOST host to bind to (default: %s)\n", params.host.c_str()); - fprintf(stderr, " -p PORT, --port PORT port to bind to (default: %d)\n", params.port); - fprintf(stderr, " -m MEM, --mem MEM backend memory size (in MB)\n"); - fprintf(stderr, " -c, --cache enable local file cache\n"); + fprintf(stderr, " -h, --help show this help message and exit\n"); + fprintf(stderr, " -t, --threads N number of threads for the CPU device (default: %d)\n", params.n_threads); + fprintf(stderr, " -d, -dev, --device comma-separated list of devices\n"); + fprintf(stderr, " -cpu enable cpu backend\n"); + fprintf(stderr, " -h, -H, --host, --Host HOST host to bind to (default: %s)\n", params.host.c_str()); + fprintf(stderr, " -p, -P, --port, --Port PORT port to bind to (default: %d)\n", params.port); + fprintf(stderr, " -c, --cache enable local file cache\n"); fprintf(stderr, "\n"); } @@ -168,7 +171,7 @@ static bool rpc_server_params_parse(int argc, char** argv, rpc_server_params& pa std::string arg; for (int i = 1; i < argc; i++) { arg = argv[i]; - if (arg == "-H" || arg == "--host") { + if (arg == "-H" || arg == "-h" || arg == "--host" || arg == "--Host") { if (++i >= argc) { return false; } @@ -184,7 +187,25 @@ static bool rpc_server_params_parse(int argc, char** argv, rpc_server_params& pa return false; } } - else if (arg == "-p" || arg == "--port") { + else if (arg == "-d" || arg == "-dev" || arg == "--device") { + if (++i >= argc) { + return false; + } + const std::regex regex{ R"([,/]+)" }; + std::string dev_str = argv[i]; + std::sregex_token_iterator iter(dev_str.begin(), dev_str.end(), regex, -1); + std::sregex_token_iterator end; + for (; iter != end; ++iter) { + try { + params.devices.push_back(*iter); + } + catch (const std::exception&) { + fprintf(stderr, "error: invalid device: %s\n", iter->str().c_str()); + return false; + } + } + } + else if (arg == "-p" || arg == "-P" || arg == "--port" || arg == "--Port") { if (++i >= argc) { return false; } @@ -196,11 +217,8 @@ static bool rpc_server_params_parse(int argc, char** argv, rpc_server_params& pa else if (arg == "-c" || arg == "--cache") { params.use_cache = true; } - else if (arg == "-m" || arg == "--mem") { - if (++i >= argc) { - return false; - } - params.backend_mem = std::stoul(argv[i]) * 1024 * 1024; + else if (arg == "-cpu") { + params.use_cpu = true; } else if (arg == "-h" || arg == "--help") { print_usage(argc, argv, params); @@ -215,11 +233,18 @@ static bool rpc_server_params_parse(int argc, char** argv, rpc_server_params& pa return true; } -static ggml_backend_t create_backend(const rpc_server_params& params) { +static ggml_backend_t create_cpu_backend(const rpc_server_params& params) { + fprintf(stderr, "%s: using CPU backend\n", __func__); + ggml_backend_t backend = ggml_backend_cpu_init(); + ggml_backend_cpu_set_n_threads(backend, params.n_threads); + return backend; +} + +static ggml_backend_t create_gpu_backend(const rpc_server_params& params, uint32_t device) { ggml_backend_t backend = NULL; #ifdef GGML_USE_CUDA - fprintf(stderr, "%s: using CUDA backend\n", __func__); - backend = ggml_backend_cuda_init(0, nullptr); // init device 0 + fprintf(stderr, "%s: using CUDA backend: CUDA%d\n", __func__, device); + backend = ggml_backend_cuda_init(device, nullptr); // init device if (!backend) { fprintf(stderr, "%s: ggml_backend_cuda_init() failed\n", __func__); } @@ -231,34 +256,113 @@ static ggml_backend_t create_backend(const rpc_server_params& params) { } #elif GGML_USE_VULKAN fprintf(stderr, "%s: using Vulkan backend\n", __func__); - backend = ggml_backend_vk_init(0); // init device 0 + backend = ggml_backend_vk_init(device); // init device 0 if (!backend) { fprintf(stderr, "%s: ggml_backend_vulkan_init() failed\n", __func__); } #elif GGML_USE_SYCL fprintf(stderr, "%s: using SYCL backend\n", __func__); - backend = ggml_backend_sycl_init(0); // init device 0 + backend = ggml_backend_sycl_init(device); // init device 0 if (!backend) { fprintf(stderr, "%s: ggml_backend_sycl_init() failed\n", __func__); } #endif - // if there aren't GPU Backends fallback to CPU backend - if (!backend) { - fprintf(stderr, "%s: using CPU backend\n", __func__); - backend = ggml_backend_cpu_init(); - ggml_backend_cpu_set_n_threads(backend, params.n_threads); - } + //if (!backend) { + // fprintf(stderr, "%s: using CPU backend\n", __func__); + // backend = ggml_backend_cpu_init(); + // ggml_backend_cpu_set_n_threads(backend, params.n_threads); + //} return backend; } -static void get_backend_memory(size_t * free_mem, size_t * total_mem) { +static int32_t find_device_idx(const std::string& str) { + std::regex pattern(R"((\d+)$)"); // Match digits at the end + std::smatch matches; + int number = -1; + if (std::regex_search(str, matches, pattern)) { + number = std::stoi(matches[1]); + } + return number; +} + +static size_t get_gpu_backend_count(const rpc_server_params& params) { + size_t count = 0; +#if defined(GGML_USE_CUDA) + count = ggml_backend_cuda_get_device_count(); +#elif defined(GGML_USE_SYCL) + count = ggml_backend_sycl_get_device_count(); +#elif defined(GGML_USE_VULKAN) + count = ggml_backend_vk_get_device_count(); +#elif defined(GGML_USE_CANN) + return ggml_backend_cann_get_device_count(); +#endif + return count; +} + +static std::vector get_devices(const rpc_server_params& params) { + std::vector devices; + if (!params.devices.empty()) { + for (auto device : params.devices) { + int32_t device_id; + ggml_backend_t dev; + if (params.use_cpu && device == "CPU" ) { + dev = create_cpu_backend(params); + } else { + device_id = find_device_idx(device); + if (device_id < 0) { + fprintf(stderr, "error: unknown device: %s\n", device.c_str()); + continue; + } + dev = create_gpu_backend(params, device_id); + } + if (dev) { + devices.push_back(dev); + } else { + fprintf(stderr, "error: unknown device: %s\n", device.c_str()); + } + } + } + else { + for (size_t i = 0; i < get_gpu_backend_count(params); i++) { + ggml_backend_t dev = create_gpu_backend(params, i); + if (dev) { + devices.push_back(dev); + } + } + // cpu backend at last + if (params.use_cpu || devices.empty()) { + ggml_backend_t dev = create_cpu_backend(params); + if (dev) { + devices.push_back(dev); + } + } + } + return devices; +} + +static void get_cpu_backend_memory(size_t * free_mem, size_t * total_mem) { +#ifdef _WIN32 + MEMORYSTATUSEX status; + status.dwLength = sizeof(status); + GlobalMemoryStatusEx(&status); + *total_mem = status.ullTotalPhys; + *free_mem = status.ullAvailPhys; +#else + long pages = sysconf(_SC_PHYS_PAGES); + long page_size = sysconf(_SC_PAGE_SIZE); + *total_mem = pages * page_size; + *free_mem = *total_mem; +#endif +} + +static void get_backend_memory(uint32_t device, size_t * free_mem, size_t * total_mem) { #ifdef GGML_USE_CUDA - ggml_backend_cuda_get_device_memory(0, free_mem, total_mem); + ggml_backend_cuda_get_device_memory(device, free_mem, total_mem); #elif GGML_USE_VULKAN - ggml_backend_vk_get_device_memory(0, free_mem, total_mem); + ggml_backend_vk_get_device_memory(device, free_mem, total_mem); #elif GGML_USE_SYCL - ggml_backend_sycl_get_device_memory(0, free_mem, total_mem); + ggml_backend_sycl_get_device_memory(device, free_mem, total_mem); #else #ifdef _WIN32 MEMORYSTATUSEX status; @@ -292,20 +396,27 @@ int main(int argc, char * argv[]) { fprintf(stderr, "\n"); } - ggml_backend_t backend = create_backend(params); - if (!backend) { - fprintf(stderr, "Failed to create backend\n"); + auto devices = get_devices(params); + if (devices.empty()) { + fprintf(stderr, "No backend found\n"); return 1; } + std::string endpoint = params.host + ":" + std::to_string(params.port); - size_t free_mem, total_mem; - if (params.backend_mem > 0) { - free_mem = params.backend_mem; - total_mem = params.backend_mem; - } - else { - get_backend_memory(&free_mem, &total_mem); + std::vector free_mem, total_mem; + for (size_t i = 0; i < devices.size(); i++) { + size_t free, total; + const char* name = ggml_backend_name(devices[i]); + if (std::string(name) == "CPU") { + get_cpu_backend_memory(&free, &total); + } else { + int32_t idx = find_device_idx(name); + get_backend_memory((uint32_t) idx, &free, &total); + } + free_mem.push_back(free); + total_mem.push_back(total); } + const char * cache_dir = nullptr; std::string cache_dir_str; if (params.use_cache) { @@ -316,14 +427,7 @@ int main(int argc, char * argv[]) { } cache_dir = cache_dir_str.c_str(); } - printf("Starting RPC server v%d.%d.%d\n", - RPC_PROTO_MAJOR_VERSION, - RPC_PROTO_MINOR_VERSION, - RPC_PROTO_PATCH_VERSION); - printf(" endpoint : %s\n", endpoint.c_str()); - printf(" local cache : %s\n", cache_dir ? cache_dir : "n/a"); - printf(" backend memory : %zu MB\n", free_mem / (1024 * 1024)); - ggml_backend_rpc_start_server(backend, endpoint.c_str(), cache_dir, free_mem, total_mem); - ggml_backend_free(backend); + ggml_backend_rpc_start_server(endpoint.c_str(), cache_dir, devices.size(), devices.data(), + free_mem.data(), total_mem.data()); return 0; } diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 549e8504..08a81448 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -7,22 +7,22 @@ extern "C" { #endif -#define RPC_PROTO_MAJOR_VERSION 2 -#define RPC_PROTO_MINOR_VERSION 0 +#define RPC_PROTO_MAJOR_VERSION 3 +#define RPC_PROTO_MINOR_VERSION 5 #define RPC_PROTO_PATCH_VERSION 1 #define GGML_RPC_MAX_SERVERS 16 // backend API -GGML_API GGML_CALL ggml_backend_t ggml_backend_rpc_init(const char * endpoint); +GGML_API GGML_CALL ggml_backend_t ggml_backend_rpc_init(const char * endpoint, uint32_t device); GGML_API GGML_CALL bool ggml_backend_is_rpc(ggml_backend_t backend); -GGML_API GGML_CALL ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint); +GGML_API GGML_CALL ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, uint32_t device); -GGML_API GGML_CALL void ggml_backend_rpc_get_device_memory(const char * endpoint, size_t * free, size_t * total); +GGML_API GGML_CALL uint32_t ggml_backend_rpc_get_device_count(const char* endpoint); -GGML_API GGML_CALL void ggml_backend_rpc_start_server(ggml_backend_t backend, const char * endpoint, - const char * cache_dir, - size_t free_mem, size_t total_mem); +GGML_API GGML_CALL void ggml_backend_rpc_get_device_memory(const char * endpoint, uint32_t device, size_t * free, size_t * total); + +GGML_API GGML_CALL void ggml_backend_rpc_start_server(const char * endpoint, const char* cache_dir, size_t device, ggml_backend_t * devices, size_t* free_mem, size_t* total_mem); #ifdef __cplusplus } diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 2ed282e2..813f4467 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -999,15 +999,6 @@ GGML_CALL static ggml_backend_t ggml_backend_reg_cpu_init(const char * params, v GGML_UNUSED(user_data); } -#ifdef GGML_USE_RPC -GGML_CALL static ggml_backend_t ggml_backend_reg_rpc_init(const char* params, void* user_data) { - return ggml_backend_rpc_init((const char*)user_data); - - GGML_UNUSED(params); - GGML_UNUSED(user_data); -} -#endif - // multi-buffer buffer struct ggml_backend_multi_buffer_context { @@ -2159,6 +2150,7 @@ void ggml_backend_sched_reset(ggml_backend_sched_t sched) { bool ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph) { GGML_ASSERT((int)sched->hash_set.size >= measure_graph->n_nodes + measure_graph->n_leafs); + ggml_backend_sched_synchronize(sched); ggml_backend_sched_split_graph(sched, measure_graph); @@ -2167,7 +2159,6 @@ bool ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph * } ggml_backend_sched_reset(sched); - ggml_backend_sched_synchronize(sched); return true; } diff --git a/ggml/src/ggml-rpc.cpp b/ggml/src/ggml-rpc.cpp index 4c44ae54..b9fd6f76 100644 --- a/ggml/src/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc.cpp @@ -1,7 +1,7 @@ #include "ggml-rpc.h" #include "ggml.h" #include "ggml-backend-impl.h" - +#include "ggml-cpp.h" #include #include #include @@ -30,14 +30,18 @@ #include namespace fs = std::filesystem; + +static constexpr size_t MAX_CHUNK_SIZE = 1024ull * 1024ull * 1024ull; // 1 GiB + #define UNUSED GGML_UNUSED #define GGML_DEBUG 0 #if (GGML_DEBUG >= 1) #define GGML_PRINT_DEBUG(...) printf(__VA_ARGS__) #else -#define GGML_PRINT_DEBUG(...) +#define LOG_DBG(...) #endif +#define GGML_LOG_ERROR(...) printf( __VA_ARGS__) #ifdef _WIN32 typedef SOCKET sockfd_t; @@ -51,7 +55,7 @@ struct socket_t { sockfd_t fd; socket_t(sockfd_t fd) : fd(fd) {} ~socket_t() { - GGML_PRINT_DEBUG("[%s] closing socket %d\n", __func__, this->fd); + LOG_DBG("[%s] closing socket %d\n", __func__, this->fd); #ifdef _WIN32 closesocket(this->fd); #else @@ -60,6 +64,9 @@ struct socket_t { } }; +// macro for nicer error messages on server crash +#define RPC_STATUS_ASSERT(x) if (!(x)) GGML_ABORT("Remote RPC server crashed or returned malformed response") + // ggml_tensor is serialized into rpc_tensor #pragma pack(push, 1) struct rpc_tensor { @@ -80,8 +87,8 @@ struct rpc_tensor { char padding[4]; }; + static_assert(sizeof(rpc_tensor) % 8 == 0, "rpc_tensor size must be multiple of 8"); -static std::unordered_map rpc_server_map; // RPC commands enum rpc_cmd { RPC_CMD_ALLOC_BUFFER = 0, @@ -99,6 +106,8 @@ enum rpc_cmd { RPC_CMD_INIT_TENSOR, RPC_CMD_GET_ALLOC_SIZE, RPC_CMD_HELLO, + RPC_CMD_DEVICE_COUNT, + RPC_CMD_GRAPH_RECOMPUTE, RPC_CMD_COUNT, }; @@ -111,7 +120,12 @@ struct rpc_msg_hello_rsp { uint8_t patch; }; +struct rpc_msg_device_count_rsp { + uint32_t device_count; +}; + struct rpc_msg_get_alloc_size_req { + uint32_t device; rpc_tensor tensor; }; @@ -124,6 +138,7 @@ struct rpc_msg_init_tensor_req { }; struct rpc_msg_alloc_buffer_req { + uint32_t device; uint64_t size; }; @@ -132,10 +147,18 @@ struct rpc_msg_alloc_buffer_rsp { uint64_t remote_size; }; +struct rpc_msg_get_alignment_req { + uint32_t device; +}; + struct rpc_msg_get_alignment_rsp { uint64_t alignment; }; +struct rpc_msg_get_max_size_req { + uint32_t device; +}; + struct rpc_msg_get_max_size_rsp { uint64_t max_size; }; @@ -157,6 +180,12 @@ struct rpc_msg_buffer_clear_req { uint8_t value; }; +struct rpc_msg_set_tensor_hash_req { + rpc_tensor tensor; + uint64_t offset; + uint64_t hash; +}; + struct rpc_msg_set_tensor_hash_rsp { uint8_t result; }; @@ -176,14 +205,19 @@ struct rpc_msg_copy_tensor_rsp { uint8_t result; }; -struct rpc_msg_graph_compute_rsp { - uint8_t result; +struct rpc_msg_get_device_memory_req { + uint32_t device; }; struct rpc_msg_get_device_memory_rsp { uint64_t free_mem; uint64_t total_mem; }; + +struct rpc_msg_graph_recompute_req { + uint32_t device; +}; + #pragma pack(pop) // RPC data structures @@ -194,16 +228,45 @@ static ggml_guid_t ggml_backend_rpc_guid() { struct ggml_backend_rpc_buffer_type_context { std::string endpoint; + uint32_t device; std::string name; - size_t alignment; - size_t max_size; + size_t alignment; + size_t max_size; +}; + +struct graph_cache { + + bool is_cached(const ggml_cgraph * cgraph) { + if ((int)last_graph.size() != cgraph->n_nodes) { + return false; + } + for (int i = 0; i < cgraph->n_nodes; i++) { + if (memcmp(&last_graph[i], cgraph->nodes[i], sizeof(ggml_tensor)) != 0) { + return false; + } + } + return true; + } + + void add(const ggml_cgraph * cgraph) { + last_graph.resize(cgraph->n_nodes); + for (int i = 0; i < cgraph->n_nodes; i++) { + memcpy(&last_graph[i], cgraph->nodes[i], sizeof(ggml_tensor)); + } + } + + std::vector last_graph; }; struct ggml_backend_rpc_context { std::string endpoint; + uint32_t device; std::string name; + graph_cache gc; }; +static std::unordered_map rpc_server_map; + struct ggml_backend_rpc_buffer_context { std::shared_ptr sock; //std::unordered_map base_cache; @@ -318,30 +381,41 @@ static std::shared_ptr create_server_socket(const char * host, int por return sock; } -static bool send_data(sockfd_t sockfd, const void * data, size_t size) { +static bool send_data(sockfd_t sockfd, const void* data, size_t size) { size_t bytes_sent = 0; while (bytes_sent < size) { - ssize_t n = send(sockfd, (const char *)data + bytes_sent, size - bytes_sent, 0); + size_t size_to_send = std::min(size - bytes_sent, MAX_CHUNK_SIZE); + ssize_t n = send(sockfd, (const char*)data + bytes_sent, size_to_send, 0); if (n < 0) { + fprintf(stderr,"send failed (bytes_sent=%zu, size_to_send=%zu)\n", + bytes_sent, size_to_send); return false; } - bytes_sent += n; + bytes_sent += (size_t)n; } return true; } -static bool recv_data(sockfd_t sockfd, void * data, size_t size) { +static bool recv_data(sockfd_t sockfd, void* data, size_t size) { size_t bytes_recv = 0; while (bytes_recv < size) { - ssize_t n = recv(sockfd, (char *)data + bytes_recv, size - bytes_recv, 0); - if (n <= 0) { + size_t size_to_recv = std::min(size - bytes_recv, MAX_CHUNK_SIZE); + ssize_t n = recv(sockfd, (char*)data + bytes_recv, size_to_recv, 0); + if (n < 0) { + fprintf(stderr, "recv failed (bytes_recv=%zu, size_to_recv=%zu)\n", + bytes_recv, size_to_recv); return false; } - bytes_recv += n; + if (n == 0) { + LOG_DBG("recv returned 0 (peer closed?)\n"); + return false; + } + bytes_recv += (size_t)n; } return true; } + static bool send_msg(sockfd_t sockfd, const void* msg, size_t msg_size) { if (!send_data(sockfd, &msg_size, sizeof(msg_size))) { return false; @@ -422,11 +496,12 @@ static bool send_rpc_cmd(const std::shared_ptr & sock, enum rpc_cmd cm return true; } + // RPC client-side implementation static bool check_server_version(const std::shared_ptr& sock) { rpc_msg_hello_rsp response; bool status = send_rpc_cmd(sock, RPC_CMD_HELLO, nullptr, 0, &response, sizeof(response)); - GGML_ASSERT(status); + RPC_STATUS_ASSERT(status); if (response.major != RPC_PROTO_MAJOR_VERSION || response.minor > RPC_PROTO_MINOR_VERSION) { fprintf(stderr, "RPC server version mismatch: %d.%d.%d\n", response.major, response.minor, response.patch); return false; @@ -473,7 +548,7 @@ static std::shared_ptr get_socket(const std::string & endpoint) { if (!check_server_version(sock)) { return nullptr; } - GGML_PRINT_DEBUG("[%s] connected to %s, sockfd=%d\n", __func__, endpoint.c_str(), sock->fd); + LOG_DBG("[%s] connected to %s, sockfd=%d\n", __func__, endpoint.c_str(), sock->fd); sockets[endpoint] = sock; return sock; } @@ -488,7 +563,7 @@ static void ggml_backend_rpc_buffer_free_buffer(ggml_backend_buffer_t buffer) { ggml_backend_rpc_buffer_context* ctx = (ggml_backend_rpc_buffer_context*)buffer->context; rpc_msg_free_buffer_req request = { ctx->remote_ptr }; bool status = send_rpc_cmd(ctx->sock, RPC_CMD_FREE_BUFFER, &request, sizeof(request), nullptr, 0); - GGML_ASSERT(status); + RPC_STATUS_ASSERT(status); delete ctx; } @@ -500,7 +575,7 @@ static void* ggml_backend_rpc_buffer_get_base(ggml_backend_buffer_t buffer) { rpc_msg_buffer_get_base_req request = { ctx->remote_ptr }; rpc_msg_buffer_get_base_rsp response; bool status = send_rpc_cmd(ctx->sock, RPC_CMD_BUFFER_GET_BASE, &request, sizeof(request), &response, sizeof(response)); - GGML_ASSERT(status); + RPC_STATUS_ASSERT(status); ctx->base_ptr = reinterpret_cast(response.base_ptr); return ctx->base_ptr; } @@ -554,7 +629,7 @@ GGML_CALL static void ggml_backend_rpc_buffer_init_tensor(ggml_backend_buffer_t request.tensor = serialize_tensor(tensor); bool status = send_rpc_cmd(ctx->sock, RPC_CMD_INIT_TENSOR, &request, sizeof(request), nullptr, 0); - GGML_ASSERT(status); + RPC_STATUS_ASSERT(status); } } @@ -562,16 +637,16 @@ static void ggml_backend_rpc_buffer_set_tensor(ggml_backend_buffer_t buffer, ggm ggml_backend_rpc_buffer_context* ctx = (ggml_backend_rpc_buffer_context*)buffer->context; rpc_tensor rpc_tensor = serialize_tensor(tensor); if (size > HASH_THRESHOLD) { - // input serialization format: | rpc_tensor | offset (8 bytes) | hash (8 bytes) - size_t input_size = sizeof(rpc_tensor) + sizeof(uint64_t) + sizeof(uint64_t); - std::vector input(input_size, 0); - uint64_t hash = fnv_hash((const uint8_t*)data, size); - memcpy(input.data(), &rpc_tensor, sizeof(rpc_tensor)); - memcpy(input.data() + sizeof(rpc_tensor), &offset, sizeof(offset)); - memcpy(input.data() + sizeof(rpc_tensor) + sizeof(offset), &hash, sizeof(hash)); + rpc_msg_set_tensor_hash_req request; + request.tensor = rpc_tensor; + request.offset = offset; + request.hash = fnv_hash((const uint8_t*)data, size); + + + rpc_msg_set_tensor_hash_rsp response; - bool status = send_rpc_cmd(ctx->sock, RPC_CMD_SET_TENSOR_HASH, input.data(), input.size(), &response, sizeof(response)); - GGML_ASSERT(status); + bool status = send_rpc_cmd(ctx->sock, RPC_CMD_SET_TENSOR_HASH, &request, sizeof(request), &response, sizeof(response)); + RPC_STATUS_ASSERT(status); if (response.result) { // the server has the same data, no need to send it return; @@ -584,7 +659,7 @@ static void ggml_backend_rpc_buffer_set_tensor(ggml_backend_buffer_t buffer, ggm memcpy(input.data() + sizeof(rpc_tensor), &offset, sizeof(offset)); memcpy(input.data() + sizeof(rpc_tensor) + sizeof(offset), data, size); bool status = send_rpc_cmd(ctx->sock, RPC_CMD_SET_TENSOR, input.data(), input.size()); - GGML_ASSERT(status); + RPC_STATUS_ASSERT(status); } @@ -595,34 +670,40 @@ static void ggml_backend_rpc_buffer_get_tensor(ggml_backend_buffer_t buffer, con request.offset = offset; request.size = size; bool status = send_rpc_cmd(ctx->sock, RPC_CMD_GET_TENSOR, &request, sizeof(request), data, size); - GGML_ASSERT(status); + RPC_STATUS_ASSERT(status); } +static bool ggml_backend_buffer_is_rpc(ggml_backend_buffer_t buffer) { + return buffer->iface.free_buffer == ggml_backend_rpc_buffer_free_buffer; +} static bool ggml_backend_rpc_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor* src, ggml_tensor* dst) { - // check if src and dst are on the same server - ggml_backend_buffer_t src_buffer = src->buffer; - ggml_backend_rpc_buffer_context* src_ctx = (ggml_backend_rpc_buffer_context*)src_buffer->context; - ggml_backend_buffer_t dst_buffer = dst->buffer; - ggml_backend_rpc_buffer_context* dst_ctx = (ggml_backend_rpc_buffer_context*)dst_buffer->context; - if (src_ctx->sock != dst_ctx->sock) { - return false; + if (ggml_backend_buffer_is_rpc(src->buffer)) { + // check if src and dst are on the same server + ggml_backend_buffer_t src_buffer = src->buffer; + ggml_backend_rpc_buffer_context* src_ctx = (ggml_backend_rpc_buffer_context*)src_buffer->context; + ggml_backend_buffer_t dst_buffer = dst->buffer; + ggml_backend_rpc_buffer_context* dst_ctx = (ggml_backend_rpc_buffer_context*)dst_buffer->context; + if (src_ctx->sock != dst_ctx->sock) { + return false; + } + ggml_backend_rpc_buffer_context* ctx = (ggml_backend_rpc_buffer_context*)buffer->context; + rpc_msg_copy_tensor_req request; + request.src = serialize_tensor(src); + request.dst = serialize_tensor(dst); + rpc_msg_copy_tensor_rsp response; + bool status = send_rpc_cmd(ctx->sock, RPC_CMD_COPY_TENSOR, &request, sizeof(request), &response, sizeof(response)); + RPC_STATUS_ASSERT(status); + return response.result; } - ggml_backend_rpc_buffer_context* ctx = (ggml_backend_rpc_buffer_context*)buffer->context; - rpc_msg_copy_tensor_req request; - request.src = serialize_tensor(src); - request.dst = serialize_tensor(dst); - rpc_msg_copy_tensor_rsp response; - bool status = send_rpc_cmd(ctx->sock, RPC_CMD_COPY_TENSOR, &request, sizeof(request), &response, sizeof(response)); - GGML_ASSERT(status); - return response.result; + return false; } static void ggml_backend_rpc_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { ggml_backend_rpc_buffer_context* ctx = (ggml_backend_rpc_buffer_context*)buffer->context; rpc_msg_buffer_clear_req request = { ctx->remote_ptr, value }; bool status = send_rpc_cmd(ctx->sock, RPC_CMD_BUFFER_CLEAR, &request, sizeof(request), nullptr, 0); - GGML_ASSERT(status); + RPC_STATUS_ASSERT(status); } @@ -644,14 +725,19 @@ GGML_CALL static const char * ggml_backend_rpc_buffer_type_name(ggml_backend_buf return buft_ctx->name.c_str(); } +static std::string create_rpc_name(std::string endpoint, uint32_t device) { + std::string dev_name = "RPC" + std::to_string(device) + "[" + std::string(endpoint) + "]"; + return dev_name; +} + static ggml_backend_buffer_t ggml_backend_rpc_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { ggml_backend_rpc_buffer_type_context* buft_ctx = (ggml_backend_rpc_buffer_type_context*)buft->context; - rpc_msg_alloc_buffer_req request = { size }; + rpc_msg_alloc_buffer_req request = { buft_ctx->device, size }; rpc_msg_alloc_buffer_rsp response; auto sock = get_socket(buft_ctx->endpoint); - std::string name= "RPC[" + std::string(buft_ctx->endpoint) + "]"; + std::string name = create_rpc_name(buft_ctx->endpoint, buft_ctx->device);// "RPC[" + std::string(buft_ctx->endpoint) + "]"; bool status = send_rpc_cmd(sock, RPC_CMD_ALLOC_BUFFER, &request, sizeof(request), &response, sizeof(response)); - GGML_ASSERT(status); + RPC_STATUS_ASSERT(status); if (response.remote_ptr != 0) { ggml_backend_buffer_t buffer = ggml_backend_buffer_init(buft, ggml_backend_rpc_buffer_interface, @@ -665,10 +751,11 @@ static ggml_backend_buffer_t ggml_backend_rpc_buffer_type_alloc_buffer(ggml_back } -static size_t get_alignment(const std::shared_ptr& sock) { +static size_t get_alignment(const std::shared_ptr& sock, uint32_t device) { + rpc_msg_get_alignment_req request = { device }; rpc_msg_get_alignment_rsp response; - bool status = send_rpc_cmd(sock, RPC_CMD_GET_ALIGNMENT, nullptr, 0, &response, sizeof(response)); - GGML_ASSERT(status); + bool status = send_rpc_cmd(sock, RPC_CMD_GET_ALIGNMENT, &request, sizeof(request), &response, sizeof(response)); + RPC_STATUS_ASSERT(status); return response.alignment; } @@ -678,10 +765,11 @@ GGML_CALL static size_t ggml_backend_rpc_buffer_type_get_alignment(ggml_backend_ return buft_ctx->alignment; } -static size_t get_max_size(const std::shared_ptr& sock) { +static size_t get_max_size(const std::shared_ptr& sock, uint32_t device) { + rpc_msg_get_max_size_req request = { device }; rpc_msg_get_max_size_rsp response; - bool status = send_rpc_cmd(sock, RPC_CMD_GET_MAX_SIZE, nullptr, 0, &response, sizeof(response)); - GGML_ASSERT(status); + bool status = send_rpc_cmd(sock, RPC_CMD_GET_MAX_SIZE, &request, sizeof(request), &response, sizeof(response)); + RPC_STATUS_ASSERT(status); return response.max_size; } @@ -697,12 +785,12 @@ GGML_CALL static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend auto sock = get_socket(buft_ctx->endpoint); rpc_msg_get_alloc_size_req request; - + request.device = buft_ctx->device; request.tensor = serialize_tensor(tensor); rpc_msg_get_alloc_size_rsp response; bool status = send_rpc_cmd(sock, RPC_CMD_GET_ALLOC_SIZE, &request, sizeof(request), &response, sizeof(response)); - GGML_ASSERT(status); + RPC_STATUS_ASSERT(status); return response.alloc_size; } @@ -736,7 +824,7 @@ GGML_CALL static void ggml_backend_rpc_free(ggml_backend_t backend) { GGML_CALL static ggml_backend_buffer_type_t ggml_backend_rpc_get_default_buffer_type(ggml_backend_t backend) { ggml_backend_rpc_context * ctx = (ggml_backend_rpc_context *)backend->context; - return ggml_backend_rpc_buffer_type(ctx->endpoint.c_str()); + return ggml_backend_rpc_buffer_type(ctx->endpoint.c_str(), ctx->device); } GGML_CALL static void ggml_backend_rpc_synchronize(ggml_backend_t backend) { @@ -759,7 +847,7 @@ static void add_tensor(ggml_tensor * tensor, std::vector & tensors, tensors.push_back(serialize_tensor(tensor)); } -static void serialize_graph(const ggml_cgraph * cgraph, std::vector & output) { +static void serialize_graph(uint32_t device, const ggml_cgraph* cgraph, std::vector& output) { uint32_t n_nodes = cgraph->n_nodes; std::vector tensors; std::unordered_set visited; @@ -767,29 +855,45 @@ static void serialize_graph(const ggml_cgraph * cgraph, std::vector & o add_tensor(cgraph->nodes[i], tensors, visited); } // serialization format: - // | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(rpc_tensor)) | + // | device (4 bytes) | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(rpc_tensor)) | uint32_t n_tensors = tensors.size(); - int output_size = sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t) + n_tensors * sizeof(rpc_tensor); + int output_size = 2 * sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t) + n_tensors * sizeof(rpc_tensor); output.resize(output_size, 0); - memcpy(output.data(), &n_nodes, sizeof(n_nodes)); + uint8_t* dest = output.data(); + memcpy(dest, &device, sizeof(device)); + dest += sizeof(device); + memcpy(dest, &n_nodes, sizeof(n_nodes)); + dest += sizeof(n_nodes); for (uint32_t i = 0; i < n_nodes; i++) { - memcpy(output.data() + sizeof(n_nodes) + i * sizeof(uint64_t), &cgraph->nodes[i], sizeof(uint64_t)); + memcpy(dest + i * sizeof(uint64_t), &cgraph->nodes[i], sizeof(uint64_t)); } - uint32_t * out_ntensors = (uint32_t *)(output.data() + sizeof(n_nodes) + n_nodes * sizeof(uint64_t)); - *out_ntensors = n_tensors; - rpc_tensor * out_tensors = (rpc_tensor *)(output.data() + sizeof(n_nodes) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t)); + dest += n_nodes * sizeof(uint64_t); + memcpy(dest, &n_tensors, sizeof(n_tensors)); + dest += sizeof(n_tensors); + rpc_tensor* out_tensors = (rpc_tensor*)dest; memcpy(out_tensors, tensors.data(), n_tensors * sizeof(rpc_tensor)); } static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, ggml_cgraph* cgraph) { ggml_backend_rpc_context* rpc_ctx = (ggml_backend_rpc_context*)backend->context; - std::vector input; - serialize_graph(cgraph, input); - rpc_msg_graph_compute_rsp response; - auto sock = get_socket(rpc_ctx->endpoint); - bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size(), &response, sizeof(response)); - GGML_ASSERT(status); - return (enum ggml_status)response.result; + + GGML_ASSERT(cgraph->n_nodes > 0); + bool reuse = rpc_ctx->gc.is_cached(cgraph); + if (reuse) { + rpc_msg_graph_recompute_req request; + request.device = rpc_ctx->device; + auto sock = get_socket(rpc_ctx->endpoint); + bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_RECOMPUTE, &request, sizeof(request)); + RPC_STATUS_ASSERT(status); + } else { + rpc_ctx->gc.add(cgraph); + std::vector input; + serialize_graph(rpc_ctx->device, cgraph, input); + auto sock = get_socket(rpc_ctx->endpoint); + bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size()); + RPC_STATUS_ASSERT(status); + } + return GGML_STATUS_SUCCESS; } GGML_CALL static bool ggml_backend_rpc_supports_op(ggml_backend_t backend, const ggml_tensor * op) { @@ -804,8 +908,8 @@ GGML_CALL static bool ggml_backend_rpc_supports_buft(ggml_backend_t backend, ggm return false; } ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context; - ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context; - return buft_ctx->endpoint == rpc_ctx->endpoint; + ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context; + return buft_ctx->name == rpc_ctx->name; } static ggml_backend_i ggml_backend_rpc_interface = { @@ -831,12 +935,13 @@ static ggml_backend_i ggml_backend_rpc_interface = { /* .event_synchronize = */ NULL, }; -GGML_API GGML_CALL ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint) { +GGML_API GGML_CALL ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, uint32_t device) { static std::mutex mutex; std::lock_guard lock(mutex); + std::string dev_name = create_rpc_name(endpoint, device); // NOTE: buffer types are allocated and never freed; this is by design static std::unordered_map buft_map; - auto it = buft_map.find(endpoint); + auto it = buft_map.find(dev_name); if (it != buft_map.end()) { return it->second; } @@ -845,26 +950,28 @@ GGML_API GGML_CALL ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const fprintf(stderr, "Failed to connect to %s\n", endpoint); return nullptr; } - size_t alignment = get_alignment(sock); - size_t max_size = get_max_size(sock); + size_t alignment = get_alignment(sock, device); + size_t max_size = get_max_size(sock, device); ggml_backend_rpc_buffer_type_context * buft_ctx = new ggml_backend_rpc_buffer_type_context { /* .endpoint = */ endpoint, - /* .name = */ "RPC[" + std::string(endpoint) + "]", + /* .device = */ device, + /* .name = */ dev_name , /* .alignment = */ alignment, /* .max_size = */ max_size }; - rpc_server_map[endpoint] = "RPC[" + std::string(endpoint) + "]"; + rpc_server_map[dev_name] = buft_ctx; ggml_backend_buffer_type_t buft = new ggml_backend_buffer_type { /* .iface = */ ggml_backend_rpc_buffer_type_interface, /* .context = */ buft_ctx }; - buft_map[endpoint] = buft; + buft_map[dev_name] = buft; return buft; } // backend registry GGML_CALL static ggml_backend_t ggml_backend_reg_rpc_init(const char* params, void* user_data) { - ggml_backend_t cuda_backend = ggml_backend_rpc_init((const char*)user_data); + auto rpc_ctx = (ggml_backend_rpc_buffer_type_context *) user_data; + ggml_backend_t cuda_backend = ggml_backend_rpc_init(rpc_ctx->endpoint.c_str(), rpc_ctx->device); return cuda_backend; GGML_UNUSED(params); @@ -878,19 +985,23 @@ GGML_CALL int ggml_backend_rpc_reg_devices() { int device_count = (int)rpc_server_map.size(); int i = 0; for (auto& it : rpc_server_map) - { - std::string endpoint = it.first; - auto name = "RPC[" + std::string(endpoint) + "]"; - ggml_backend_register(name.c_str(), ggml_backend_reg_rpc_init, ggml_backend_rpc_buffer_type(endpoint.c_str()), &(endpoint)); + { + std::string name = it.second->name; + std::string endpoint = std::string(it.second->endpoint); + uint32_t device = it.second->device; + ggml_backend_register(name.c_str(), ggml_backend_reg_rpc_init, ggml_backend_rpc_buffer_type(endpoint.c_str(), device), &(it.second)); i++; } return device_count; } -GGML_CALL ggml_backend_t ggml_backend_rpc_init(const char * endpoint) { +GGML_CALL ggml_backend_t ggml_backend_rpc_init(const char * endpoint, uint32_t device) { + std::string dev_name = create_rpc_name(endpoint, device); ggml_backend_rpc_context * ctx = new ggml_backend_rpc_context { - /* .endpoint = */ endpoint, - /* .name = */ "RPC[" + std::string(endpoint) + "]", + /* .endpoint = */ endpoint, + /* .device = */ device, + /* .name = */ dev_name, + /* .gc = */ {}, }; ggml_backend_t backend = new ggml_backend { @@ -905,46 +1016,57 @@ GGML_API GGML_CALL bool ggml_backend_is_rpc(ggml_backend_t backend) { return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_rpc_guid()); } -static void get_device_memory(const std::shared_ptr& sock, size_t* free, size_t* total) { +static void get_device_memory(const std::shared_ptr& sock, uint32_t device, size_t* free, size_t* total) { + rpc_msg_get_device_memory_req request; + request.device = device; rpc_msg_get_device_memory_rsp response; - bool status = send_rpc_cmd(sock, RPC_CMD_GET_DEVICE_MEMORY, nullptr, 0, &response, sizeof(response)); - GGML_ASSERT(status); + bool status = send_rpc_cmd(sock, RPC_CMD_GET_DEVICE_MEMORY, &request, sizeof(request), &response, sizeof(response)); + RPC_STATUS_ASSERT(status); *free = response.free_mem; *total = response.total_mem; } -GGML_API GGML_CALL void ggml_backend_rpc_get_device_memory(const char * endpoint, size_t * free, size_t * total) { +void ggml_backend_rpc_get_device_memory(const char* endpoint, uint32_t device, size_t* free, size_t* total) { auto sock = get_socket(endpoint); if (sock == nullptr) { *free = 0; *total = 0; return; } - get_device_memory(sock, free, total); + get_device_memory(sock, device, free, total); } + // RPC server-side implementation class rpc_server { public: - rpc_server(ggml_backend_t backend, const char* cache_dir) - : backend(backend), cache_dir(cache_dir) { + rpc_server(std::vector all_backends, const char * cache_dir) + : backends(std::move(all_backends)), cache_dir(cache_dir) { + stored_graphs.resize(backends.size()); } ~rpc_server(); - void hello(rpc_msg_hello_rsp& response); - void alloc_buffer(const rpc_msg_alloc_buffer_req& request, rpc_msg_alloc_buffer_rsp& response); - void get_alignment(rpc_msg_get_alignment_rsp& response); - void get_max_size(rpc_msg_get_max_size_rsp& response); - bool buffer_get_base(const rpc_msg_buffer_get_base_req& request, rpc_msg_buffer_get_base_rsp& response); - bool free_buffer(const rpc_msg_free_buffer_req& request); - bool buffer_clear(const rpc_msg_buffer_clear_req& request); - bool set_tensor(const std::vector& input); - bool set_tensor_hash(const std::vector& input, rpc_msg_set_tensor_hash_rsp& response); - bool get_tensor(const rpc_msg_get_tensor_req& request, std::vector& response); - bool copy_tensor(const rpc_msg_copy_tensor_req& request, rpc_msg_copy_tensor_rsp& response); - bool graph_compute(const std::vector& input, rpc_msg_graph_compute_rsp& response); - bool init_tensor(const rpc_msg_init_tensor_req& request); - bool get_alloc_size(const rpc_msg_get_alloc_size_req& request, rpc_msg_get_alloc_size_rsp& response); + void hello(rpc_msg_hello_rsp & response); + bool alloc_buffer(const rpc_msg_alloc_buffer_req & request, rpc_msg_alloc_buffer_rsp & response); + bool get_alignment(const rpc_msg_get_alignment_req & request, rpc_msg_get_alignment_rsp & response); + bool get_max_size(const rpc_msg_get_max_size_req & request, rpc_msg_get_max_size_rsp & response); + bool buffer_get_base(const rpc_msg_buffer_get_base_req & request, rpc_msg_buffer_get_base_rsp & response); + bool free_buffer(const rpc_msg_free_buffer_req & request); + bool buffer_clear(const rpc_msg_buffer_clear_req & request); + bool set_tensor(const std::vector & input); + bool set_tensor_hash(const rpc_msg_set_tensor_hash_req & request, rpc_msg_set_tensor_hash_rsp & response); + bool get_tensor(const rpc_msg_get_tensor_req & request, std::vector & response); + bool copy_tensor(const rpc_msg_copy_tensor_req & request, rpc_msg_copy_tensor_rsp & response); + bool graph_compute(const std::vector & input); + bool graph_recompute(const rpc_msg_graph_recompute_req & request); + bool init_tensor(const rpc_msg_init_tensor_req & request); + bool get_alloc_size(const rpc_msg_get_alloc_size_req & request, rpc_msg_get_alloc_size_rsp & response); + bool get_device_memory(const rpc_msg_get_device_memory_req & request, rpc_msg_get_device_memory_rsp & response); + + struct stored_graph { + ggml_context_ptr ctx_ptr; + ggml_cgraph * graph; + }; private: bool get_cached_file(uint64_t hash, std::vector& data); @@ -955,19 +1077,25 @@ private: std::unordered_map & tensor_map); - ggml_backend_t backend; + std::vector backends; const char* cache_dir; std::unordered_set buffers; + // store the last computed graph for each backend + std::vector stored_graphs; }; void rpc_server::hello(rpc_msg_hello_rsp& response) { response.major = RPC_PROTO_MAJOR_VERSION; response.minor = RPC_PROTO_MINOR_VERSION; response.patch = RPC_PROTO_PATCH_VERSION; - GGML_PRINT_DEBUG("[%s] version: %d.%d.%d\n", __func__, response.major, response.minor, response.patch); + LOG_DBG("[%s] version: %d.%d.%d\n", __func__, response.major, response.minor, response.patch); } bool rpc_server::get_alloc_size(const rpc_msg_get_alloc_size_req& request, rpc_msg_get_alloc_size_rsp& response) { + uint32_t dev_id = request.device; + if (dev_id >= backends.size()) { + return false; + } ggml_backend_buffer_type_t buft; struct ggml_init_params params { /*.mem_size =*/ ggml_tensor_overhead(), @@ -975,59 +1103,76 @@ bool rpc_server::get_alloc_size(const rpc_msg_get_alloc_size_req& request, rpc_m /*.no_alloc =*/ true, }; - struct ggml_context* ctx = ggml_init(params); + ggml_context_ptr ctx_ptr{ ggml_init(params) }; + GGML_ASSERT(ctx_ptr != nullptr); + ggml_context* ctx = ctx_ptr.get(); ggml_tensor* tensor = deserialize_tensor(ctx, &request.tensor); if (tensor == nullptr) { GGML_ABORT("Null tensor pointer passed to server get_alloc_size function.\n"); - ggml_free(ctx); return false; } - if (tensor->buffer == nullptr) { //No buffer allocated. - buft = ggml_backend_get_default_buffer_type(backend); + buft = ggml_backend_get_default_buffer_type(backends[dev_id]); } else { buft = tensor->buffer->buft; } - + LOG_DBG("[%s] device: %d, buffer: %p, data: %p\n", __func__, dev_id, (void*)tensor->buffer, tensor->data); response.alloc_size = ggml_backend_buft_get_alloc_size(buft, tensor); - ggml_free(ctx); return true; } -void rpc_server::alloc_buffer(const rpc_msg_alloc_buffer_req& request, rpc_msg_alloc_buffer_rsp& response) { - ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(backend); + +bool rpc_server::alloc_buffer(const rpc_msg_alloc_buffer_req& request, rpc_msg_alloc_buffer_rsp& response) { + uint32_t dev_id = request.device; + if (dev_id >= backends.size()) { + return false; + } + ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(backends[dev_id]); ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(buft, request.size); response.remote_ptr = 0; response.remote_size = 0; if (buffer != nullptr) { response.remote_ptr = reinterpret_cast(buffer); response.remote_size = buffer->size; - GGML_PRINT_DEBUG("[%s] size: %" PRIu64 " -> remote_ptr: %" PRIx64 ", remote_size: %" PRIu64 "\n", __func__, request.size, response.remote_ptr, response.remote_size); + LOG_DBG("[%s] device: %d, size: %" PRIu64 " -> remote_ptr: %" PRIx64 ", remote_size: %" PRIu64 "\n", + __func__, dev_id, request.size, response.remote_ptr, response.remote_size); buffers.insert(buffer); } else { - GGML_ABORT("[%s] size: %" PRIu64 " -> failed\n", __func__, request.size); + LOG_DBG("[%s] device: %d, size: %" PRIu64 " -> failed\n", __func__, dev_id, request.size); } -} -void rpc_server::get_alignment(rpc_msg_get_alignment_rsp& response) { - ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(backend); - size_t alignment = ggml_backend_buft_get_alignment(buft); - GGML_PRINT_DEBUG("[%s] alignment: %lu\n", __func__, alignment); - response.alignment = alignment; + return true; } -void rpc_server::get_max_size(rpc_msg_get_max_size_rsp& response) { - ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(backend); +bool rpc_server::get_alignment(const rpc_msg_get_alignment_req& request, rpc_msg_get_alignment_rsp& response) { + uint32_t dev_id = request.device; + if (dev_id >= backends.size()) { + return false; + } + ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(backends[dev_id]); + size_t alignment = ggml_backend_buft_get_alignment(buft); + LOG_DBG("[%s] device: %d, alignment: %lu\n", __func__, dev_id, alignment); + response.alignment = alignment; + return true; +} + +bool rpc_server::get_max_size(const rpc_msg_get_max_size_req& request, rpc_msg_get_max_size_rsp& response) { + uint32_t dev_id = request.device; + if (dev_id >= backends.size()) { + return false; + } + ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(backends[dev_id]); size_t max_size = ggml_backend_buft_get_max_size(buft); - GGML_PRINT_DEBUG("[%s] max_size: %lu\n", __func__, max_size); + LOG_DBG("[%s] device: %d, max_size: %lu\n", __func__, dev_id, max_size); response.max_size = max_size; + return true; } bool rpc_server::buffer_get_base(const rpc_msg_buffer_get_base_req& request, rpc_msg_buffer_get_base_rsp& response) { - GGML_PRINT_DEBUG("[%s] remote_ptr: %" PRIx64 "\n", __func__, request.remote_ptr); + LOG_DBG("[%s] remote_ptr: %" PRIx64 "\n", __func__, request.remote_ptr); ggml_backend_buffer_t buffer = reinterpret_cast(request.remote_ptr); if (buffers.find(buffer) == buffers.end()) { GGML_ABORT("[%s] buffer not found\n", __func__); @@ -1039,7 +1184,7 @@ bool rpc_server::buffer_get_base(const rpc_msg_buffer_get_base_req& request, rpc } bool rpc_server::free_buffer(const rpc_msg_free_buffer_req& request) { - GGML_PRINT_DEBUG("[%s] remote_ptr: %" PRIx64 "\n", __func__, request.remote_ptr); + LOG_DBG("[%s] remote_ptr: %" PRIx64 "\n", __func__, request.remote_ptr); ggml_backend_buffer_t buffer = reinterpret_cast(request.remote_ptr); if (buffers.find(buffer) == buffers.end()) { GGML_ABORT("[%s] buffer not found\n", __func__); @@ -1051,7 +1196,7 @@ bool rpc_server::free_buffer(const rpc_msg_free_buffer_req& request) { } bool rpc_server::buffer_clear(const rpc_msg_buffer_clear_req& request) { - GGML_PRINT_DEBUG("[%s] remote_ptr: %" PRIx64 ", value: %u\n", __func__, request.remote_ptr, request.value); + LOG_DBG("[%s] remote_ptr: %" PRIx64 ", value: %u\n", __func__, request.remote_ptr, request.value); ggml_backend_buffer_t buffer = reinterpret_cast(request.remote_ptr); if (buffers.find(buffer) == buffers.end()) { GGML_ABORT("[%s] buffer not found\n", __func__); @@ -1064,7 +1209,7 @@ bool rpc_server::buffer_clear(const rpc_msg_buffer_clear_req& request) { ggml_tensor * rpc_server::deserialize_tensor(struct ggml_context * ctx, const rpc_tensor * tensor) { // Validate tensor type before using it if (tensor->type >= GGML_TYPE_COUNT) { - GGML_PRINT_DEBUG("[%s] invalid tensor type received: %u\n", __func__, tensor->type); + LOG_DBG("[%s] invalid tensor type received: %u\n", __func__, tensor->type); return nullptr; } @@ -1073,7 +1218,7 @@ ggml_tensor * rpc_server::deserialize_tensor(struct ggml_context * ctx, const rp // ggml_new_tensor_4d might fail if dimensions are invalid, although less likely to crash than invalid type if (result == nullptr) { - GGML_PRINT_DEBUG("[%s] ggml_new_tensor_4d failed for type %u\\n", __func__, tensor->type); + LOG_DBG("[%s] ggml_new_tensor_4d failed for type %u\\n", __func__, tensor->type); return nullptr; } @@ -1120,14 +1265,15 @@ bool rpc_server::set_tensor(const std::vector& input) { /*.mem_buffer =*/ NULL, /*.no_alloc =*/ true, }; - struct ggml_context* ctx = ggml_init(params); + ggml_context_ptr ctx_ptr{ ggml_init(params) }; + GGML_ASSERT(ctx_ptr != nullptr); + ggml_context* ctx = ctx_ptr.get(); ggml_tensor* tensor = deserialize_tensor(ctx, in_tensor); if (tensor == nullptr) { GGML_ABORT("[%s] error deserializing tensor\n", __func__); - ggml_free(ctx); return false; } - GGML_PRINT_DEBUG("[%s] buffer: %p, data: %p, offset: %" PRIu64 ", size: %zu\n", __func__, (void*)tensor->buffer, tensor->data, offset, size); + LOG_DBG("[%s] buffer: %p, data: %p, offset: %" PRIu64 ", size: %zu\n", __func__, (void*)tensor->buffer, tensor->data, offset, size); // sanitize tensor->data { @@ -1135,7 +1281,7 @@ bool rpc_server::set_tensor(const std::vector& input) { const size_t p1 = p0 + ggml_backend_buffer_get_size(tensor->buffer); if (in_tensor->data + offset < p0 || in_tensor->data + offset >= p1 || size >(p1 - in_tensor->data - offset)) { - GGML_PRINT_DEBUG("[%s] tensor data region (data=0x%" PRIx64 ", offset=%" PRIu64 ", size=%zu) out of buffer bounds [0x%zx, 0x%zx)\n", + LOG_DBG("[%s] tensor data region (data=0x%" PRIx64 ", offset=%" PRIu64 ", size=%zu) out of buffer bounds [0x%zx, 0x%zx)\n", __func__, in_tensor->data, offset, size, p0, p1); return false; } @@ -1153,7 +1299,6 @@ bool rpc_server::set_tensor(const std::vector& input) { printf("[%s] saved to '%s'\n", __func__, cache_file.c_str()); } ggml_backend_tensor_set(tensor, data, offset, size); - ggml_free(ctx); return true; } @@ -1177,18 +1322,10 @@ bool rpc_server::get_cached_file(uint64_t hash, std::vector& data) { return true; } -bool rpc_server::set_tensor_hash(const std::vector& input, rpc_msg_set_tensor_hash_rsp& response) +bool rpc_server::set_tensor_hash(const rpc_msg_set_tensor_hash_req& request, rpc_msg_set_tensor_hash_rsp& response) { - // serialization format: | rpc_tensor | offset (8 bytes) | hash (8 bytes) | - if (input.size() != sizeof(rpc_tensor) + 16) { - return false; - } - const rpc_tensor* in_tensor = (const rpc_tensor*)input.data(); - uint64_t offset; - memcpy(&offset, input.data() + sizeof(rpc_tensor), sizeof(offset)); - const uint64_t* hash = (const uint64_t*)(input.data() + sizeof(rpc_tensor) + sizeof(offset)); std::vector cached_file; - if (!get_cached_file(*hash, cached_file)) { + if (!get_cached_file(request.hash, cached_file)) { response.result = 0; return true; } @@ -1198,29 +1335,32 @@ bool rpc_server::set_tensor_hash(const std::vector& input, rpc_msg_set_ /*.mem_buffer =*/ NULL, /*.no_alloc =*/ true, }; - struct ggml_context* ctx = ggml_init(params); - ggml_tensor* tensor = deserialize_tensor(ctx, in_tensor); - if (tensor == nullptr) { - GGML_ABORT("[%s] error deserializing tensor\n", __func__); - ggml_free(ctx); + ggml_context_ptr ctx_ptr{ ggml_init(params) }; + GGML_ASSERT(ctx_ptr != nullptr); + ggml_context* ctx = ctx_ptr.get(); + ggml_tensor* tensor = deserialize_tensor(ctx, &request.tensor); + if (tensor == nullptr || tensor->buffer == nullptr) { + GGML_LOG_ERROR("[%s] error deserializing tensor\n", __func__); return false; } - GGML_PRINT_DEBUG("[%s] buffer: %p, data: %p, offset: %" PRIu64 ", size: %zu, hash: %" PRIx64 "\n", __func__, (void*)tensor->buffer, tensor->data, offset, size, *hash); + LOG_DBG("[%s] buffer: %p, data: %p, offset: %" PRIu64 ", size: %zu, hash: %" PRIx64 "\n", + __func__, (void*)tensor->buffer, tensor->data, request.offset, size, request.hash); // sanitize tensor->data { const size_t p0 = (size_t)ggml_backend_buffer_get_base(tensor->buffer); const size_t p1 = p0 + ggml_backend_buffer_get_size(tensor->buffer); - if (in_tensor->data + offset < p0 || in_tensor->data + offset >= p1 || size >(p1 - in_tensor->data - offset)) { - GGML_PRINT_DEBUG("[%s] tensor data region (data=0x%" PRIx64 ", offset=%" PRIu64 ", size=%zu, hash=0x%" PRIx64 ") out of buffer bounds [0x%zx, 0x%zx)\n", - __func__, in_tensor->data, offset, size, *hash, p0, p1); + if (request.tensor.data + request.offset < p0 + || request.tensor.data + request.offset >= p1 + || size >(p1 - request.tensor.data - request.offset)) { + GGML_LOG_ERROR("[%s] tensor data region (data=0x%" PRIx64 ", offset=%" PRIu64 ", size=%zu, hash=0x%" PRIx64 ") out of buffer bounds [0x%zx, 0x%zx)\n", + __func__, request.tensor.data, request.offset, size, request.hash, p0, p1); return false; } } - ggml_backend_tensor_set(tensor, cached_file.data(), offset, size); + ggml_backend_tensor_set(tensor, cached_file.data(), request.offset, size); response.result = 1; - ggml_free(ctx); return true; } @@ -1230,11 +1370,13 @@ bool rpc_server::init_tensor(const rpc_msg_init_tensor_req& request) { /*.mem_buffer =*/ NULL, /*.no_alloc =*/ true, }; - struct ggml_context* ctx = ggml_init(params); + + ggml_context_ptr ctx_ptr{ ggml_init(params) }; + GGML_ASSERT(ctx_ptr != nullptr); + ggml_context* ctx = ctx_ptr.get(); ggml_tensor* tensor = deserialize_tensor(ctx, &request.tensor); if (tensor == nullptr) { - GGML_PRINT_DEBUG("Null tensor pointer passed to server init_tensor function.\n"); - ggml_free(ctx); + LOG_DBG("Null tensor pointer passed to server init_tensor function.\n"); return false; } @@ -1244,18 +1386,16 @@ bool rpc_server::init_tensor(const rpc_msg_init_tensor_req& request) { buffer->iface.init_tensor(buffer, tensor); } else { - GGML_PRINT_DEBUG("Null buffer for tensor passed to init_tensor function\n"); + LOG_DBG("Null buffer for tensor passed to init_tensor function\n"); } if (tensor->extra != nullptr) { // This pointer can either be passed around client/server, or probably better stored server-side and kept track of. // Currently unimplemented. - GGML_PRINT_DEBUG("tensor->extra populated by the backend, this is currently unsupported.\n"); - ggml_free(ctx); + LOG_DBG("tensor->extra populated by the backend, this is currently unsupported.\n"); return false; } - ggml_free(ctx); return true; } @@ -1265,14 +1405,15 @@ bool rpc_server::get_tensor(const rpc_msg_get_tensor_req& request, std::vectorbuffer, tensor->data, request.offset, request.size); + LOG_DBG("[%s] buffer: %p, data: %p, offset: %" PRIu64 ", size: %" PRIu64 "\n", __func__, (void*)tensor->buffer, tensor->data, request.offset, request.size); // sanitize tensor->data { @@ -1282,7 +1423,7 @@ bool rpc_server::get_tensor(const rpc_msg_get_tensor_req& request, std::vector= p1 || request.size > (p1 - request.tensor.data - request.offset)) { - GGML_PRINT_DEBUG("[%s] requested tensor region (data=0x%" PRIx64 ", offset=%" PRIu64 ", size=%" PRIu64 ") out of buffer bounds [0x%zx, 0x%zx)\n", + LOG_DBG("[%s] requested tensor region (data=0x%" PRIx64 ", offset=%" PRIu64 ", size=%" PRIu64 ") out of buffer bounds [0x%zx, 0x%zx)\n", __func__, request.tensor.data, request.offset, request.size, p0, p1); return false; } @@ -1290,7 +1431,6 @@ bool rpc_server::get_tensor(const rpc_msg_get_tensor_req& request, std::vectorbuffer); if (dst_data + src_size > dst_base + dst_buf_sz) { - GGML_PRINT_DEBUG("[%s] out-of-bounds write in rpc_server::copy_tensor:\n" + LOG_DBG("[%s] out-of-bounds write in rpc_server::copy_tensor:\n" " write range : [0x%" PRIx64 ", 0x%" PRIx64 "]\n" " buffer base: [0x%" PRIx64 ", 0x%" PRIx64 "]\n", __func__, @@ -1322,15 +1463,13 @@ bool rpc_server::copy_tensor(const rpc_msg_copy_tensor_req& request, rpc_msg_cop dst_data + src_size, dst_base, dst_base + dst_buf_sz); - ggml_free(ctx); return false; } - GGML_PRINT_DEBUG("[%s] src->buffer: %p, dst->buffer: %p\n", + LOG_DBG("[%s] src->buffer: %p, dst->buffer: %p\n", __func__, (void*)src->buffer, (void*)dst->buffer); response.result = ggml_backend_buffer_copy_tensor(src, dst); - ggml_free(ctx); return true; } @@ -1361,7 +1500,7 @@ ggml_tensor* rpc_server::create_node(uint64_t id, result->src[i] = create_node(tensor->src[i], ctx, tensor_ptrs, tensor_map); // If the recursive call failed for a non-zero ID, propagate the error if (result->src[i] == nullptr) { - GGML_PRINT_DEBUG("[%s] failed to create source node %d (src_id=%" PRIu64 ") for node id %" PRIu64 "\n", + LOG_DBG("[%s] failed to create source node %d (src_id=%" PRIu64 ") for node id %" PRIu64 "\n", __func__, i, tensor->src[i], id); // Must return nullptr to signal failure up the call stack return nullptr; @@ -1376,7 +1515,7 @@ ggml_tensor* rpc_server::create_node(uint64_t id, result->view_src = create_node(tensor->view_src, ctx, tensor_ptrs, tensor_map); // If the recursive call failed for a non-zero ID, propagate the error if (result->view_src == nullptr) { - GGML_PRINT_DEBUG("[%s] failed to create view_src node (view_src_id=%" PRIu64 ") for node id %" PRIu64 "\n", + LOG_DBG("[%s] failed to create view_src node (view_src_id=%" PRIu64 ") for node id %" PRIu64 "\n", __func__, tensor->view_src, id); // Must return nullptr to signal failure up the call stack return nullptr; @@ -1386,25 +1525,35 @@ ggml_tensor* rpc_server::create_node(uint64_t id, return result; } -bool rpc_server::graph_compute(const std::vector& input, rpc_msg_graph_compute_rsp& response) { +bool rpc_server::graph_compute(const std::vector& input) { // serialization format: - // | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(rpc_tensor)) | - if (input.size() < sizeof(uint32_t)) { + // | device (4 bytes) | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(rpc_tensor)) | + if (input.size() < 2 * sizeof(uint32_t)) { + return false; + } + const uint8_t* src = input.data(); + uint32_t device; + memcpy(&device, src, sizeof(device)); + src += sizeof(device); + if (device >= backends.size()) { return false; } uint32_t n_nodes; - memcpy(&n_nodes, input.data(), sizeof(n_nodes)); - if (input.size() < sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t)) { + memcpy(&n_nodes, src, sizeof(n_nodes)); + src += sizeof(n_nodes); + if (input.size() < 2 * sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t)) { return false; } - const uint64_t* nodes = (const uint64_t*)(input.data() + sizeof(n_nodes)); + const uint64_t* nodes = (const uint64_t*)src; + src += n_nodes * sizeof(uint64_t); uint32_t n_tensors; - memcpy(&n_tensors, input.data() + sizeof(n_nodes) + n_nodes * sizeof(uint64_t), sizeof(n_tensors)); - if (input.size() < sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t) + n_tensors * sizeof(rpc_tensor)) { + memcpy(&n_tensors, src, sizeof(n_tensors)); + src += sizeof(n_tensors); + if (input.size() < 2 * sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t) + n_tensors * sizeof(rpc_tensor)) { return false; } - const rpc_tensor* tensors = (const rpc_tensor*)(input.data() + sizeof(n_nodes) + n_nodes * sizeof(uint64_t) + sizeof(n_tensors)); - GGML_PRINT_DEBUG("[%s] n_nodes: %u, n_tensors: %u\n", __func__, n_nodes, n_tensors); + const rpc_tensor* tensors = (const rpc_tensor*)src; + LOG_DBG("[%s] device: %u, n_nodes: %u, n_tensors: %u\n", __func__, device, n_nodes, n_tensors); size_t buf_size = ggml_tensor_overhead() * (n_nodes + n_tensors) + ggml_graph_overhead_custom(n_nodes, false); @@ -1413,7 +1562,9 @@ bool rpc_server::graph_compute(const std::vector& input, rpc_msg_graph_ /*.mem_buffer =*/ NULL, /*.no_alloc =*/ true, }; - struct ggml_context* ctx = ggml_init(params); + ggml_context_ptr ctx_ptr{ ggml_init(params) }; + GGML_ASSERT(ctx_ptr != nullptr); + ggml_context* ctx = ctx_ptr.get(); struct ggml_cgraph* graph = ggml_new_graph_custom(ctx, n_nodes, false); graph->n_nodes = n_nodes; std::unordered_map tensor_ptrs; @@ -1430,24 +1581,42 @@ bool rpc_server::graph_compute(const std::vector& input, rpc_msg_graph_ // If id was 0, create_node returning nullptr is expected. // If id was non-zero and create_node returned nullptr, it indicates a deserialization error. if (graph->nodes[i] == nullptr && id != 0) { - GGML_PRINT_DEBUG("[%s] failed to create graph node %d (id=%" PRId64 ")\n", __func__, i, id); + GGML_LOG_ERROR("[%s] failed to create graph node %d (id=%" PRId64 ")\n", __func__, i, id); return false; } } - ggml_status status = ggml_backend_graph_compute(backend, graph); - response.result = status; - ggml_free(ctx); + ggml_status status = ggml_backend_graph_compute(backends[device], graph); + GGML_ASSERT(status == GGML_STATUS_SUCCESS && "Unsuccessful graph computations are not supported with RPC"); + stored_graphs[device].ctx_ptr.swap(ctx_ptr); + stored_graphs[device].graph = graph; return true; } + +bool rpc_server::graph_recompute(const rpc_msg_graph_recompute_req & request) { + uint32_t device = request.device; + if (device >= backends.size()) { + return false; + } + if (stored_graphs[device].graph == nullptr) { + return false; + } + ggml_cgraph * graph = stored_graphs[device].graph; + LOG_DBG("[%s] device: %u\n", __func__, device); + ggml_status status = ggml_backend_graph_compute(backends[device], graph); + GGML_ASSERT(status == GGML_STATUS_SUCCESS && "Unsuccessful graph computations are not supported with RPC"); + return true; +} + + rpc_server::~rpc_server() { for (auto buffer : buffers) { ggml_backend_buffer_free(buffer); } } -static void rpc_serve_client(ggml_backend_t backend, const char* cache_dir, - sockfd_t sockfd, size_t free_mem, size_t total_mem) { - rpc_server server(backend, cache_dir); +static void rpc_serve_client(const std::vector& backends, const char* cache_dir, + sockfd_t sockfd, const std::vector& free_mem, const std::vector& total_mem) { + rpc_server server(backends, cache_dir); uint8_t cmd; if (!recv_data(sockfd, &cmd, 1)) { return; @@ -1479,13 +1648,26 @@ static void rpc_serve_client(ggml_backend_t backend, const char* cache_dir, // HELLO command is handled above return; } + case RPC_CMD_DEVICE_COUNT: { + if (!recv_msg(sockfd, nullptr, 0)) { + return; + } + rpc_msg_device_count_rsp response; + response.device_count = backends.size(); + if (!send_msg(sockfd, &response, sizeof(response))) { + return; + } + break; + } case RPC_CMD_ALLOC_BUFFER: { rpc_msg_alloc_buffer_req request; if (!recv_msg(sockfd, &request, sizeof(request))) { return; } rpc_msg_alloc_buffer_rsp response; - server.alloc_buffer(request, response); + if (!server.alloc_buffer(request, response)) { + return; + } if (!send_msg(sockfd, &response, sizeof(response))) { return; } @@ -1504,22 +1686,28 @@ static void rpc_serve_client(ggml_backend_t backend, const char* cache_dir, break; } case RPC_CMD_GET_ALIGNMENT: { - if (!recv_msg(sockfd, nullptr, 0)) { + rpc_msg_get_alignment_req request; + if (!recv_msg(sockfd, &request, sizeof(request))) { return; } rpc_msg_get_alignment_rsp response; - server.get_alignment(response); + if (!server.get_alignment(request, response)) { + return; + } if (!send_msg(sockfd, &response, sizeof(response))) { return; } break; } case RPC_CMD_GET_MAX_SIZE: { - if (!recv_msg(sockfd, nullptr, 0)) { + rpc_msg_get_max_size_req request; + if (!recv_msg(sockfd, &request, sizeof(request))) { return; } rpc_msg_get_max_size_rsp response; - server.get_max_size(response); + if (!server.get_max_size(request, response)) { + return; + } if (!send_msg(sockfd, &response, sizeof(response))) { return; } @@ -1576,12 +1764,12 @@ static void rpc_serve_client(ggml_backend_t backend, const char* cache_dir, break; } case RPC_CMD_SET_TENSOR_HASH: { - std::vector input; - if (!recv_msg(sockfd, input)) { + rpc_msg_set_tensor_hash_req request; + if (!recv_msg(sockfd, &request, sizeof(request))) { return; } rpc_msg_set_tensor_hash_rsp response; - if (!server.set_tensor_hash(input, response)) { + if (!server.set_tensor_hash(request, response)) { return; } if (!send_msg(sockfd, &response, sizeof(response))) { @@ -1635,22 +1823,25 @@ static void rpc_serve_client(ggml_backend_t backend, const char* cache_dir, if (!recv_msg(sockfd, input)) { return; } - rpc_msg_graph_compute_rsp response; - if (!server.graph_compute(input, response)) { - return; - } - if (!send_msg(sockfd, &response, sizeof(response))) { + if (!server.graph_compute(input)) { return; } break; } case RPC_CMD_GET_DEVICE_MEMORY: { - if (!recv_msg(sockfd, nullptr, 0)) { + rpc_msg_get_device_memory_req request; + if (!recv_msg(sockfd, &request, sizeof(request))) { + return; + } + auto dev_id = request.device; + if (dev_id >= backends.size()) { return; } rpc_msg_get_device_memory_rsp response; - response.free_mem = free_mem; - response.total_mem = total_mem; + response.free_mem = free_mem[dev_id]; + response.total_mem = total_mem[dev_id]; + LOG_DBG("[get_device_mem] device: %u, free_mem: %" PRIu64 ", total_mem: %" PRIu64 "\n", dev_id, + response.free_mem, response.total_mem); if (!send_msg(sockfd, &response, sizeof(response))) { return; } @@ -1665,9 +1856,32 @@ static void rpc_serve_client(ggml_backend_t backend, const char* cache_dir, } -void ggml_backend_rpc_start_server(ggml_backend_t backend, const char* endpoint, +GGML_API GGML_CALL void ggml_backend_rpc_start_server(const char* endpoint, const char* cache_dir, - size_t free_mem, size_t total_mem) { + size_t n_devices, ggml_backend_t * devices, + size_t * free_mem, size_t * total_mem) { + + if (n_devices == 0 || devices == nullptr || free_mem == nullptr || total_mem == nullptr) { + fprintf(stderr, "Invalid arguments to ggml_backend_rpc_start_server\n"); + return; + } + std::vector backends; + std::vector free_mem_vec(free_mem, free_mem + n_devices); + std::vector total_mem_vec(total_mem, total_mem + n_devices); + printf("Starting RPC server v%d.%d.%d\n", + RPC_PROTO_MAJOR_VERSION, + RPC_PROTO_MINOR_VERSION, + RPC_PROTO_PATCH_VERSION); + printf(" endpoint : %s\n", endpoint); + printf(" local cache : %s\n", cache_dir ? cache_dir : "n/a"); + printf("Using devices:\n"); + for (size_t i = 0; i < n_devices; i++) { + auto dev = devices[i]; + backends.push_back(dev); + const char* name = ggml_backend_name(devices[i]); + printf(" %8s: %10zu MiB total, %10zu MiB free\n", name, + total_mem_vec[i] / 1024 / 1024, free_mem_vec[i] / 1024 / 1024); + } std::string host; int port; if (!parse_endpoint(endpoint, host, port)) { @@ -1694,13 +1908,24 @@ void ggml_backend_rpc_start_server(ggml_backend_t backend, const char* endpoint, fprintf(stderr, "Failed to accept client connection\n"); return; } - printf("Accepted client connection, free_mem=%zu, total_mem=%zu\n", free_mem, total_mem); + printf("Accepted client connection\n"); fflush(stdout); - rpc_serve_client(backend, cache_dir, client_socket->fd, free_mem, total_mem); + rpc_serve_client(backends, cache_dir, client_socket->fd, free_mem_vec, total_mem_vec); printf("Client connection closed\n"); fflush(stdout); } #ifdef _WIN32 WSACleanup(); #endif + for (auto backend : backends) { + ggml_backend_free(backend); + } +} + +GGML_API GGML_CALL uint32_t ggml_backend_rpc_get_device_count(const char* endpoint) { + auto sock = get_socket(endpoint); + rpc_msg_device_count_rsp response; + bool status = send_rpc_cmd(sock, RPC_CMD_DEVICE_COUNT, nullptr, 0, &response, sizeof(response)); + RPC_STATUS_ASSERT(status); + return response.device_count; } diff --git a/src/llama-model.h b/src/llama-model.h index 769f0513..342d6df6 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -271,6 +271,11 @@ struct llama_layer { struct llama_lora_adapter; +struct rpc_device { + std::string endpoint; + uint32_t device; +}; + struct llama_model { e_model type = MODEL_UNKNOWN; llm_arch arch = LLM_ARCH_UNKNOWN; @@ -299,7 +304,7 @@ struct llama_model { int main_gpu; int n_gpu_layers; - std::vector rpc_servers; + std::vector rpc_servers; std::vector devices; // gguf metadata diff --git a/src/llama.cpp b/src/llama.cpp index 7b1b4e69..587ff62f 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -125,19 +125,6 @@ // helpers // -// trim whitespace from the beginning and end of a string -//static std::string trim(const std::string & str) { -// Fails for Chinese character -// size_t start = 0; -// size_t end = str.size(); -// while (start < end && isspace(str[start])) { -// start += 1; -// } -// while (end > start && isspace(str[end - 1])) { -// end -= 1; -// } -// return str.substr(start, end - start); -//} static bool is_utf8_whitespace(uint8_t c) { // Basic ASCII whitespace @@ -155,38 +142,35 @@ static std::string trim(const std::string & str) { } -static std::vector llama_string_split(const std::string& str, const std::string& delimiter) { +static std::vector string_split(const std::string& str, const std::string& delimiter) { std::vector parts; size_t start = 0; size_t end = str.find(delimiter); - while (end != std::string::npos) { parts.push_back(str.substr(start, end - start)); start = end + delimiter.length(); end = str.find(delimiter, start); } - parts.push_back(str.substr(start)); - return parts; - } // extract ip and port from RPC[ip:port] for rpc and keep other device names -static std::vector extract_ip_from_rpc_device(std::vector devices) { - std::vector rpc_servers; - std::regex pattern("RPC\\[(.*?)\\]"); - std::smatch matches; - for (auto device : devices) { - if (std::regex_search(device, matches, pattern)) { - rpc_servers.push_back(matches[1]); - } else { - rpc_servers.push_back(device); +static std::vector extract_device_from_rpc_device(std::vector devices) { + std::vector rpc_servers; + for (auto & device : devices) { + rpc_device rpc; + auto value = string_split(device, "|"); + if (value.size() == 2) { + rpc.device = std::stoi(value[1]); + rpc.endpoint = value[0]; } + rpc_servers.push_back(rpc); } return rpc_servers; } + enum llm_chat_template { LLM_CHAT_TEMPLATE_CHATML, LLM_CHAT_TEMPLATE_LLAMA_2, @@ -445,8 +429,10 @@ static ggml_backend_buffer_type_t llama_default_buffer_type_offload(const llama_ int dev_count = (int)llama_get_device_count(model); int rpc_count = (int)model.rpc_servers.size(); if (gpu >= dev_count - rpc_count) { - const char * endpoint = model.rpc_servers[gpu - dev_count + rpc_count].c_str(); - return ggml_backend_rpc_buffer_type(endpoint); + int rpc_idx = gpu - dev_count + rpc_count; + rpc_device rpc = model.rpc_servers[rpc_idx]; + const char * endpoint = rpc.endpoint.c_str(); + return ggml_backend_rpc_buffer_type(endpoint, rpc.device); } #endif #if defined(GGML_USE_METAL) @@ -504,8 +490,9 @@ static size_t llama_get_device_memory(const llama_model & model, int device) { if (device >= dev_count - rpc_count) { size_t total; size_t free; - const char * endpoint = model.rpc_servers[device - dev_count + rpc_count].c_str(); - ggml_backend_rpc_get_device_memory(endpoint, &free, &total); + rpc_device rpc = model.rpc_servers[device - dev_count + rpc_count]; + const char * endpoint = rpc.endpoint.c_str(); + ggml_backend_rpc_get_device_memory(endpoint, rpc.device, &free, &total); return free; } #endif @@ -1694,11 +1681,23 @@ static bool llm_load_tensors( int act_gpu_layers = std::min(n_gpu_layers, (int)n_layer + 1); for (int i = i_gpu_start; i < n_layer; ++i) { int layer_gpu = std::upper_bound(splits.begin(), splits.begin() + device_count, float(i - i_gpu_start)/act_gpu_layers) - splits.begin(); +#ifndef NDEBUG + ggml_backend_buffer_type_t buft = llama_default_buffer_type_offload(model, model.devices[layer_gpu]); + const char* name = ggml_backend_buft_name(buft); + LLAMA_LOG_DEBUG("load_tensors: layers %3d assigned to backend %s\n", i, + name); +#endif model.buft_layer[i] = llama_default_buffer_type_offload(model, model.devices[layer_gpu]); } // assign the output layer if (n_gpu_layers > n_layer) { int layer_gpu = std::upper_bound(splits.begin(), splits.begin() + device_count, float(act_gpu_layers - 1)/act_gpu_layers) - splits.begin(); +#ifndef NDEBUG + ggml_backend_buffer_type_t buft = llama_default_buffer_type_offload(model, model.devices[layer_gpu]); + const char* name = ggml_backend_buft_name(buft); + LLAMA_LOG_DEBUG("load_tensors: output layers assigned to backend %s\n", + name); +#endif model.buft_output = llama_default_buffer_type_offload(model, model.devices[layer_gpu]); } else { model.buft_output = llama_default_buffer_type_cpu(true); @@ -4016,17 +4015,11 @@ int64_t llama_time_us(void) { return ggml_time_us(); } -static int32_t find_device_idx(const std::string& str) { - std::regex pattern(R"((\d+)$)"); // Match digits at the end - std::smatch matches; - int number = -1; - if (std::regex_search(str, matches, pattern)) { - number = std::stoi(matches[1]); - } - return number; +static std::string create_rpc_name(std::string endpoint, uint32_t device) { + std::string dev_name = "RPC" + std::to_string(device) + "[" + std::string(endpoint) + "]"; + return dev_name; } - struct llama_model * llama_load_model_from_file( const char * path_model, struct llama_model_params params) { @@ -4058,8 +4051,7 @@ struct llama_model * llama_load_model_from_file( std::vector params_devices; if (params.devices && !striequals(params.devices, "")) { - params_devices = llama_string_split(params.devices, ","); - params_devices = extract_ip_from_rpc_device(params_devices); + params_devices = string_split(params.devices, ","); } std::map buffer_names; @@ -4075,20 +4067,21 @@ struct llama_model * llama_load_model_from_file( gpu_names.push_back(std::string(name)); } if (has_rpc) { - model->rpc_servers = llama_string_split(params.rpc_servers, ","); + model->rpc_servers = extract_device_from_rpc_device(string_split(params.rpc_servers, ",")); for (auto rpc : model->rpc_servers) { - buffer_names.insert({ rpc, idx}); + buffer_names.insert({ create_rpc_name(rpc.endpoint, rpc.device), idx}); idx++; } } std::vector device_names; if (params_devices.size()) { device_names = params_devices; - } - else { + } else { // add RPC servers at the front of the list to minimize the network transfers if (has_rpc) { - device_names = model->rpc_servers; + for (auto& it : model->rpc_servers) { + device_names.push_back(create_rpc_name(it.endpoint, it.device)); + } } device_names.insert(device_names.end(), gpu_names.begin(), gpu_names.end()); } @@ -4096,8 +4089,7 @@ struct llama_model * llama_load_model_from_file( for (auto & device : device_names) { if (buffer_names.count(device)) { model->devices.push_back(buffer_names[device]); - } - else { + } else { LLAMA_LOG_ERROR("%s backend not available.\n", device.c_str()); } } @@ -4451,10 +4443,11 @@ struct llama_context * llama_new_context_with_model( #if defined(GGML_USE_RPC) if (model->n_gpu_layers > 0) { - for (const auto & endpoint : model->rpc_servers) { - ggml_backend_t backend = ggml_backend_rpc_init(endpoint.c_str()); + for (const auto & device : model->rpc_servers) { + ggml_backend_t backend = ggml_backend_rpc_init(device.endpoint.c_str(), device.device); if (backend == nullptr) { - LLAMA_LOG_ERROR("%s: failed to initialize RPC to '%s'\n", __func__, endpoint.c_str()); + LLAMA_LOG_ERROR("%s: failed to initialize RPC%d to '%s'\n", __func__, device.device, + device.endpoint.c_str()); llama_free(ctx); return nullptr; }