diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index e84e86297..0b1f0d351 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -117,6 +117,7 @@ struct Conditioner { virtual SDCondition get_learned_condition(int n_threads, const ConditionerParams& conditioner_params) = 0; virtual void get_param_tensors(std::map& tensors) = 0; + virtual void get_param_tensor_ops(std::map& tensor_ops) {} virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {} virtual void set_stream_layers_enabled(bool enabled) {} virtual void set_runtime_backends(const std::vector& backends) {} @@ -1664,6 +1665,10 @@ struct AnimaConditioner : public Conditioner { llm->get_param_tensors(tensors, "text_encoders.llm"); } + void get_param_tensor_ops(std::map& tensor_ops) override { + llm->get_param_tensor_ops(tensor_ops); + } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { llm->set_max_graph_vram_bytes(max_vram_bytes); } @@ -1847,6 +1852,10 @@ struct LLMEmbedder : public Conditioner { } } + void get_param_tensor_ops(std::map& tensor_ops) override { + llm->get_param_tensor_ops(tensor_ops); + } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { llm->set_max_graph_vram_bytes(max_vram_bytes); if (byt5) { @@ -2828,6 +2837,10 @@ struct LTXAVEmbedder : public Conditioner { projector->get_param_tensors(tensors, "text_embedding_projection"); } + void get_param_tensor_ops(std::map& tensor_ops) override { + llm->get_param_tensor_ops(tensor_ops); + } + void set_flash_attention_enabled(bool enabled) override { llm->set_flash_attention_enabled(enabled); projector->set_flash_attention_enabled(enabled); diff --git a/src/core/ggml_extend.hpp b/src/core/ggml_extend.hpp index d8e017795..d14d2875b 100644 --- a/src/core/ggml_extend.hpp +++ b/src/core/ggml_extend.hpp @@ -1753,7 +1753,7 @@ struct GGMLRunner { std::vector graph_cut_layer_split_backend_vram_limits_; std::vector extra_runtime_backends; // borrowed (SDBackendManager-owned) - ggml_backend_sched_t sched = nullptr; // owned, multi-device only + ggml_backend_sched_t sched = nullptr; // owned ggml_backend_t cpu_fallback_backend = nullptr; // owned, sched requires a trailing CPU backend bool multi_device_eval_callback_warned = false; @@ -2147,8 +2147,22 @@ struct GGMLRunner { return !extra_runtime_backends.empty(); } + bool graph_requires_backend_fallback(ggml_cgraph* gf) const { + if (gf == nullptr || sd_backend_is_cpu(runtime_backend)) { + return false; + } + const int n_nodes = ggml_graph_n_nodes(gf); + for (int i = 0; i < n_nodes; ++i) { + ggml_tensor* node = ggml_graph_node(gf, i); + if (node != nullptr && !ggml_backend_supports_op(runtime_backend, node)) { + return true; + } + } + return false; + } + bool alloc_compute_buffer(ggml_cgraph* gf) { - if (is_multi_device()) { + if (sched != nullptr || is_multi_device() || graph_requires_backend_fallback(gf)) { // The sched replaces the gallocr. Do NOT ggml_backend_sched_reserve // the graph here: reserve runs split_graph, which rewires the // graph's src pointers to sched-internal copy tensors, and the @@ -2156,6 +2170,10 @@ struct GGMLRunner { // rewired graph, silently corrupting every cross-backend input. A // graph must be split at most once; the alloc in execute_graph // performs the real allocation. + if (compute_allocr != nullptr) { + ggml_gallocr_free(compute_allocr); + compute_allocr = nullptr; + } return ensure_sched(gf); } if (compute_allocr != nullptr) { @@ -2753,7 +2771,7 @@ struct GGMLRunner { }; ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer); - if (is_multi_device()) { + if (sched != nullptr) { ggml_backend_sched_reset(sched); pin_multi_device_nodes(gf); // reset clears the pins; re-apply before alloc if (!ggml_backend_sched_alloc_graph(sched, gf)) { @@ -2774,9 +2792,9 @@ struct GGMLRunner { } ggml_status status; - if (is_multi_device()) { + if (sched != nullptr) { if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) { - LOG_WARN("%s: eval callback is not supported with multiple runtime backends; ignoring", + LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring", get_desc().c_str()); multi_device_eval_callback_warned = true; } @@ -3018,12 +3036,9 @@ struct GGMLRunner { // do copy after alloc graph void set_backend_tensor_data(ggml_tensor* tensor, const void* data) { - if (is_multi_device()) { - // The sched only assigns a backend (and thus a buffer) to tensors - // that participate in the graph; flag standalone data tensors as - // inputs so they get one. - ggml_set_input(tensor); - } + // The scheduler only allocates standalone data tensors when they are + // marked as graph inputs. The flag is harmless for single-backend graphs. + ggml_set_input(tensor); backend_tensor_data_map[tensor] = data; } @@ -3240,6 +3255,11 @@ class GGMLBlock { virtual void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {} + virtual enum ggml_op param_usage_op(const std::string& name) const { + (void)name; + return GGML_OP_NONE; + } + public: void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") { if (prefix.size() > 0) { @@ -3290,6 +3310,18 @@ class GGMLBlock { } } + void get_param_tensor_ops(std::map& tensor_ops) { + for (auto& pair : blocks) { + pair.second->get_param_tensor_ops(tensor_ops); + } + for (auto& pair : params) { + enum ggml_op op = param_usage_op(pair.first); + if (op != GGML_OP_NONE) { + tensor_ops[pair.second] = op; + } + } + } + virtual std::string get_desc() { return "GGMLBlock"; } @@ -3417,6 +3449,10 @@ class Embedding : public UnaryBlock { params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings); } + enum ggml_op param_usage_op(const std::string& name) const override { + return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE; + } + public: Embedding(int64_t num_embeddings, int64_t embedding_dim) : embedding_dim(embedding_dim), diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index 60e2d9c15..c0651a4a8 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -1657,6 +1657,10 @@ namespace LLM { model.get_param_tensors(tensors, prefix); } + void get_param_tensor_ops(std::map& tensor_ops) { + model.get_param_tensor_ops(tensor_ops); + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input_ids, ggml_tensor* input_pos, diff --git a/src/model_loader.cpp b/src/model_loader.cpp index ccd1e7d96..f1436b8a7 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -1053,7 +1053,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, std::atomic tensor_idx(0); std::atomic failed(false); std::vector workers; - std::mutex rpc_backend_mutex; + std::mutex backend_tensor_set_mutex; for (int i = 0; i < n_threads; ++i) { workers.emplace_back([&, file_path, is_zip]() { @@ -1214,17 +1214,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, if (dst_tensor->buffer != nullptr && !ggml_backend_buffer_is_host(dst_tensor->buffer)) { t0 = ggml_time_ms(); - // RPC backends require serialized access to prevent concurrency issues - const char* buffer_type_name = ggml_backend_buft_name(ggml_backend_buffer_get_type(dst_tensor->buffer)); - bool is_rpc_buffer = buffer_type_name != nullptr && - std::string(buffer_type_name).find("RPC") != std::string::npos; - - if (is_rpc_buffer) { - std::lock_guard lock(rpc_backend_mutex); - ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor)); - } else { - ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor)); - } + std::lock_guard lock(backend_tensor_set_mutex); + ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor)); t1 = ggml_time_ms(); copy_to_backend_time_ms.fetch_add(t1 - t0); diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 33bc18373..2f44b0003 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -53,6 +53,48 @@ static bool backend_supports_host_buffer(ggml_backend_t backend) { return props.caps.buffer_from_host_ptr; } +static bool device_supports_param_op(ggml_backend_dev_t device, + ggml_tensor* weight, + enum ggml_op op, + ggml_backend_buffer_type_t buft) { + if (op == GGML_OP_NONE) { + return true; + } + if (device == nullptr || weight == nullptr || buft == nullptr || weight->buffer != nullptr) { + return false; + } + + ggml_init_params params; + params.mem_size = ggml_tensor_overhead() * 2; + params.mem_buffer = nullptr; + params.no_alloc = true; + ggml_context* ctx = ggml_init(params); + if (ctx == nullptr) { + return false; + } + + ggml_tensor* op_tensor = nullptr; + if (op == GGML_OP_GET_ROWS) { + ggml_tensor* indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + op_tensor = ggml_get_rows(ctx, weight, indices); + } + if (op_tensor == nullptr) { + ggml_free(ctx); + return false; + } + + weight->buffer = ggml_backend_buft_alloc_buffer(buft, 0); + if (weight->buffer == nullptr) { + ggml_free(ctx); + return false; + } + bool supported = ggml_backend_dev_supports_op(device, op_tensor); + ggml_backend_buffer_free(weight->buffer); + weight->buffer = nullptr; + ggml_free(ctx); + return supported; +} + ModelManager::~ModelManager() { release_all(); } @@ -135,7 +177,8 @@ bool ModelManager::register_param_tensors(const std::string& desc, ggml_backend_t params_backend, size_t* registered_tensor_size, bool allow_split_buffer, - bool params_follow_compute_backend) { + bool params_follow_compute_backend, + const std::map* tensor_ops) { if (desc.empty()) { LOG_ERROR("model manager tensor desc is empty"); return false; @@ -168,6 +211,12 @@ bool ModelManager::register_param_tensors(const std::string& desc, state->params_backend = params_backend; state->allow_split_buffer = allow_split_buffer; state->params_follow_compute_backend = params_follow_compute_backend; + if (tensor_ops != nullptr) { + auto op_it = tensor_ops->find(tensor); + if (op_it != tensor_ops->end()) { + state->usage_op = op_it->second; + } + } new_states.push_back(std::move(state)); } @@ -844,6 +893,22 @@ ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorStat if (params_buft == nullptr) { params_buft = ggml_backend_get_default_buffer_type(state.params_backend); } + if (state.usage_op != GGML_OP_NONE && + state.compute_backend != nullptr) { + ggml_backend_dev_t compute_dev = ggml_backend_get_device(state.compute_backend); + if (device_supports_param_op(compute_dev, state.tensor, state.usage_op, params_buft)) { + return params_buft; + } + + ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); + params_buft = cpu_dev != nullptr ? ggml_backend_dev_buffer_type(cpu_dev) : nullptr; + if (!device_supports_param_op(cpu_dev, state.tensor, state.usage_op, params_buft)) { + LOG_ERROR("model manager has no compatible buffer for tensor '%s' used by %s", + state.name.c_str(), + ggml_op_name(state.usage_op)); + return nullptr; + } + } return params_buft; } diff --git a/src/model_manager.h b/src/model_manager.h index 85e982de2..881a17714 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -39,6 +39,7 @@ class ModelManager : public RunnerWeightManager { bool allow_split_buffer = false; bool params_follow_compute_backend = false; bool metadata_validated = false; + enum ggml_op usage_op = GGML_OP_NONE; int active_prepare_count = 0; @@ -132,7 +133,8 @@ class ModelManager : public RunnerWeightManager { ggml_backend_t params_backend, size_t* registered_tensor_size = nullptr, bool allow_split_buffer = false, - bool params_follow_compute_backend = false); + bool params_follow_compute_backend = false, + const std::map* tensor_ops = nullptr); bool unregister_param_tensors(const std::string& desc, size_t* registered_tensor_size = nullptr); diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 5f8dfa8d4..de3088178 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -320,7 +320,11 @@ class StableDiffusionGGML { return true; } std::map group_tensors; + std::map tensor_ops; model->get_param_tensors(group_tensors); + if constexpr (std::is_base_of_v) { + model->get_param_tensor_ops(tensor_ops); + } if (model_manager == nullptr) { return true; } @@ -337,6 +341,7 @@ class StableDiffusionGGML { module, module_backends, std::move(group_tensors), + tensor_ops, residency_mode, params_mem_size); } @@ -345,6 +350,7 @@ class StableDiffusionGGML { module, module_backends, std::move(group_tensors), + tensor_ops, residency_mode, params_mem_size); } @@ -358,7 +364,10 @@ class StableDiffusionGGML { residency_mode, backend_for(module), params_backend_for(module), - params_mem_size); + params_mem_size, + false, + false, + &tensor_ops); } template @@ -367,6 +376,7 @@ class StableDiffusionGGML { SDBackendModule module, const std::vector& module_backends, std::map group_tensors, + const std::map& tensor_ops, ModelManager::ResidencyMode residency_mode, size_t* params_mem_size) { ggml_backend_t main_backend = module_backends[0]; @@ -378,6 +388,7 @@ class StableDiffusionGGML { module, module_backends, std::move(group_tensors), + tensor_ops, residency_mode, params_mem_size); }; @@ -452,7 +463,9 @@ class StableDiffusionGGML { main_backend, params_backend_for(module), params_mem_size, - /*allow_split_buffer=*/true)) { + /*allow_split_buffer=*/true, + false, + &tensor_ops)) { return false; } return model_manager->register_param_tensors(desc, @@ -460,7 +473,10 @@ class StableDiffusionGGML { residency_mode, main_backend, params_backend_for(module), - params_mem_size); + params_mem_size, + false, + false, + &tensor_ops); } // Register graph-cut layer-split tensors on the primary backend first. @@ -472,6 +488,7 @@ class StableDiffusionGGML { SDBackendModule module, const std::vector& module_backends, std::map group_tensors, + const std::map& tensor_ops, ModelManager::ResidencyMode residency_mode, size_t* params_mem_size) { bool has_cpu_device = false; @@ -493,7 +510,10 @@ class StableDiffusionGGML { residency_mode, module_backends[0], params_backend_for(module), - params_mem_size); + params_mem_size, + false, + false, + &tensor_ops); } model->set_runtime_backends(module_backends); @@ -518,7 +538,8 @@ class StableDiffusionGGML { initial_params_backend, params_mem_size, false, - params_follow_runtime); + params_follow_runtime, + &tensor_ops); } bool unload_control_net() {