From: Meshnet Subject: [PATCH] llama: add dense owned-range loading seam diff --git a/include/llama.h b/include/llama.h index a311ac20..1f9459cf 100644 --- a/include/llama.h +++ b/include/llama.h @@ -292,6 +292,19 @@ extern "C" { ggml_backend_buffer_type_t buft; }; + // Immutable report for the project-owned dense-Llama range-loading seam. + // The bounds are inclusive/exclusive and are populated only after the + // model has registered and allocated its owned tensors. + struct llama_meshnet_range_report { + int32_t start_layer; + int32_t end_layer; + uint64_t mapped_bytes; + uint64_t resident_bytes; + uint64_t registered_bytes; + bool has_token_embeddings; + bool has_output_head; + }; + struct llama_model_params { @@ -319,6 +332,12 @@ extern "C" { const struct llama_model_kv_override * kv_overrides; + int32_t meshnet_owned_layer_start; + int32_t meshnet_owned_layer_end; + // Keep the booleans together to avoid misalignment during copy-by-value. @@ -616,6 +635,13 @@ extern "C" { LLAMA_API uint64_t llama_model_size(const struct llama_model * model); + LLAMA_API bool llama_model_meshnet_range_report( + const struct llama_model * model, + struct llama_meshnet_range_report * out); + // Get the default chat template. Returns nullptr if not available diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d8748138..4d2a3ec1 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1015,6 +1015,9 @@ struct llama_model::impl { std::vector tensor_split_owned; + llama_meshnet_range_report meshnet_range_report = {}; + bool has_meshnet_range_report = false; }; @@ -1236,6 +1239,19 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { const bool use_mmap_buffer = true; + const bool meshnet_range_requested = params.meshnet_owned_layer_start != 0 || params.meshnet_owned_layer_end != 0; + const int meshnet_start = params.meshnet_owned_layer_start; + const int meshnet_end = params.meshnet_owned_layer_end; + if (meshnet_range_requested) { + if (arch != LLM_ARCH_LLAMA) { + throw std::runtime_error("Meshnet owned range currently supports dense Llama only"); + } + if (meshnet_start < 0 || meshnet_end <= meshnet_start || meshnet_end > static_cast(hparams.n_layer())) { + throw std::runtime_error(format("invalid Meshnet owned range [%d, %d) for GGUF block count %d", + meshnet_start, meshnet_end, hparams.n_layer())); + } + } @@ -1336,7 +1352,9 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { - for (int i = 0; i < n_layer_all; ++i) { + const int optional_scale_start = meshnet_range_requested ? meshnet_start : 0; + const int optional_scale_end = meshnet_range_requested ? meshnet_end : n_layer_all; + for (int i = optional_scale_start; i < optional_scale_end; ++i) { @@ -1487,7 +1505,7 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { - ml.done_getting_tensors(); + ml.done_getting_tensors(meshnet_range_requested); @@ -1613,8 +1631,11 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { + uint64_t meshnet_mapped_bytes = 0; + uint64_t meshnet_resident_bytes = 0; for (auto & [_, bufs] : pimpl->ctxs_bufs) { for (auto & buf: bufs) { + meshnet_resident_bytes += ggml_backend_buffer_get_size(buf.get()); @@ -1637,6 +1658,35 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } + if (meshnet_range_requested) { + uint64_t registered_bytes = 0; + for (const auto & [_, tensor] : tensors_by_name) registered_bytes += ggml_nbytes(tensor); + if (ml.use_mmap) for (const auto & [first, last] : ml.mmaps_used) if (last > first) meshnet_mapped_bytes += last - first; + const auto registered = [this](const ggml_tensor * tensor) { + return tensor != nullptr && std::any_of(tensors_by_name.begin(), tensors_by_name.end(), + [tensor](const auto & entry) { return entry.second == tensor; }); + }; + const auto registered_name = [this](const char * name) { + return std::any_of(tensors_by_name.begin(), tensors_by_name.end(), + [name](const auto & entry) { return entry.first == name; }); + }; + pimpl->meshnet_range_report = { meshnet_start, meshnet_end, meshnet_mapped_bytes, meshnet_resident_bytes, + registered_bytes, registered_name("token_embd.weight"), registered(output_norm) && registered(output) }; + pimpl->has_meshnet_range_report = true; + } return true; @@ -1711,6 +1761,14 @@ uint64_t llama_model::n_elements() const { } +bool llama_model::meshnet_range_report(llama_meshnet_range_report * out) const { + if (out == nullptr || !pimpl->has_meshnet_range_report) return false; + *out = pimpl->meshnet_range_report; + return true; +} @@ -2308,6 +2366,8 @@ llama_model_params llama_model_default_params() { /*.kv_overrides =*/ nullptr, + /*.meshnet_owned_layer_start =*/ 0, + /*.meshnet_owned_layer_end =*/ 0, @@ -2641,6 +2701,10 @@ uint64_t llama_model_size(const llama_model * model) { } +bool llama_model_meshnet_range_report(const llama_model * model, llama_meshnet_range_report * out) { + return model != nullptr && model->meshnet_range_report(out); +} diff --git a/src/llama-model.h b/src/llama-model.h index 45b054ce..1b3f9bd0 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -652,6 +652,8 @@ struct llama_model { + bool meshnet_range_report(llama_meshnet_range_report * out) const; + diff --git a/src/models/llama.cpp b/src/models/llama.cpp index 4bfebc88..b4f25aed 100644 --- a/src/models/llama.cpp +++ b/src/models/llama.cpp @@ -34,18 +34,26 @@ void llama_model_llama::load_arch_hparams(llama_model_loader & ml) { - tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + const bool meshnet_range_requested = params.meshnet_owned_layer_start != 0 || params.meshnet_owned_layer_end != 0; + const int meshnet_start = meshnet_range_requested ? params.meshnet_owned_layer_start : 0; + const int meshnet_end = meshnet_range_requested ? params.meshnet_owned_layer_end : n_layer; - // output - output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); - output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); + if (!meshnet_range_requested || meshnet_start == 0) tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + if (!meshnet_range_requested || meshnet_end == n_layer) { + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); - // if output is NULL, init from the input tok embed - if (output == NULL) { - output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + if (output == NULL) output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); } - for (int i = 0; i < n_layer; ++i) { + for (int i = meshnet_start; i < meshnet_end; ++i) { @@ -102,6 +110,25 @@ llama_model_llama::graph::graph(const llama_model & model, const llm_grap + llama_meshnet_range_report meshnet_report = {}; + if (model.meshnet_range_report(&meshnet_report)) { + if (meshnet_report.start_layer != 0) throw std::runtime_error("Meshnet dense-Llama graph requires a head endpoint adapter"); + if (meshnet_report.end_layer != n_layer) throw std::runtime_error("Meshnet dense-Llama graph requires a tail endpoint adapter"); + if (!meshnet_report.has_token_embeddings) throw std::runtime_error("Meshnet dense-Llama head range is missing token embeddings"); + if (!meshnet_report.has_output_head) throw std::runtime_error("Meshnet dense-Llama tail range is missing final norm or output head"); + } ggml_tensor * cur; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 855295c1..9a7be6ee 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -193,6 +193,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) + llama_build_and_test(test-meshnet-range-ownership.cpp) diff --git a/tests/test-meshnet-range-ownership.cpp b/tests/test-meshnet-range-ownership.cpp new file mode 100644 index 00000000..7b58ebf8 --- /dev/null +++ b/tests/test-meshnet-range-ownership.cpp @@ -0,0 +1,6 @@ +#include "ggml.h" +#include "gguf.h" +#include "llama.h" +#include "../src/llama-model.h" +#include +#include