From a7cfd8864ff796d8271c7f851f43a371e0628873 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Fri, 28 Aug 2026 01:06:40 +0200 Subject: [PATCH] force lazy tensor on cpu if lazy is on --- src/llama-model-loader.cpp | 31 ++++++++++++++++++++++++------- 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index d9241022cf..39dfc85127 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1073,6 +1073,9 @@ static ggml_backend_buffer_type_t select_weight_buft(const llama_hparams & hpara struct ggml_tensor * llama_model_loader::create_tensor( const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output, const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags) { + // set below, before buft_for_tensor() runs + bool is_lazy = false; + auto ctx_for_buft = [&](ggml_backend_buffer_type_t buft) -> ggml_context * { auto it = ctx_map.find(buft); if (it == ctx_map.end()) { @@ -1160,6 +1163,16 @@ struct ggml_tensor * llama_model_loader::create_tensor( } } + // if the tensor is marked as "lazy", we always keep it on CPU no matter what + // use case: PLE / engrams embd tensors + if (is_lazy) { + auto * cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); + if (!cpu_dev) { + throw std::runtime_error("no CPU backend found"); + } + return ggml_backend_dev_buffer_type(cpu_dev); + } + // select the buffer type for this tensor const buft_list_t * buft_list; switch (info.layer) { @@ -1287,16 +1300,20 @@ struct ggml_tensor * llama_model_loader::create_tensor( return NULL; } - if ((flags & TENSOR_READ_LAZY) && use_mmap && tensor_read_lazy != LLAMA_TENSOR_READ_LAZY_OFF) { + if ((flags & TENSOR_READ_LAZY) && tensor_read_lazy != LLAMA_TENSOR_READ_LAZY_OFF) { // in auto mode, small tensors are cheap enough to keep resident constexpr size_t auto_lazy_min_size = 4ull * 1024 * 1024 * 1024; - if (tensor_read_lazy == LLAMA_TENSOR_READ_LAZY_ON || ggml_nbytes(cur) > auto_lazy_min_size) { - const auto & w = require_weight(tn.str().c_str()); - lazy_tensor_ranges[w.idx].emplace_back(w.offs, w.offs + ggml_nbytes(cur)); - LLAMA_LOG_INFO("%s: tensor %s (size = %zu MiB) lazy read enabled\n", - __func__, tn.str().c_str(), ggml_nbytes(cur)/1024/1024); - } + // note: this must not depend on use_mmap, or the memory-fit pass (which loads with no_alloc and no mmap) + is_lazy = tensor_read_lazy == LLAMA_TENSOR_READ_LAZY_ON || ggml_nbytes(cur) > auto_lazy_min_size; + } + + if (is_lazy && use_mmap) { + const auto & w = require_weight(tn.str().c_str()); + lazy_tensor_ranges[w.idx].emplace_back(w.offs, w.offs + ggml_nbytes(cur)); + + LLAMA_LOG_INFO("%s: tensor %s (size = %zu MiB) lazy read enabled\n", + __func__, tn.str().c_str(), ggml_nbytes(cur)/1024/1024); } ggml_tensor t_meta = *cur;