misc : prevent RAM peaking at model loading stage (#27483)

This commit is contained in:
Tarek Dakhran
2026-09-03 09:32:24 +02:00
committed by GitHub
parent 4aa6ffba25
commit 5ec4eab69e
2 changed files with 28 additions and 2 deletions
+20 -2
View File
@@ -1497,7 +1497,6 @@ bool llama_model_loader::load_all_data(
}
GGML_ASSERT(size_data != 0 && "call init_mappings() first");
std::vector<no_init<uint8_t>> read_buf;
std::vector<std::future<std::pair<ggml_tensor *, bool>>> validation_result;
// 4 staging buffers for async uploads, each sized 1MB seems to be a good default for single NVMe drives.
@@ -1598,7 +1597,25 @@ bool llama_model_loader::load_all_data(
ggml_backend_name(upload_backend));
}
std::vector<ggml_tensor *> tensors;
for (struct ggml_tensor * cur = ggml_get_first_tensor(ctx); cur != NULL; cur = ggml_get_next_tensor(ctx, cur)) {
tensors.push_back(cur);
}
// without mmap, tensors in non-host buffers are staged through a temporary buffer sized like the tensor
// load them biggest-first so the largest staging buffer is allocated while the fewest weights are resident
if (!use_mmap) {
std::stable_sort(tensors.begin(), tensors.end(), [](const ggml_tensor * a, const ggml_tensor * b) {
const bool staged_a = a->buffer && !ggml_backend_buffer_is_host(a->buffer);
const bool staged_b = b->buffer && !ggml_backend_buffer_is_host(b->buffer);
if (staged_a != staged_b) {
return staged_a;
}
return staged_a && ggml_nbytes(a) > ggml_nbytes(b);
});
}
for (struct ggml_tensor * cur : tensors) {
const auto * weight = get_weight(ggml_get_name(cur));
if (weight == nullptr) {
// this can happen with split experts models
@@ -1711,7 +1728,8 @@ bool llama_model_loader::load_all_data(
buffer_idx %= n_buffers;
}
} else {
read_buf.resize(n_size);
// scoped to one tensor so only one staging buffer is alive at a time
std::vector<no_init<uint8_t>> read_buf(n_size);
file->seek(weight->offs, SEEK_SET);
file->read_raw(read_buf.data(), n_size);
ggml_backend_tensor_set(cur, read_buf.data(), 0, n_size);
+8
View File
@@ -1816,6 +1816,14 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
return true;
}
// without mmap, load non-host buffers first: their tensors go through a staging buffer, which is cheapest while the fewest weights are resident
if (!ml.use_mmap) {
std::stable_partition(ctx_buf_maps.begin(), ctx_buf_maps.end(), [](const auto & ctx_buf_map) {
const auto & buf_map = ctx_buf_map.second;
return !buf_map.empty() && !ggml_backend_buffer_is_host(buf_map.begin()->second);
});
}
// load tensor data
for (auto & [ctx, buf_map] : ctx_buf_maps) {
if (!ml.load_all_data(ctx, buf_map, use_mlock ? &pimpl->mlock_mmaps : NULL, params.progress_callback, params.progress_callback_user_data)) {