nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

use mmap to reduce unbounded memory usage for split-tensor graph parallell loading (#2102)

This commit is contained in:
Farmadupe 2026-07-09 17:27:18 +01:00 committed by GitHub
parent 3bb0e9f09c
commit c32c3819f7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 14 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
src/llama-model-loader.cpp
View File

@ -1078,6 +1078,8 @@ bool llama_model_loader::load_all_data(
std::vector<void*> host_ptrs; std::vector<void*> host_ptrs;
std::vector<ggml_backend_event_t> events; std::vector<ggml_backend_event_t> events;
std::vector<std::unique_ptr<llama_mmap>> split_mappings(files.size());
ggml_backend_t cuda_backend = nullptr; ggml_backend_t cuda_backend = nullptr;
if (!use_mmap && !check_tensors) { if (!use_mmap && !check_tensors) {
// When not using mmaped io use async uploads from pinned memory to GPU memory. // When not using mmaped io use async uploads from pinned memory to GPU memory.
@ -1197,17 +1199,21 @@ bool llama_model_loader::load_all_data(
const char * buffer_name = ggml_backend_buffer_name(cur->buffer); const char * buffer_name = ggml_backend_buffer_name(cur->buffer);
const bool is_probably_split_mode_graph = std::strncmp(buffer_name, GGML_CUDA_NAME, strlen(GGML_CUDA_NAME)) == 0; const bool is_probably_split_mode_graph = std::strncmp(buffer_name, GGML_CUDA_NAME, strlen(GGML_CUDA_NAME)) == 0;
if (is_probably_split_mode_graph) { if (is_probably_split_mode_graph) {
auto & read_buf = read_bufs[thread_idx]; llama_mmap * mapping;
if (read_buf.capacity() > n_size) { {
read_buf = std::vector<no_init<uint8_t>>(); std::lock_guard<std::mutex> lock(load_mutex);
auto & m = split_mappings[weight->idx];
if (!m) {
m.reset(new llama_mmap(files.at(weight->idx).get(), 0, ggml_is_numa()));
} }
read_buf.resize(n_size); mapping = m.get();
file->seek(weight->offs, SEEK_SET); }
file->read_raw(read_buf.data(), n_size); uint8_t * data = (uint8_t *) mapping->addr() + weight->offs;
ggml_backend_tensor_set(cur, read_buf.data(), 0, n_size); ggml_backend_tensor_set(cur, data, 0, n_size);
if (check_tensors && !ggml_validate_row_data(cur->type, read_buf.data(), n_size)) { if (check_tensors && !ggml_validate_row_data(cur->type, data, n_size)) {
throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(cur))); throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(cur)));
} }
mapping->dontneed_fragment(weight->offs, weight->offs + n_size);
return n_size; return n_size;
} }
#endif #endif