feat: initial support for memory-mapping model weights

Author: wbruna

src/model.cpp

@@ -2,6 +2,7 @@
 #include <atomic>
 #include <chrono>
 #include <cstdarg>
+#include <cstdint>
 #include <fstream>
 #include <functional>
 #include <mutex>
@@ -768,6 +769,99 @@ void ModelLoader::process_model_files(bool enable_mmap) {
     LOG_INFO("model files processing completed in %.2fs", process_time_ms / 1000.f);
 }
 
+std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
+                                                       std::set<std::string> ignore_tensors)
+{
+    process_model_files(true);
+
+    std::vector<MmapTensorStore> result;
+    uint64_t mapped_bytes = 0;
+    size_t mapped_tensors = 0;
+
+    LOG_DEBUG("memory-mapping tensors...");
+
+    int64_t t_start = ggml_time_ms();
+
+    for (const auto& fdata : file_data) {
+        if (!fdata.mmapped) continue;
+
+        const std::vector<TensorStorage>& file_tensors = fdata.tensors;
+        std::shared_ptr<MmapWrapper> mmapped = fdata.mmapped;
+
+        uint8_t * mmap_data = const_cast<uint8_t*>(mmapped->data());
+
+        ggml_backend_buffer_t buf_mmap = ggml_backend_cpu_buffer_from_ptr(mmap_data, mmapped->size());
+        if (!buf_mmap) {
+            LOG_WARN("mmap: failed to create backend buffer for file %s", fdata.path.c_str());
+            continue;
+        }
+        ggml_backend_buffer_set_usage(buf_mmap, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+
+        size_t file_mapped_bytes   = 0;
+        size_t file_mapped_tensors = 0;
+
+        for (const auto& tensor_storage : file_tensors) {
+            const std::string& name = tensor_storage.name;
+
+            bool is_ignored = false;
+            for (const auto& ignore_prefix : ignore_tensors) {
+                if (starts_with(name, ignore_prefix)) {
+                    is_ignored = true;
+                    break;
+                }
+            }
+            if (is_ignored)
+                continue;
+
+            auto it = tensors.find(name);
+            if (it == tensors.end())
+                continue;
+
+            ggml_tensor* dst_tensor = it->second;
+            if (dst_tensor == nullptr)
+                continue;
+
+            if (tensor_storage.type != dst_tensor->type)
+                continue;
+
+            size_t tensor_size   = tensor_storage.nbytes();
+            size_t tensor_offset = tensor_storage.offset;
+
+            if (tensor_storage.ne[0] != dst_tensor->ne[0] ||
+                tensor_storage.ne[1] != dst_tensor->ne[1] ||
+                tensor_storage.ne[2] != dst_tensor->ne[2] ||
+                tensor_storage.ne[3] != dst_tensor->ne[3] ||
+                tensor_size          != ggml_nbytes(dst_tensor)) {
+                // let load_tensors worry about this
+                continue;
+            }
+
+            dst_tensor->buffer = buf_mmap;
+            dst_tensor->data   = mmap_data + tensor_offset;
+
+            file_mapped_bytes += tensor_size;
+            file_mapped_tensors++;
+        }
+
+        if (file_mapped_bytes > 0) {
+            mapped_tensors += file_mapped_tensors;
+            mapped_bytes   += file_mapped_bytes;
+            result.push_back({mmapped, buf_mmap});
+        }
+    }
+
+    int64_t t_end = ggml_time_ms();
+    int64_t duration_ms = t_end - t_start;
+
+    LOG_INFO("memory-mapped %zu tensors in %zu files (%.2f MB), taking %.2fs",
+             mapped_tensors,
+             result.size(),
+             mapped_bytes / (1024.0 * 1024.0),
+             duration_ms / 1000.0);
+
+    return result;
+}
+
 bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads_p, bool enable_mmap) {
 
     process_model_files(enable_mmap);
@@ -860,6 +954,12 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
                         continue;
                     }
 
+                    // skip mmapped tensors
+                    if (dst_tensor->buffer != nullptr
+                        && ggml_backend_buffer_get_usage(dst_tensor->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
+                        continue;
+                    }
+
                     size_t nbytes_to_read = tensor_storage.nbytes_to_read();
 
                     auto read_data = [&](char* buf, size_t n) {

src/model.h

@@ -199,6 +199,11 @@ struct ModelFileData {
     bool is_zip;
 };
 
+struct MmapTensorStore {
+    std::shared_ptr<MmapWrapper> mmapped;
+    ggml_backend_buffer_t buffer;
+};
+
 class ModelLoader {
 protected:
     SDVersion version_ = VERSION_COUNT;
@@ -228,6 +233,8 @@ class ModelLoader {
     String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
     void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = "");
     void process_model_files(bool enable_mmap = false);
+    std::vector<MmapTensorStore> mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
+                                              std::set<std::string> ignore_tensors = {});
     bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads = 0, bool use_mmap = false);
     bool load_tensors(std::map<std::string, ggml_tensor*>& tensors,
                       std::set<std::string> ignore_tensors = {},

src/stable-diffusion.cpp

@@ -107,6 +107,7 @@ static float get_cache_reuse_threshold(const sd_cache_params_t& params) {
 
 class StableDiffusionGGML {
 public:
+    std::vector<MmapTensorStore> mmap_tensor_store;
     ggml_backend_t backend             = nullptr;  // general backend
     ggml_backend_t clip_backend        = nullptr;
     ggml_backend_t control_net_backend = nullptr;
@@ -825,6 +826,16 @@ class StableDiffusionGGML {
             ignore_tensors.insert("text_encoders.llm.multi_modal_projector.");
         }
 
+        if (sd_ctx_params->enable_mmap) {
+            if (!(offload_params_to_cpu || ggml_backend_is_cpu(backend))) {
+                LOG_DEBUG("cannot memory-map model weights: only supported with CPU or --offload-to-cpu");
+            } else if (apply_lora_immediately) {
+                LOG_DEBUG("cannot memory-map model weights: only supported with --lora-apply-mode at_runtime");
+            } else {
+                mmap_tensor_store = model_loader.mmap_tensors(tensors, ignore_tensors);
+            }
+        }
+
         if (clip_vision) {
             clip_vision->alloc_params_buffer();
         }