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

docs: clarify split loading usage

Browse Source
This commit is contained in:
Joseph Liaw
2025-06-19 19:55:56 +08:00
parent 1e7ae71ce5
commit 6a898ef9d8
6 changed files with 91 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+42
View File
@@ -222,6 +222,48 @@ 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).
### Load Multiple GGUF Splits
`gguf-split` can divide a model into several files, e.g., running
```shell
./gguf-split --split-max-tensors 128 llama.gguf llama.gguf
```
creates files `llama.gguf-00001-of-00004.gguf` … `llama.gguf-00004-of-00004.gguf`.
To load such a split model, pass the first split (e.g., `llama.gguf-00001-of-00004.gguf`) via `-m` (used only for metadata) and list *all*
required splits using `--splits` (indexes start at zero).
Split indexes can be assigned arbitrarily to ranks, and a rank may load multiple splits or share splits with others.
Examples:
- **4 splits, 2 ranks**
```shell
# rank 0 loads splits 0 and 1
./llama-cli -m llama.gguf-00001-of-00004.gguf --splits 0,1 \
--world 2 --rank 0
# rank 1 loads splits 2 and 3
./llama-cli -m llama.gguf-00001-of-00004.gguf --splits 2,3 \
--world 2 --rank 1
```
- **2 splits, 4 ranks**
```shell
# ranks 0 and 1 load split 0
./llama-cli -m model.gguf-00001-of-00002.gguf --splits 0 --world 4 --rank 0
./llama-cli -m model.gguf-00001-of-00002.gguf --splits 0 --world 4 --rank 1
# ranks 2 and 3 load split 1
./llama-cli -m model.gguf-00001-of-00002.gguf --splits 1 --world 4 --rank 2
./llama-cli -m model.gguf-00001-of-00002.gguf --splits 1 --world 4 --rank 3
```
Ensure `--world` and `--n-layer-window` (or `--n-gpu-layers`) are set so each
rank only processes its assigned layers. The loader reports an error if a needed
tensor is missing from the provided splits.
This split-loading mechanism only checks that every rank receives the tensors
for the layers it will run. For reliable control over the mapping of layers to
ranks, explicitly set `-lw`/`--n-layer-window` so you know exactly which split
IDs to include in `--splits`.
### (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
+7
View File
@@ -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
+1
View File
@@ -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
+1
View File
@@ -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
+2
View File
@@ -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
+38 -25
View File
@@ -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,21 @@ 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[0] != '\0') {
for (int s : string_split<int>(split_str, ',')) {
if (s >= 0) split_list.insert((uint32_t)s);
}
}
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 +4795,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 +4819,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 +4847,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 +9561,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 +19648,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 +20222,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,