furyhawk/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-07-21 10:15:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Parallelize weight loading for weights targeted at anonymous host ram and also GPU (#2057)

Browse Source 
* Parallelize weight loading for weights targeted at anonymous host memory

* Parallelize Cuda weight loading

* Add parallel loading for sm graph
This commit is contained in:
Farmadupe
2026-07-07 09:52:19 +02:00
committed by GitHub
parent 05cba319e0
commit 93c91fad63
3 changed files with 186 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-mmap.cpp
+11 -2
View File
@@ -76,7 +76,7 @@ struct llama_file::impl {
return ret;
}
impl(const char * fname, const char * mode) : path(fname) {
impl(const char * fname, const char * mode) : path(fname), mode(mode) {
fp = ggml_fopen(fname, mode);
if (fp == NULL) {
throw std::runtime_error(format("failed to open %s: %s", fname, strerror(errno)));
@@ -163,7 +163,7 @@ struct llama_file::impl {
}
}
#else
impl(const char * fname, const char * mode) : path(fname) {
impl(const char * fname, const char * mode) : path(fname), mode(mode) {
fp = ggml_fopen(fname, mode);
if (fp == NULL) {
throw std::runtime_error(format("failed to open %s: %s", fname, strerror(errno)));
@@ -244,11 +244,20 @@ struct llama_file::impl {
FILE * fp;
size_t size;
std::string mode;
};
llama_file::llama_file(const char * fname, const char * mode) : pimpl(std::make_unique<impl>(fname, mode)) {}
llama_file::~llama_file() = default;
std::unique_ptr<llama_file> llama_file::clone() const {
//can only clone readable file pointers without truncating.
GGML_ASSERT(!pimpl->mode.empty() && pimpl->mode[0] == 'r'
&& pimpl->mode.find('+') == std::string::npos);
return std::make_unique<llama_file>(pimpl->path.c_str(), pimpl->mode.c_str());
}
size_t llama_file::tell() const { return pimpl->tell(); }
size_t llama_file::size() const { return pimpl->size; }
src/llama-mmap.h
+2
View File
@@ -31,6 +31,8 @@ struct llama_file {
void write_u32(uint32_t val) const;
const std::string & get_path() const;
std::unique_ptr<llama_file> clone() const;
private:
struct impl;
std::unique_ptr<impl> pimpl;
src/llama-model-loader.cpp
+173 -66
View File
@@ -3,6 +3,7 @@
#include "llama-mmap.h"
#include "llama-model.h"
#include "ggml.h"
#include <memory>
//#include "ggml-backend.h"
#ifdef GGML_USE_CUDA
@@ -24,6 +25,10 @@
#include <future>
#include <regex>
#include <algorithm>
#include <thread>
#include <atomic>
#include <mutex>
#include <cstring>
#if defined(_WIN32)
#define WIN32_LEAN_AND_MEAN
@@ -1060,19 +1065,20 @@ bool llama_model_loader::load_all_data(
void * progress_callback_user_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;
// Number of worker threads for cuda and host tensor loading.
const int n_workers = 8;
std::vector<std::vector<no_init<uint8_t>>> read_bufs(n_workers);
#if defined(GGML_USE_CUDA)
// 4 staging buffers for async uploads, each sized 1MB seems to be a good default for single NVMe drives.
// NVMe raid configurations might require more / larger buffers.
constexpr size_t n_buffers = 4;
constexpr size_t buffer_size = 1 * 1024 * 1024; // 1MB
// One pinned staging buffer per worker for async uploads
constexpr size_t buffer_size = 16 * 1024 * 1024; // 16MB
std::vector<ggml_backend_buffer_t> host_buffers;
std::vector<void*> host_ptrs;
std::vector<ggml_backend_event_t> events;
size_t buffer_idx = 0; // buffer to use for async loads
ggml_backend_t cuda_backend = nullptr;
if (!use_mmap && !check_tensors) {
@@ -1092,31 +1098,39 @@ bool llama_model_loader::load_all_data(
// If the cuda backend is active create pinned memory buffers and events for synchronisation.
if (cuda_backend) {
for (size_t idx = 0; idx < n_buffers; ++idx) {
for (size_t idx = 0; idx < (size_t)n_workers; ++idx) {
host_buffers.emplace_back(ggml_backend_buft_alloc_buffer(llama_default_buffer_type_cpu(true), buffer_size));
host_ptrs.emplace_back(ggml_backend_buffer_get_base(host_buffers[idx]));
events.emplace_back(ggml_backend_event_new(cuda_backend));
}
// Force creation of the upload stream now
ggml_backend_event_record(events[0]);
ggml_backend_event_synchronize(events[0]);
}
}
#endif
for (struct ggml_tensor * cur = ggml_get_first_tensor(ctx); cur != NULL; cur = ggml_get_next_tensor(ctx, cur)) {
// Tensors are loaded in a threadpool. A lot of the ops are serialized
// and this is the mutex we use.
std::mutex load_mutex;
// Load model weights into a backing buffer:
// * mmap: host page cache Serial
// * host: system ram Parallel
// * cuda: GPU Parallel
// * --split-mode graph: GPU Parallel
// * rest: All other backends, Serial
auto load_tensor = [&](ggml_tensor * cur, int thread_idx) -> size_t {
const auto * weight = get_weight(ggml_get_name(cur));
if (weight == nullptr) {
// this can happen with split experts models
continue;
}
if (progress_callback) {
if (!progress_callback((float) size_done / size_data, progress_callback_user_data)) {
return false;
}
}
size_t n_size = ggml_nbytes(cur);
GGML_ASSERT(weight != nullptr);
GGML_ASSERT(weight->idx < files.size());
const size_t n_size = ggml_nbytes(cur);
const auto file = files.at(weight->idx)->clone();
// mmap. Serialized.
if (use_mmap) {
std::lock_guard<std::mutex> lock(load_mutex);
const auto & mapping = mappings.at(weight->idx);
ggml_backend_buffer_t buf_mmap = nullptr;
if (bufs_mmap.count(weight->idx)) {
@@ -1144,59 +1158,152 @@ bool llama_model_loader::load_all_data(
} else {
ggml_backend_tensor_set(cur, data, 0, n_size);
}
} else {
GGML_ASSERT(weight->idx < files.size());
const auto & file = files.at(weight->idx);
if (ggml_backend_buffer_is_host(cur->buffer)) {
file->seek(weight->offs, SEEK_SET);
file->read_raw(cur->data, n_size);
if (check_tensors) {
validation_result.emplace_back(std::async(std::launch::async, [cur, n_size] {
return std::make_pair(cur, ggml_validate_row_data(cur->type, cur->data, n_size));
}));
}
} else {
#if defined(GGML_USE_CUDA)
// If cuda_backend is valid load the tensor in chunks to pinned memory and upload the buffers asynchronously to the GPU.
if (cuda_backend) {
file->seek(weight->offs, SEEK_SET);
size_t bytes_read = 0;
while (bytes_read < n_size) {
size_t read_iteration = std::min<size_t>(buffer_size, n_size - bytes_read);
ggml_backend_event_synchronize(events[buffer_idx]);
file->read_raw(host_ptrs[buffer_idx], read_iteration);
ggml_backend_tensor_set_async(cuda_backend, cur, host_ptrs[buffer_idx], bytes_read, read_iteration);
ggml_backend_event_record(events[buffer_idx]);
bytes_read += read_iteration;
++buffer_idx;
buffer_idx %= n_buffers;
}
}
else
#endif
{
read_buf.resize(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);
if (check_tensors && !ggml_validate_row_data(cur->type, read_buf.data(), n_size)) {
throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(cur)));
}
}
}
return n_size;
}
size_done += n_size;
// host. Parallel.
if (ggml_backend_buffer_is_host(cur->buffer)) {
file->seek(weight->offs, SEEK_SET);
file->read_raw(cur->data, n_size);
if (check_tensors) {
std::lock_guard<std::mutex> lock(load_mutex);
validation_result.emplace_back(std::async(std::launch::async, [cur, n_size] {
return std::make_pair(cur, ggml_validate_row_data(cur->type, cur->data, n_size));
}));
}
return n_size;
}
// cuda. Parallel
#if defined(GGML_USE_CUDA)
// If cuda_backend is valid load the tensor in chunks to pinned memory and upload the buffers asynchronously to the GPU.
if (cuda_backend) {
file->seek(weight->offs, SEEK_SET);
size_t bytes_read = 0;
while (bytes_read < n_size) {
size_t read_iteration = std::min<size_t>(buffer_size, n_size - bytes_read);
ggml_backend_event_synchronize(events[thread_idx]);
file->read_raw(host_ptrs[thread_idx], read_iteration);
ggml_backend_tensor_set_async(cuda_backend, cur, host_ptrs[thread_idx], bytes_read, read_iteration);
ggml_backend_event_record(events[thread_idx]);
bytes_read += read_iteration;
}
return n_size;
}
// --split-mode graph. Parallel
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;
if (is_probably_split_mode_graph) {
auto & read_buf = read_bufs[thread_idx];
if (read_buf.capacity() > n_size) {
read_buf = std::vector<no_init<uint8_t>>();
}
read_buf.resize(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);
if (check_tensors && !ggml_validate_row_data(cur->type, read_buf.data(), n_size)) {
throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(cur)));
}
return n_size;
}
#endif
// rest. Serialized.
{
std::lock_guard<std::mutex> lock(load_mutex);
auto & read_buf = read_bufs[thread_idx];
read_buf.resize(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);
if (check_tensors && !ggml_validate_row_data(cur->type, read_buf.data(), n_size)) {
throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(cur)));
}
return n_size;
}
};
// An iterator that the threadpool shares to get the next tensor for loading.
// Each thread loops over this function until there are no more tensors left to load.
auto make_next_tensor = [&]() {
auto cursor = std::make_shared<std::atomic<ggml_tensor *>>(ggml_get_first_tensor(ctx));
auto next_tensor = [this, ctx, cursor]() -> ggml_tensor * {
while (true) {
ggml_tensor * cur = cursor->load();
if (!cur) {
return nullptr;
}
if (!cursor->compare_exchange_weak(cur, ggml_get_next_tensor(ctx, cur))) {
continue;
}
// With split experts models get_weight can return nullptr.
if (get_weight(ggml_get_name(cur))) {
return cur;
}
}
};
return next_tensor;
};
auto next_tensor = make_next_tensor();
std::atomic<size_t> loaded{size_done}; // bytes loaded so far, for the progress bar
std::atomic<bool> cancelled{false};
std::atomic<bool> failed{false};
std::exception_ptr first_exception;
// threadpool worker.
auto worker = [&](int thread_idx) {
try {
while (!cancelled.load() && !failed.load()) {
if (progress_callback) {
const size_t done = loaded.load(std::memory_order_relaxed);
std::lock_guard<std::mutex> lock(load_mutex);
if (!progress_callback((float) done / size_data, progress_callback_user_data)) {
cancelled.store(true);
break;
}
}
ggml_tensor * cur = next_tensor();
if (!cur) {
break;
}
const size_t n_size = load_tensor(cur, thread_idx);
loaded.fetch_add(n_size, std::memory_order_relaxed);
}
} catch (...) {
std::lock_guard<std::mutex> lock(load_mutex);
if (!failed.exchange(true)) {
first_exception = std::current_exception();
}
}
};
std::vector<std::thread> pool;
for (int thread_idx = 0; thread_idx < n_workers; ++thread_idx) {
pool.emplace_back(worker, thread_idx);
}
for (auto & t : pool) {
t.join();
}
size_done = loaded.load();
if (first_exception) {
std::rethrow_exception(first_exception);
}
if (cancelled.load()) {
return false;
}
#if defined(GGML_USE_CUDA)
// free temporary resources used for async cuda uploads
if (cuda_backend) {
for (size_t idx = 0; idx < n_buffers;++idx) {
for (size_t idx = 0; idx < (size_t)n_workers;++idx) {
ggml_backend_event_synchronize(events[idx]);
ggml_backend_event_free(events[idx]);
ggml_backend_buffer_free(host_buffers[idx]);