Merge pull request #3 from DandinPower/feat/enable-rank-specific-gguf-file-loading

Feat/enable rank specific gguf file loading

README.md

@@ -222,6 +222,47 @@ 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).
 
+### Effecient Multi-Split GGUF Support in `prima.cpp`
+
+If your model is split into multiple GGUF files using `gguf-split`, for example by running:
+
+```shell
+./gguf-split --split-max-tensors 128 llama.gguf llama.gguf
+```
+
+You may end up with files like:
+
+```
+llama.gguf-00001-of-00004.gguf
+llama.gguf-00002-of-00004.gguf
+llama.gguf-00003-of-00004.gguf
+llama.gguf-00004-of-00004.gguf
+```
+
+Unlike `llama.cpp`, which typically requires loading all splits, `prima.cpp` only needs to load the GGUF splits that contain the weights for the layers assigned to the current rank. Downloading or loading all splits is unnecessary and inefficient. Please note that prima.cpp enforces a layer assignment policy where the first and last layers are always allocated to the master rank. 
+
+* **Rank 0** (the master node) always owns the **embedding-related layers** and the **last portion** of transformer blocks.
+* **Rank 1** (the first non-master node) always starts from **layer 0**.
+
+For example, if `--n-layer-window` (`-lw`) is set to `"8,8"` and `--world` is `2`, the layer distribution is:
+
+* Rank 0 (master) owns: **embedding related layers** and **layers 8–15**
+* Rank 1 (server) owns: **layers 0–7**
+
+Assuming the model is split into 4 parts, the proper usage of `--splits` would be:
+
+#### Example: 4 Splits, 2 Ranks
+
+```bash
+# Rank 0 (master) loads splits 0, 2 and 3
+./llama-cli -m llama.gguf-00001-of-00004.gguf --splits 0,2,3 --world 2 --rank 0 <OTHER FLAGS>
+
+# Rank 1 (server) loads splits 0 and 1
+./llama-cli -m llama.gguf-00001-of-00004.gguf --splits 0,1 --world 2 --rank 1 <OTHER FLAGS>
+```
+
+If a required tensor is missing from the specified splits, the loader will raise an error. To ensure correct mapping and avoid runtime errors, it is recommended to  **explicitly set `--n-layer-window`** so that you can precisely determine which layer belongs to which rank, and in turn which split files to include in the `--splits` flag.
+
 ### (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:
 

common/arg.cpp

@@ -1510,6 +1510,13 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
         }
     ).set_env("LLAMA_ARG_RPC"));
 #endif
+    add_opt(llama_arg(
+        {"--splits"}, "LIST",
+        "comma separated list of GGUF split indexes to load (add 0 to load tensors from the first split)",
+        [](gpt_params & params, const std::string & value) {
+            params.gguf_splits = value;
+        }
+    ));
     add_opt(llama_arg(
         {"--mlock"},
         "force system to keep model in RAM rather than swapping or compressing",

common/common.cpp

@@ -1943,6 +1943,7 @@ struct llama_model_params llama_model_params_from_gpt_params(const gpt_params &
     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;

common/common.h

@@ -217,6 +217,7 @@ struct gpt_params {
     std::string lookup_cache_dynamic = ""; // path of dynamic ngram cache file for lookup decoding          // NOLINT
     std::string logits_file          = ""; // file for saving *all* logits                                  // NOLINT
     std::string rpc_servers          = ""; // comma separated list of RPC servers                           // NOLINT
+    std::string gguf_splits          = ""; // comma separated list of GGUF split indexes                 // NOLINT
 
     std::vector<std::string> in_files;   // all input files
     std::vector<std::string> antiprompt; // strings upon which more user input is prompted (a.k.a. reverse prompts)

include/llama.h

@@ -297,6 +297,8 @@ extern "C" {
 
         // comma separated list of RPC servers to use for offloading
         const char * rpc_servers;
+        // comma separated list of GGUF split indexes to load; include 0 if the first split's tensors are needed
+        const char * gguf_splits;
 
         // Called with a progress value between 0.0 and 1.0. Pass NULL to disable.
         // If the provided progress_callback returns true, model loading continues.

src/llama.cpp

@@ -4722,6 +4722,9 @@ struct llama_model_loader {
     llama_fver  fver;
 
     llama_mmaps mappings;
+    uint32_t my_rank = 0;
+    uint32_t n_world = 1;
+    std::set<uint32_t> split_list;
 
     // Holds information on a model weight
     struct llama_tensor_weight {
@@ -4754,12 +4757,29 @@ struct llama_model_loader {
     std::string arch_name;
     LLM_KV      llm_kv    = LLM_KV(LLM_ARCH_UNKNOWN);
 
-    llama_model_loader(const std::string & fname, bool use_mmap, bool check_tensors, const struct llama_model_kv_override * param_overrides_p) {
+    llama_model_loader(const std::string & fname, bool use_mmap, bool check_tensors, const struct llama_model_kv_override * param_overrides_p,
+                       uint32_t rank = 0, uint32_t world = 1, const char * split_str = nullptr) {
+        my_rank = rank;
+        n_world = world;
         int trace = 0;
         if (getenv("LLAMA_TRACE")) {
             trace = atoi(getenv("LLAMA_TRACE"));
         }
 
+        if (split_str && *split_str) {
+            std::string tmp{split_str};
+            std::stringstream ss(tmp);
+            std::string tok;
+            while (std::getline(ss, tok, ',')) {
+                try {
+                    auto v = std::stoi(tok);
+                    if (v >= 0) split_list.insert(static_cast<uint32_t>(v));
+                } catch (...) {
+                    // ignore bad tokens
+                }
+            }
+        }
+
         if (param_overrides_p != nullptr) {
             for (const struct llama_model_kv_override * p = param_overrides_p; p->key[0] != 0; p++) {
                 kv_overrides.insert({std::string(p->key), *p});
@@ -4783,21 +4803,23 @@ struct llama_model_loader {
         files.emplace_back(new llama_file(fname.c_str(), "rb"));
         contexts.emplace_back(ctx);
 
-        // Save tensors data offset of the main file.
+        // Save tensors data offset of the main file if this split is requested
         // For subsidiary files, `meta` tensor data offset must not be used,
         // so we build a unified tensors index for weights.
-        for (ggml_tensor * cur = ggml_get_first_tensor(ctx); cur; cur = ggml_get_next_tensor(ctx, cur)) {
-            weights.emplace_back(files.back().get(), 0, cur->name, meta, cur);
+        if (split_list.empty() || split_list.count(0) > 0) {
+            for (ggml_tensor * cur = ggml_get_first_tensor(ctx); cur; cur = ggml_get_next_tensor(ctx, cur)) {
+                weights.emplace_back(files.back().get(), 0, cur->name, meta, cur);
+            }
         }
         uint16_t n_split = 0;
         get_key(llm_kv(LLM_KV_SPLIT_COUNT), n_split, false);
 
-        // Load additional GGML contexts
+        // Split file verification and loading additional splits
         if (n_split > 1) {
             uint16_t idx = 0;
             get_key(llm_kv(LLM_KV_SPLIT_NO), idx);
             if (idx != 0) {
-                throw std::runtime_error(format("illegal split file: %d, model must be loaded with the first split", idx));
+                throw std::runtime_error(format("illegal split file: %d, first split expected", idx));
             }
 
             char split_prefix[PATH_MAX] = {0};
@@ -4805,13 +4827,21 @@ struct llama_model_loader {
                 throw std::runtime_error(format("invalid split file: %s", fname.c_str()));
             }
 
-            if (trace > 0) {
-                LLAMA_LOG_INFO("%s: loading additional %d GGUFs\n", __func__, n_split);
+            if (split_list.empty()) {
+                for (idx = 0; idx < n_split; ++idx) {
+                    split_list.insert(idx);
+                }
             }
 
             char split_path[PATH_MAX] = {0};
-            for (idx = 1; idx < n_split; idx++) {
-                llama_split_path(split_path, sizeof(split_path), split_prefix, idx, n_split);
+            for (uint16_t sid : split_list) {
+                if (sid == 0) {
+                    continue;
+                }
+                if (sid >= n_split) {
+                    throw std::runtime_error(format("split index %d out of range", sid));
+                }
+                llama_split_path(split_path, sizeof(split_path), split_prefix, sid, n_split);
 
                 struct gguf_init_params split_params = {
                     /*.no_alloc = */ true,
@@ -4825,25 +4855,14 @@ struct llama_model_loader {
                 files.emplace_back(new llama_file(split_path, "rb"));
                 contexts.emplace_back(ctx);
 
-                // Save tensors data offset info of the shard.
                 for (ggml_tensor * cur = ggml_get_first_tensor(ctx); cur; cur = ggml_get_next_tensor(ctx, cur)) {
-                    weights.emplace_back(files.back().get(), idx, cur->name, ctx_gguf, cur);
+                    weights.emplace_back(files.back().get(), sid, cur->name, ctx_gguf, cur);
                 }
 
                 gguf_free(ctx_gguf);
             }
 
-            get_key(llm_kv(LLM_KV_SPLIT_TENSORS_COUNT), n_tensors);
-
-            // sanity check
-            {
-                const int n_tensors_loaded = (int) weights.size();
-                if (n_tensors != n_tensors_loaded) {
-                    throw std::runtime_error(format("corrupted model: %d tensors expected but %d found", n_tensors, n_tensors_loaded));
-                }
-            }
-
-            LLAMA_LOG_INFO("%s: additional %d GGUFs metadata loaded.\n",  __func__, n_split - 1);
+            get_key(llm_kv(LLM_KV_SPLIT_TENSORS_COUNT), n_tensors, false);
         }
 
         n_kv      = gguf_get_n_kv(meta);
@@ -9550,7 +9569,8 @@ int llm_load_tensors(
 // Returns 0 on success, -1 on error, and -2 on cancellation via llama_progress_callback
 static llama_model_loader * llama_model_load_impl(const std::string & fname, llama_model & model, llama_model_params & params) {
     try {
-        llama_model_loader * ml = new llama_model_loader(fname, params.use_mmap, params.check_tensors, params.kv_overrides);
+        llama_model_loader * ml = new llama_model_loader(fname, params.use_mmap, params.check_tensors, params.kv_overrides,
+                                                         params.rank, params.n_world, params.gguf_splits);
 
         model.hparams.vocab_only = params.vocab_only;
 
@@ -19636,7 +19656,7 @@ static void llama_model_quantize_internal(const std::string & fname_inp, const s
         auto v = (std::vector<llama_model_kv_override>*)params->kv_overrides;
         kv_overrides = v->data();
     }
-    llama_model_loader ml(fname_inp, use_mmap, /*check_tensors*/ true, kv_overrides);
+    llama_model_loader ml(fname_inp, use_mmap, /*check_tensors*/ true, kv_overrides, 0, 1, nullptr);
     ml.init_mappings(false); // no prefetching
 
     llama_model model;
@@ -20210,6 +20230,7 @@ struct llama_model_params llama_model_default_params() {
         /*.main_gpu                    =*/ 0,
         /*.tensor_split                =*/ nullptr,
         /*.rpc_servers                 =*/ nullptr,
+        /*.gguf_splits                 =*/ nullptr,
         /*.progress_callback           =*/ nullptr,
         /*.progress_callback_user_data =*/ nullptr,
         /*.kv_overrides                =*/ nullptr,