diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 7663797ba00..b4d979c9311 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1456,6 +1456,48 @@ void llama_model_loader::get_mapping_range(size_t * first, size_t * last, void * } } +void llama_model_loader::get_mapping_ranges(std::vector> & ranges, void ** addr, int idx, ggml_context * ctx) const { + GGML_ASSERT(!mappings.empty()); + const auto & mapping = mappings.at(idx); + + ranges.clear(); + *addr = mapping->addr(); + + for (ggml_tensor * tensor = ggml_get_first_tensor(ctx); tensor; tensor = ggml_get_next_tensor(ctx, tensor)) { + const auto * weight = get_weight(ggml_get_name(tensor)); + if (!weight || weight->idx != idx) { + continue; + } + ranges.emplace_back(weight->offs, weight->offs + ggml_nbytes(tensor)); + } + + if (ranges.empty()) { + return; + } + + std::sort(ranges.begin(), ranges.end()); + + // below this a gap is padding, and splitting on it would cost a buffer for nothing + constexpr size_t min_gap = 32ull*1024*1024; + + // never trade one oversized buffer for very many small ones + constexpr size_t max_ranges = 64; + + size_t n = 0; + for (size_t i = 1; i < ranges.size(); ++i) { + if (ranges[i].first > ranges[n].second + min_gap) { + ranges[++n] = ranges[i]; + } else { + ranges[n].second = std::max(ranges[n].second, ranges[i].second); + } + } + ranges.resize(n + 1); + + if (ranges.size() > max_ranges) { + ranges = { { ranges.front().first, ranges.back().second } }; + } +} + void llama_model_loader::unmap_weight(const llama_tensor_weight & w) const { if (!use_mmap) { return; } mappings.at(w.idx)->unmap_fragment(w.offs, w.offs + ggml_nbytes(w.tensor)); @@ -1523,7 +1565,7 @@ bool llama_model_loader::load_all_data( } // When not using mmaped io use async uploads from pinned memory to GPU memory. // First determine if the backend supports the necessary features for async uploads. - auto * buf = bufs.count(0) ? bufs.at(0) : nullptr; + auto * buf = bufs.count(0) && !bufs.at(0).empty() ? bufs.at(0).front() : nullptr; if (!buf) { LLAMA_LOG_DEBUG("%s: no buffer found for async uploads\n", func); return nullptr; @@ -1594,7 +1636,7 @@ bool llama_model_loader::load_all_data( if (upload_backend) { LLAMA_LOG_DEBUG("%s: using async uploads for device %s, buffer type %s, backend %s\n", __func__, ggml_backend_dev_name(ggml_backend_get_device(upload_backend)), - ggml_backend_buft_name(ggml_backend_buffer_get_type(bufs.at(0))), + ggml_backend_buft_name(ggml_backend_buffer_get_type(bufs.at(0).front())), ggml_backend_name(upload_backend)); } @@ -1618,10 +1660,17 @@ bool llama_model_loader::load_all_data( if (from_mapping) { const auto & mapping = mappings.at(weight->idx); ggml_backend_buffer_t buf_mmap = nullptr; + uint8_t * data = (uint8_t *) mapping->addr() + weight->offs; + if (bufs.count(weight->idx)) { - buf_mmap = bufs.at(weight->idx); + for (ggml_backend_buffer_t b : bufs.at(weight->idx)) { + uint8_t * base = (uint8_t *) ggml_backend_buffer_get_base(b); + if (data >= base && data + n_size <= base + ggml_backend_buffer_get_size(b)) { + buf_mmap = b; + break; + } + } } - uint8_t * data = (uint8_t *) mapping->addr() + weight->offs; if (check_tensors) { validation_result.emplace_back(std::async(std::launch::async, [cur, data, n_size] { diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 9e51d0ce750..6204a227457 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -16,7 +16,7 @@ #include #include -using llama_buf_map = std::unordered_map; +using llama_buf_map = std::unordered_map>; // lists of buffer types used for each layer using buft_list_t = std::vector>; @@ -244,6 +244,10 @@ struct llama_model_loader { void get_mapping_range(size_t * first, size_t * last, void ** addr, int idx, ggml_context * ctx) const; + // near-adjacent runs are merged, so a normal layout yields one range and only a large + // foreign tensor splits it + void get_mapping_ranges(std::vector> & ranges, void ** addr, int idx, ggml_context * ctx) const; + // release a weight's mmap pages void unmap_weight(const llama_tensor_weight & w) const; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index e679b24e87f..6e9192512d7 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1727,19 +1727,26 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { // this is important for metal with apple silicon: if the entire model could be mapped to a metal buffer, // then we could just use metal for all layers // this allows using partial offloading when the model size exceeds the metal buffer size, but not the RAM size + // a tensor of another buffer type can sit between this context's, and one span + // over them would map it too, so map each contiguous run separately void * addr = nullptr; - size_t first, last; // NOLINT - ml.get_mapping_range(&first, &last, &addr, idx, ctx); - if (first >= last) { + std::vector> ranges; + ml.get_mapping_ranges(ranges, &addr, idx, ctx); + if (ranges.empty()) { continue; } const size_t max_size = ggml_get_max_tensor_size(ctx); - ggml_backend_buffer_t buf = ggml_backend_dev_buffer_from_host_ptr(dev, (char *) addr + first, last - first, max_size); - if (buf == nullptr) { - throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft))); + for (const auto & [first, last] : ranges) { + if (first >= last) { + continue; + } + ggml_backend_buffer_t buf = ggml_backend_dev_buffer_from_host_ptr(dev, (char *) addr + first, last - first, max_size); + if (buf == nullptr) { + throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft))); + } + bufs.emplace_back(buf); + buf_map[idx].push_back(buf); } - bufs.emplace_back(buf); - buf_map.emplace(idx, buf); } } else { ggml_backend_buffer_t buf; @@ -1762,7 +1769,7 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } bufs.emplace_back(buf); for (uint32_t idx = 0; idx < ml.files.size(); idx++) { - buf_map.emplace(idx, buf); + buf_map[idx].push_back(buf); } }