diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index cdf9bdf78..25dc8163a 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -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(fname, mode)) {} llama_file::~llama_file() = default; +std::unique_ptr 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(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; } diff --git a/src/llama-mmap.h b/src/llama-mmap.h index 0d0576567..b7ee01685 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -31,6 +31,8 @@ struct llama_file { void write_u32(uint32_t val) const; const std::string & get_path() const; + std::unique_ptr clone() const; + private: struct impl; std::unique_ptr pimpl; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 180669d6d..c65a40988 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -3,6 +3,7 @@ #include "llama-mmap.h" #include "llama-model.h" #include "ggml.h" +#include //#include "ggml-backend.h" #ifdef GGML_USE_CUDA @@ -24,6 +25,10 @@ #include #include #include +#include +#include +#include +#include #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> read_buf; std::vector>> validation_result; + // Number of worker threads for cuda and host tensor loading. + const int n_workers = 8; + + std::vector>> 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 host_buffers; std::vector host_ptrs; std::vector 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 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(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 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(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>(); + } + 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 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>(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 loaded{size_done}; // bytes loaded so far, for the progress bar + std::atomic cancelled{false}; + std::atomic 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 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 lock(load_mutex); + if (!failed.exchange(true)) { + first_exception = std::current_exception(); + } + } + }; + + std::vector 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]);