diff --git a/docs/backend.md b/docs/backend.md index 58c2ca8b..70eef34f 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -161,6 +161,9 @@ resident allocations. Vulkan reports exceeding total memory are rejected because its heap-budget subtraction can underflow. Other backends use the cap instead of treating such reports as zero free memory. Failed checks log the reported free and total memory alongside tracked weight and runtime allocations. +With `--mmap`, device-backed mappings count toward these budgets at their full +mapped-file size, once per device buffer even when multiple parameter blocks +share it. Mappings retained in the loader cache continue to count. Components are considered in `diffusion`, `te`, `vae` order so that repeatedly used diffusion weights have priority. Each component's weights use the first diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index f3049dc2..b1bce984 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -13,6 +13,7 @@ #endif #include "core/util.h" +#include "ggml-backend-impl.h" #include "ggml-impl.h" #include "stable-diffusion.h" @@ -433,6 +434,24 @@ bool sd_backend_is_cpu(ggml_backend_t backend) { return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU; } +ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, + void* ptr, + size_t size, + size_t max_tensor_size) { + ggml_backend_buffer_t buffer = ggml_backend_dev_buffer_from_host_ptr(device, ptr, size, max_tensor_size); + if (buffer != nullptr && buffer->context == nullptr) { + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(device); + if (reg != nullptr && std::strcmp(ggml_backend_reg_name(reg), "Metal") == 0) { + // Metal can wrap a failed mapping in a non-null buffer. Its free callback also + // dereferences the missing context, so only release the outer buffer. + buffer->iface.free_buffer = nullptr; + ggml_backend_buffer_free(buffer); + return nullptr; + } + } + return buffer; +} + bool sd_backend_supports_cuda_mma(ggml_backend_t backend) { #ifdef SD_USE_CUDA if (!sd_backend_is(backend, "CUDA")) { diff --git a/src/core/ggml_extend_backend.h b/src/core/ggml_extend_backend.h index b2ebd7eb..df1d6924 100644 --- a/src/core/ggml_extend_backend.h +++ b/src/core/ggml_extend_backend.h @@ -88,6 +88,10 @@ private: bool sd_backend_is(ggml_backend_t backend, const std::string& name); bool sd_backend_is_cpu(ggml_backend_t backend); bool sd_backend_supports_cuda_mma(ggml_backend_t backend); +ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, + void* ptr, + size_t size, + size_t max_tensor_size); ggml_backend_t sd_backend_cpu_init(); bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads); ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend, diff --git a/src/model_loader.cpp b/src/model_loader.cpp index ee5aa49e..2bd9261d 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -874,7 +874,8 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) { std::vector ModelLoader::mmap_tensors(std::map& tensors, std::set ignore_tensors, - bool writable_mmap) { + bool writable_mmap, + ggml_backend_dev_t device) { std::set names; for (const auto& entry : tensors) { names.insert(entry.first); @@ -896,6 +897,39 @@ std::vector ModelLoader::mmap_tensors(std::map file_buffer = device == nullptr ? fdata.mmbuffer : nullptr; + bool file_unmappable = false; + + auto buffer_for_file = [&]() -> ggml_backend_buffer_t { + if (file_buffer || file_unmappable) { + return file_buffer.get(); + } + auto cached = fdata.device_mmbuffers.find(device); + if (cached != fdata.device_mmbuffers.end()) { + file_buffer = cached->second; + return file_buffer.get(); + } + size_t max_tensor_size = 0; + for (const auto& ts : fdata.tensors) { + max_tensor_size = std::max(max_tensor_size, static_cast(ts.nbytes())); + } + ggml_backend_buffer_t buf = sd_backend_dev_buffer_from_host_ptr(device, + fdata.mmapped->writable_data(), + fdata.mmapped->size(), + max_tensor_size); + if (buf == nullptr) { + LOG_WARN("mmap: %s cannot map '%s', loading it instead", + ggml_backend_dev_name(device), fdata.path.c_str()); + file_unmappable = true; + return nullptr; + } + LOG_INFO("mmap: mapped '%s' for %s", fdata.path.c_str(), ggml_backend_dev_name(device)); + file_buffer = std::shared_ptr(buf, ggml_backend_buffer_free); + fdata.device_mmbuffers[device] = file_buffer; + return file_buffer.get(); + }; + const std::vector& file_tensors = fdata.tensors; size_t file_mapped_bytes = 0; @@ -944,10 +978,13 @@ std::vector ModelLoader::mmap_tensors(std::map(ggml_backend_buffer_get_base(buf_mmap)); - dst_tensor->buffer = buf_mmap; - dst_tensor->data = mmap_data + tensor_offset; + ggml_backend_buffer_t buf_mmap = buffer_for_file(); + if (buf_mmap == nullptr) { + break; + } + uint8_t* mmap_data = static_cast(ggml_backend_buffer_get_base(buf_mmap)); + dst_tensor->buffer = buf_mmap; + dst_tensor->data = mmap_data + tensor_offset; file_mapped_bytes += tensor_size; file_mapped_tensors++; @@ -956,7 +993,7 @@ std::vector ModelLoader::mmap_tensors(std::map 0) { mapped_tensors += file_mapped_tensors; mapped_bytes += file_mapped_bytes; - result.push_back({fdata.mmapped, fdata.mmbuffer}); + result.push_back({fdata.mmapped, file_buffer}); } } @@ -972,6 +1009,16 @@ std::vector ModelLoader::mmap_tensors(std::map ModelLoader::get_device_mmap_buffers() const { + std::vector buffers; + for (const auto& fdata : file_data) { + for (const auto& entry : fdata.device_mmbuffers) { + buffers.push_back(entry.second.get()); + } + } + return buffers; +} + bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, bool enable_mmap, const std::set* target_tensor_names, @@ -1115,6 +1162,11 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, if (dst_tensor->buffer != nullptr && dst_tensor->buffer == fdata.mmbuffer.get()) { continue; } + if (dst_tensor->buffer != nullptr && + std::any_of(fdata.device_mmbuffers.begin(), fdata.device_mmbuffers.end(), + [&](const auto& entry) { return entry.second.get() == dst_tensor->buffer; })) { + continue; + } size_t nbytes_to_read = tensor_storage.nbytes_to_read(); diff --git a/src/model_loader.h b/src/model_loader.h index c5335b53..e51b260b 100644 --- a/src/model_loader.h +++ b/src/model_loader.h @@ -20,6 +20,8 @@ struct ModelFileData { std::vector tensors; std::shared_ptr mmapped; std::shared_ptr mmbuffer; + // mmapped wrapped by devices that can use host memory in place (buffer_from_host_ptr) + std::map> device_mmbuffers; bool is_zip; }; @@ -120,7 +122,9 @@ public: void process_model_files(bool enable_mmap = false, bool writable_mmap = true); std::vector mmap_tensors(std::map& tensors, std::set ignore_tensors = {}, - bool writable = true); + bool writable = true, + ggml_backend_dev_t device = nullptr); + std::vector get_device_mmap_buffers() const; bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, bool use_mmap = false, const std::set* target_tensor_names = nullptr, diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 930c6f49..5bdf59ca 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -780,38 +780,52 @@ bool ModelManager::validate_tensor(const TensorState& state) const { bool ModelManager::mmap_params(const std::vector& states, std::vector& created_storage_blocks) { - std::map mmap_candidates; - std::map mmap_states; + // A GPU that computes on mmapped params in place cannot address a CPU buffer, and nothing + // stages them for it, so they are mapped through a buffer of that GPU's device. + struct MmapGroup { + std::map candidates; + std::map states; + }; + std::map groups; for (TensorState* state : states) { if (state == nullptr || !can_mmap_storage(*state) || state->tensor == nullptr || state->tensor->data != nullptr || state->tensor->view_src != nullptr) { continue; } - mmap_candidates[state->name] = state->tensor; - mmap_states[state->name] = state; - } - if (mmap_candidates.empty()) { - return true; - } - - auto mmap_store = model_loader_.mmap_tensors(mmap_candidates, {}, writable_mmap_); - if (mmap_store.empty()) { - return true; - } - - auto block = std::make_unique(); - block->mmap_tensor_stores = std::move(mmap_store); - ParamsStorageBlock* raw = block.get(); - for (const auto& pair : mmap_states) { - TensorState* state = pair.second; - if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) { - block->states.push_back(state); + ggml_backend_dev_t device = nullptr; + if (!sd_backend_is_cpu(state->compute_backend) && !sd_backend_is_cpu(state->params_backend)) { + device = ggml_backend_get_device(state->compute_backend); } + MmapGroup& group = groups[device]; + group.candidates[state->name] = state->tensor; + group.states[state->name] = state; } - if (!block->states.empty()) { - params_storage_blocks_.push_back(std::move(block)); - created_storage_blocks.push_back(raw); + for (auto& [device, group] : groups) { + // Device buffers wrap read-only mappings only; params that LoRAs are merged into in place + // are loaded instead. + if (device != nullptr && writable_mmap_) { + continue; + } + auto mmap_store = model_loader_.mmap_tensors(group.candidates, {}, writable_mmap_, device); + if (mmap_store.empty()) { + continue; + } + + auto block = std::make_unique(); + block->mmap_tensor_stores = std::move(mmap_store); + ParamsStorageBlock* raw = block.get(); + for (const auto& pair : group.states) { + TensorState* state = pair.second; + if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) { + block->states.push_back(state); + } + } + + if (!block->states.empty()) { + params_storage_blocks_.push_back(std::move(block)); + created_storage_blocks.push_back(raw); + } } return true; } @@ -1353,15 +1367,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe } size_t total_size = 0; - auto add_buffer = [&](ggml_backend_buffer_t buffer) { - if (buffer == nullptr || ggml_backend_buffer_is_host(buffer)) { + std::unordered_set seen; + auto add_buffer = [&](ggml_backend_buffer_t buffer) { + if (buffer == nullptr || ggml_backend_buffer_is_host(buffer) || !seen.insert(buffer).second) { return; } ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(buffer); auto split_devices = split_buffer_devices_.find(buffer_type); const bool on_device = split_devices == split_buffer_devices_.end() - ? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device - : std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) { + ? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device + : std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) { return ggml_backend_get_device(entry.first) == compute_device; }); if (!on_device) { @@ -1371,9 +1386,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe total_size = buffer_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + buffer_size; }; + // The loader may retain device mappings after their parameter blocks are released. + for (ggml_backend_buffer_t buffer : model_loader_.get_device_mmap_buffers()) { + add_buffer(buffer); + } for (const auto& block : params_storage_blocks_) { if (block != nullptr) { add_buffer(block->buffer); + for (const auto& store : block->mmap_tensor_stores) { + add_buffer(store.mmbuffer.get()); + } } } for (const auto& block : compute_staging_blocks_) {