From: Meshnet Subject: [PATCH] llama: add dense owned-range loading seam diff --git a/include/llama.h b/include/llama.h index a311ac20..24a69978 100644 --- a/include/llama.h +++ b/include/llama.h @@ -292,6 +292,16 @@ 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; + }; + struct llama_model_params { // NULL-terminated list of devices to use for offloading (if NULL, all available devices are used) ggml_backend_dev_t * devices; @@ -319,6 +329,12 @@ extern "C" { // override key-value pairs of the model meta data const struct llama_model_kv_override * kv_overrides; + // Project-owned dense-Llama range. A zero/zero pair preserves stock + // whole-model loading. Any other pair is validated against GGUF's + // immutable block-count metadata before tensor registration. + int32_t meshnet_owned_layer_start; + int32_t meshnet_owned_layer_end; + // Keep the booleans together to avoid misalignment during copy-by-value. bool vocab_only; // only load the vocabulary, no weights bool use_mmap; // use mmap if possible @@ -616,6 +632,13 @@ extern "C" { // Returns the total size of all the tensors in the model in bytes LLAMA_API uint64_t llama_model_size(const struct llama_model * model); + // Returns false unless this model was instantiated through the Meshnet + // owned-range loader. Values are derived from registered buffers, never + // copied from caller-supplied parameters. + 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 // If name is NULL, returns the default chat template LLAMA_API const char * llama_model_chat_template(const struct llama_model * model, const char * name); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d8748138..5173279a 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1015,6 +1015,9 @@ struct llama_model::impl { bool has_tensor_overrides; std::vector tensor_split_owned; + + llama_meshnet_range_report meshnet_range_report = {}; + bool has_meshnet_range_report = false; }; llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique()) { @@ -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())); + } + } + this->ml = &ml; // to be used by create_tensor() and load_arch_tensors() LLAMA_LOG_INFO("%s: loading model tensors, this can take a while... (mmap = %s, direct_io = %s)\n", @@ -1336,7 +1352,9 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { // generic pass: load optional per-tensor/per-expert ".scale" tensors (e.g. NVFP4 scale2) // this avoids having to add scale loading to every architecture - 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) { auto & layer = layers[i]; // attention weight scales (per-tensor, shape {1}) @@ -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); // Tied NVFP4 output is valid when no separate LM-head scale tensors are present. // If sidecar scales exist, the output weight must be an actual output tensor. @@ -1613,13 +1631,32 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } // print memory requirements per buffer type + 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()); LLAMA_LOG_INFO("%s: %12s model buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf.get()), ggml_backend_buffer_get_size(buf.get()) / 1024.0 / 1024.0); } } + if (meshnet_range_requested) { + // mmap backend buffers exactly describe the mapped file spans. On + // non-mmap backends the instantiated allocation is the best exact + // resident measure and no file span is claimed as mapped. + if (ml.use_mmap) { + meshnet_mapped_bytes = meshnet_resident_bytes; + } + pimpl->meshnet_range_report = { + meshnet_start, + meshnet_end, + meshnet_mapped_bytes, + meshnet_resident_bytes, + }; + pimpl->has_meshnet_range_report = true; + } + if (ml.no_alloc) { return true; } @@ -1711,6 +1748,14 @@ uint64_t llama_model::n_elements() const { return pimpl->n_elements; } +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; +} + void llama_model::print_info() const { const std::string rope_scaling_type = llama_rope_scaling_type_name(hparams.rope_scaling_type_train); @@ -2308,6 +2353,8 @@ llama_model_params llama_model_default_params() { /*.progress_callback =*/ nullptr, /*.progress_callback_user_data =*/ nullptr, /*.kv_overrides =*/ nullptr, + /*.meshnet_owned_layer_start =*/ 0, + /*.meshnet_owned_layer_end =*/ 0, /*.vocab_only =*/ false, /*.use_mmap =*/ true, /*.use_direct_io =*/ false, @@ -2641,6 +2688,10 @@ uint64_t llama_model_size(const llama_model * model) { return model->size(); } +bool llama_model_meshnet_range_report(const llama_model * model, llama_meshnet_range_report * out) { + return model != nullptr && model->meshnet_range_report(out); +} + const char * llama_model_chat_template(const llama_model * model, const char * name) { const auto key = name ? LLM_KV(model->arch, name)(LLM_KV_TOKENIZER_CHAT_TEMPLATE) : LLM_KV(model->arch)(LLM_KV_TOKENIZER_CHAT_TEMPLATE); 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 { // total number of parameters in the model uint64_t n_elements() const; + bool meshnet_range_report(llama_meshnet_range_report * out) const; + void print_info() const; ggml_backend_dev_t dev_layer(int il) const; diff --git a/src/models/llama.cpp b/src/models/llama.cpp index 4bfebc88..68b08e4c 100644 --- a/src/models/llama.cpp +++ b/src/models/llama.cpp @@ -34,18 +34,24 @@ void llama_model_llama::load_arch_hparams(llama_model_loader & ml) { void llama_model_llama::load_arch_tensors(llama_model_loader &) { LLAMA_LOAD_LOCALS; - 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) { + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); - // 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); + // 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 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); + } } - for (int i = 0; i < n_layer; ++i) { + for (int i = meshnet_start; i < meshnet_end; ++i) { auto & layer = layers[i]; layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); 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-double-float.cpp) # SLOW llama_build_and_test(test-llama-archs.cpp) + llama_build_and_test(test-meshnet-range-ownership.cpp) endif() llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp) diff --git a/tests/test-meshnet-range-ownership.cpp b/tests/test-meshnet-range-ownership.cpp new file mode 100644 index 00000000..447496fc --- /dev/null +++ b/tests/test-meshnet-range-ownership.cpp @@ -0,0 +1,124 @@ +#include "ggml.h" +#include "gguf.h" +#include "llama.h" + +#include "../src/llama-model.h" + +#include +#include +#include +#include +#include + +namespace { + +constexpr int kLayers = 4; +constexpr int kEmbd = 8; +constexpr int kFfn = 16; +constexpr int kVocab = 16; + +void check(bool condition, const char * message) { + if (!condition) { + throw std::runtime_error(message); + } +} + +void add_tensor(gguf_context * gguf, ggml_context * tensors, const char * name, int d0, int d1 = 1) { + ggml_tensor * tensor = d1 == 1 + ? ggml_new_tensor_1d(tensors, GGML_TYPE_F32, d0) + : ggml_new_tensor_2d(tensors, GGML_TYPE_F32, d0, d1); + ggml_set_name(tensor, name); + std::memset(tensor->data, 0, ggml_nbytes(tensor)); + gguf_add_tensor(gguf, tensor); +} + +std::string write_fixture() { + const std::string path = "meshnet-dense-llama-range-fixture.gguf"; + gguf_context * gguf = gguf_init_empty(); + ggml_init_params params = { 128 * 1024, nullptr, false }; + ggml_context * tensors = ggml_init(params); + check(gguf && tensors, "failed to create dense-Llama fixture contexts"); + + gguf_set_val_str(gguf, "general.architecture", "llama"); + gguf_set_val_u32(gguf, "llama.context_length", 16); + gguf_set_val_u32(gguf, "llama.embedding_length", kEmbd); + gguf_set_val_u32(gguf, "llama.block_count", kLayers); + gguf_set_val_u32(gguf, "llama.feed_forward_length", kFfn); + gguf_set_val_u32(gguf, "llama.attention.head_count", 2); + gguf_set_val_u32(gguf, "llama.attention.head_count_kv", 2); + gguf_set_val_u32(gguf, "llama.rope.dimension_count", 4); + gguf_set_val_f32(gguf, "llama.attention.layer_norm_rms_epsilon", 1.0e-5f); + gguf_set_val_str(gguf, "tokenizer.ggml.model", "no_vocab"); + gguf_set_val_u32(gguf, "llama.vocab_size", kVocab); + + add_tensor(gguf, tensors, "token_embd.weight", kEmbd, kVocab); + add_tensor(gguf, tensors, "output_norm.weight", kEmbd); + add_tensor(gguf, tensors, "output.weight", kEmbd, kVocab); + for (int layer = 0; layer < kLayers; ++layer) { + const std::string p = "blk." + std::to_string(layer) + "."; + add_tensor(gguf, tensors, (p + "attn_norm.weight").c_str(), kEmbd); + add_tensor(gguf, tensors, (p + "attn_q.weight").c_str(), kEmbd, kEmbd); + add_tensor(gguf, tensors, (p + "attn_k.weight").c_str(), kEmbd, kEmbd); + add_tensor(gguf, tensors, (p + "attn_v.weight").c_str(), kEmbd, kEmbd); + add_tensor(gguf, tensors, (p + "attn_output.weight").c_str(), kEmbd, kEmbd); + add_tensor(gguf, tensors, (p + "ffn_norm.weight").c_str(), kEmbd); + add_tensor(gguf, tensors, (p + "ffn_gate.weight").c_str(), kEmbd, kFfn); + add_tensor(gguf, tensors, (p + "ffn_down.weight").c_str(), kFfn, kEmbd); + add_tensor(gguf, tensors, (p + "ffn_up.weight").c_str(), kEmbd, kFfn); + } + check(gguf_write_to_file(gguf, path.c_str(), false), "failed to write dense-Llama fixture"); + ggml_free(tensors); + gguf_free(gguf); + return path; +} + +int block_number(const std::string & name) { + int block = -1; + return std::sscanf(name.c_str(), "blk.%d.", &block) == 1 ? block : -1; +} + +llama_meshnet_range_report load_and_check(const std::string & path, int start, int end) { + llama_model_params params = llama_model_default_params(); + params.meshnet_owned_layer_start = start; + params.meshnet_owned_layer_end = end; + llama_model * model = llama_model_load_from_file(path.c_str(), params); + check(model != nullptr, "failed to load dense-Llama fixture"); + + llama_meshnet_range_report report = {}; + check(llama_model_meshnet_range_report(model, &report), "range report is absent"); + check(report.start_layer == start, "reported start does not match registered range"); + check(report.end_layer == end, "reported end does not match registered range"); + check(report.mapped_bytes > 0, "mmap report is empty"); + check(report.resident_bytes >= report.mapped_bytes, "resident bytes undercount mapped bytes"); + + const auto & tensors = llama_internal_get_tensor_map(model); + check(!tensors.empty(), "no tensors registered for owned range"); + for (const auto & [name, _] : tensors) { + const int block = block_number(name); + check(block >= start && block < end, "registered tensor is outside owned blk.N range"); + } + llama_model_free(model); + return report; +} + +} // namespace + +int main() { + llama_backend_init(); + const std::string fixture = write_fixture(); + + const auto head = load_and_check(fixture, 0, 1); + const auto middle = load_and_check(fixture, 1, 3); + const auto tail = load_and_check(fixture, 3, 4); + check(middle.mapped_bytes > head.mapped_bytes, "two-layer range did not map more bytes than head"); + check(middle.resident_bytes > tail.resident_bytes, "two-layer range did not allocate more bytes than tail"); + + llama_model_params invalid = llama_model_default_params(); + invalid.meshnet_owned_layer_start = 3; + invalid.meshnet_owned_layer_end = 5; + check(llama_model_load_from_file(fixture.c_str(), invalid) == nullptr, "invalid range loaded"); + + std::remove(fixture.c_str()); + llama_backend_free(); + return 0; +}