llama: map each contiguous run of a context's tensors, not one span over all of them

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<std::pair<size_t, size_t>> & 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));
@@ -1522,7 +1564,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;
@@ -1593,7 +1635,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));
     }
 
@@ -1635,10 +1677,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] {

src/llama-model-loader.h

@@ -16,7 +16,7 @@
 #include <stdexcept>
 #include <unordered_map>
 
-using llama_buf_map = std::unordered_map<uint32_t, ggml_backend_buffer_t>;
+using llama_buf_map = std::unordered_map<uint32_t, std::vector<ggml_backend_buffer_t>>;
 
 // lists of buffer types used for each layer
 using buft_list_t = std::vector<std::pair<ggml_backend_dev_t, ggml_backend_buffer_type_t>>;
@@ -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<std::pair<size_t, size_t>> & ranges, void ** addr, int idx, ggml_context * ctx) const;
+
     // release a weight's mmap pages
     void unmap_weight(const llama_tensor_weight & w) const;
 

src/llama-model.cpp

@@ -1738,19 +1738,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<std::pair<size_t, size_t>> 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;
@@ -1773,7 +1780,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);
             }
         }
 
@@ -1820,7 +1827,13 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
     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);
+            // one entry is a run of buffers now, not a single buffer, and every
+            // buffer in a context comes from the same buft, so the first one
+            // answers the question for all of them
+            if (buf_map.empty() || buf_map.begin()->second.empty()) {
+                return false;
+            }
+            return !ggml_backend_buffer_is_host(buf_map.begin()->second.front());
         });
     }