Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

llama: improve TENSOR_READ_LAZY handling (#27837)

Browse Source 
* force lazy tensor on cpu if lazy is on

* llama: improve TENSOR_READ_LAZY handling
This commit is contained in:
Xuan-Son Nguyen
2026-08-30 16:59:48 +02:00
committed by GitHub
parent f1793c1c4e
commit 2578138397
4 changed files with 118 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-model-loader.cpp
+61 -19
View File
@@ -1070,11 +1070,52 @@ static ggml_backend_buffer_type_t select_weight_buft(const llama_hparams & hpara
return nullptr;
}
ggml_backend_buffer_type_t llama_model_loader::lazy_read::buft() {
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);
}
bool llama_model_loader::lazy_read::add(const std::string & name, const ggml_tensor * t, const llama_tensor_weight * w) {
if (mode == LLAMA_LAZY_MODE_OFF) {
return false;
}
// do not lazy-read small tensors, it has significant overhead and is not worth it
constexpr size_t auto_min_size = 4ull * 1024 * 1024 * 1024;
if (mode != LLAMA_LAZY_MODE_ON && ggml_nbytes(t) <= auto_min_size) {
return false;
}
if (!llama_mmap::SUPPORTED) {
LLAMA_LOG_WARN("%s: mmap is not available, so tensor %s (size = %zu MiB) is loaded into RAM in full\n",
__func__, name.c_str(), ggml_nbytes(t)/1024/1024);
return false;
}
if (w) {
ranges[w->idx].emplace_back(w->offs, w->offs + ggml_nbytes(t));
tensors.insert(name);
LLAMA_LOG_INFO("%s: tensor %s (size = %zu MiB) lazy read enabled\n",
__func__, name.c_str(), ggml_nbytes(t)/1024/1024);
}
return true;
}
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<int64_t> & 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);
const ctx_key key { buft, is_lazy };
auto it = ctx_map.find(key);
if (it == ctx_map.end()) {
// one ggml context per buffer type
int max_n_tensors = n_tensors;
@@ -1096,7 +1137,7 @@ struct ggml_tensor * llama_model_loader::create_tensor(
throw std::runtime_error(format("failed to create ggml context"));
}
ctx_map.emplace(buft, ctx);
ctx_map.emplace(key, ctx);
return ctx;
}
@@ -1160,6 +1201,10 @@ struct ggml_tensor * llama_model_loader::create_tensor(
}
}
if (is_lazy) {
return lazy_read::buft();
}
// select the buffer type for this tensor
const buft_list_t * buft_list;
switch (info.layer) {
@@ -1287,16 +1332,9 @@ struct ggml_tensor * llama_model_loader::create_tensor(
return NULL;
}
if ((flags & TENSOR_READ_LAZY) && use_mmap && lazy_mode != LLAMA_LAZY_MODE_OFF) {
// in auto mode, small tensors are cheap enough to keep resident
constexpr size_t auto_lazy_min_size = 4ull * 1024 * 1024 * 1024;
if (lazy_mode == LLAMA_LAZY_MODE_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);
}
if (flags & TENSOR_READ_LAZY) {
// the decision must not depend on the load mode, or the memory-fit pass (no_alloc, no mmap)
is_lazy = lazy.add(tn.str(), cur, no_alloc ? nullptr : &require_weight(tn.str().c_str()));
}
ggml_tensor t_meta = *cur;
@@ -1363,7 +1401,8 @@ void llama_model_loader::done_getting_tensors(bool partial) const {
}
void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps) {
if (use_mmap) {
// note: read_lazy also requires mmap; this condition make sure it's usable even when --load-mode is not set to mmap
if (use_mmap || lazy.any()) {
mappings.reserve(files.size());
mmaps_used.reserve(files.size());
for (uint32_t idx = 0; idx < files.size(); idx++) {
@@ -1380,11 +1419,10 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps
}
}
const auto it_lazy = lazy_tensor_ranges.find(idx);
static const llama_mmap::ranges no_lazy_ranges;
const size_t prefetch_size = prefetch && use_mmap ? -1 : 0;
std::unique_ptr<llama_mmap> mapping = std::make_unique<llama_mmap>(file.get(), prefetch ? -1 : 0, is_numa,
it_lazy != lazy_tensor_ranges.end() ? it_lazy->second : no_lazy_ranges);
std::unique_ptr<llama_mmap> mapping = std::make_unique<llama_mmap>(file.get(), prefetch_size, is_numa,
lazy.for_file(idx));
mmaps_used.emplace_back(mapping->size(), 0);
if (mlock_mmaps) {
std::unique_ptr<llama_mlock> mlock_mmap(new llama_mlock());
@@ -1575,7 +1613,9 @@ bool llama_model_loader::load_all_data(
size_t n_size = ggml_nbytes(cur);
if (use_mmap) {
const bool from_mapping = use_mmap || lazy.has(cur);
if (from_mapping) {
const auto & mapping = mappings.at(weight->idx);
ggml_backend_buffer_t buf_mmap = nullptr;
if (bufs.count(weight->idx)) {
@@ -1592,7 +1632,9 @@ bool llama_model_loader::load_all_data(
GGML_ASSERT(buf_mmap || cur->data); // either we have a buffer to allocate the tensor in, or it is already allocated
if (buf_mmap && cur->data == nullptr) {
ggml_backend_tensor_alloc(buf_mmap, cur, data);
if (lmlocks) {
// locking a lazy tensor would fault all of it in, which is what lazy avoids
if (lmlocks && !lazy.has(cur)) {
const auto & lmlock = lmlocks->at(weight->idx);
lmlock->grow_to(weight->offs + n_size);
}