From a9acb28674d8960b0b8868c4402e0dee5ea7a5ab Mon Sep 17 00:00:00 2001 From: leejet Date: Tue, 7 Jul 2026 23:02:57 +0800 Subject: [PATCH] feat: drive layer split from graph-cut segments --- src/conditioning/conditioner.hpp | 126 ++++++++++++ src/core/ggml_extend.hpp | 178 ++++++++++++++++- src/core/layer_split_partition.cpp | 298 ++++++++++++++++------------- src/core/layer_split_partition.h | 34 +++- src/model_manager.cpp | 68 ++++++- src/model_manager.h | 18 +- src/stable-diffusion.cpp | 89 ++++----- src/weight_manager.h | 4 + 8 files changed, 610 insertions(+), 205 deletions(-) diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index 3d315985b..5d4ad7fbb 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -118,6 +118,8 @@ struct Conditioner { 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) {} + virtual void set_graph_cut_layer_split_enabled(bool enabled) {} + virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) {} virtual void get_layer_split_param_tensors(std::map& tensors) {} virtual void set_flash_attention_enabled(bool enabled) = 0; virtual void set_weight_adapter(const std::shared_ptr& adapter) {} @@ -181,6 +183,27 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { } } + void set_runtime_backends(const std::vector& backends) override { + text_model->set_runtime_backends(backends); + if (sd_version_is_sdxl(version)) { + text_model2->set_runtime_backends(backends); + } + } + + void set_graph_cut_layer_split_enabled(bool enabled) override { + text_model->set_graph_cut_layer_split_enabled(enabled); + if (sd_version_is_sdxl(version)) { + text_model2->set_graph_cut_layer_split_enabled(enabled); + } + } + + void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) override { + text_model->set_graph_cut_layer_split_backend_vram_limits(limits); + if (sd_version_is_sdxl(version)) { + text_model2->set_graph_cut_layer_split_backend_vram_limits(limits); + } + } + void set_flash_attention_enabled(bool enabled) override { text_model->set_flash_attention_enabled(enabled); if (sd_version_is_sdxl(version)) { @@ -639,11 +662,41 @@ struct SD3CLIPEmbedder : public Conditioner { } void set_runtime_backends(const std::vector& backends) override { + if (clip_l) { + clip_l->set_runtime_backends(backends); + } + if (clip_g) { + clip_g->set_runtime_backends(backends); + } if (t5) { t5->set_runtime_backends(backends); } } + void set_graph_cut_layer_split_enabled(bool enabled) override { + if (clip_l) { + clip_l->set_graph_cut_layer_split_enabled(enabled); + } + if (clip_g) { + clip_g->set_graph_cut_layer_split_enabled(enabled); + } + if (t5) { + t5->set_graph_cut_layer_split_enabled(enabled); + } + } + + void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) override { + if (clip_l) { + clip_l->set_graph_cut_layer_split_backend_vram_limits(limits); + } + if (clip_g) { + clip_g->set_graph_cut_layer_split_backend_vram_limits(limits); + } + if (t5) { + t5->set_graph_cut_layer_split_backend_vram_limits(limits); + } + } + void get_layer_split_param_tensors(std::map& tensors) override { if (t5) { t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); @@ -1010,11 +1063,32 @@ struct FluxCLIPEmbedder : public Conditioner { } void set_runtime_backends(const std::vector& backends) override { + if (clip_l) { + clip_l->set_runtime_backends(backends); + } if (t5) { t5->set_runtime_backends(backends); } } + void set_graph_cut_layer_split_enabled(bool enabled) override { + if (clip_l) { + clip_l->set_graph_cut_layer_split_enabled(enabled); + } + if (t5) { + t5->set_graph_cut_layer_split_enabled(enabled); + } + } + + void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) override { + if (clip_l) { + clip_l->set_graph_cut_layer_split_backend_vram_limits(limits); + } + if (t5) { + t5->set_graph_cut_layer_split_backend_vram_limits(limits); + } + } + void get_layer_split_param_tensors(std::map& tensors) override { if (t5) { t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); @@ -1278,6 +1352,18 @@ struct T5CLIPEmbedder : public Conditioner { } } + void set_graph_cut_layer_split_enabled(bool enabled) override { + if (t5) { + t5->set_graph_cut_layer_split_enabled(enabled); + } + } + + void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) override { + if (t5) { + t5->set_graph_cut_layer_split_backend_vram_limits(limits); + } + } + void get_layer_split_param_tensors(std::map& tensors) override { if (t5) { t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); @@ -1482,6 +1568,18 @@ struct MiniT2IConditioner : public Conditioner { } } + void set_graph_cut_layer_split_enabled(bool enabled) override { + if (t5) { + t5->set_graph_cut_layer_split_enabled(enabled); + } + } + + void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) override { + if (t5) { + t5->set_graph_cut_layer_split_backend_vram_limits(limits); + } + } + void get_layer_split_param_tensors(std::map& tensors) override { if (t5) { t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); @@ -1576,6 +1674,14 @@ struct AnimaConditioner : public Conditioner { llm->set_runtime_backends(backends); } + void set_graph_cut_layer_split_enabled(bool enabled) override { + llm->set_graph_cut_layer_split_enabled(enabled); + } + + void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) override { + llm->set_graph_cut_layer_split_backend_vram_limits(limits); + } + void get_layer_split_param_tensors(std::map& tensors) override { llm->get_param_tensors(tensors, "text_encoders.llm"); } @@ -1729,6 +1835,18 @@ struct LLMEmbedder : public Conditioner { llm->set_runtime_backends(backends); } + void set_graph_cut_layer_split_enabled(bool enabled) override { + if (llm) { + llm->set_graph_cut_layer_split_enabled(enabled); + } + } + + void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) override { + if (llm) { + llm->set_graph_cut_layer_split_backend_vram_limits(limits); + } + } + void get_layer_split_param_tensors(std::map& tensors) override { llm->get_param_tensors(tensors, "text_encoders.llm"); } @@ -2406,6 +2524,14 @@ struct LTXAVEmbedder : public Conditioner { llm->set_runtime_backends(backends); } + void set_graph_cut_layer_split_enabled(bool enabled) override { + llm->set_graph_cut_layer_split_enabled(enabled); + } + + void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) override { + llm->set_graph_cut_layer_split_backend_vram_limits(limits); + } + void get_layer_split_param_tensors(std::map& tensors) override { llm->get_param_tensors(tensors, "text_encoders.llm"); } diff --git a/src/core/ggml_extend.hpp b/src/core/ggml_extend.hpp index 1dcf423f1..36067ce3c 100644 --- a/src/core/ggml_extend.hpp +++ b/src/core/ggml_extend.hpp @@ -21,10 +21,12 @@ #include #include #include +#include #include #include "core/ggml_extend_backend.h" #include "core/ggml_graph_cut.h" +#include "core/layer_split_partition.h" #include "ggml-alloc.h" #include "ggml-backend.h" #include "ggml.h" @@ -1745,6 +1747,8 @@ struct GGMLRunner { size_t max_graph_vram_bytes = 0; bool stream_layers_enabled = false; size_t observed_max_effective_budget_ = 0; + bool graph_cut_layer_split_enabled = false; + 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 @@ -1776,6 +1780,9 @@ struct GGMLRunner { sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_; std::unordered_set params_tensor_set_; + std::unordered_map graph_cut_layer_split_assignments_; + std::unordered_map graph_cut_layer_split_node_assignments_; + bool graph_cut_layer_split_primary_notice_logged_ = false; template static sd::Tensor take_or_empty(std::optional> tensor) { @@ -1874,6 +1881,20 @@ struct GGMLRunner { params_tensor_set_dirty_ = false; } + ggml_tensor* canonical_param_tensor(ggml_tensor* tensor) { + if (tensor == nullptr) { + return nullptr; + } + if (params_tensor_set_.find(tensor) != params_tensor_set_.end()) { + return tensor; + } + if (tensor->view_src != nullptr && + params_tensor_set_.find(tensor->view_src) != params_tensor_set_.end()) { + return tensor->view_src; + } + return nullptr; + } + std::vector collect_used_param_tensors(ggml_cgraph* gf) { std::vector used_params; rebuild_params_tensor_set(); @@ -1886,12 +1907,8 @@ struct GGMLRunner { seen_params.reserve(static_cast(n_leafs)); for (int i = 0; i < n_leafs; ++i) { ggml_tensor* leaf = sd::ggml_graph_cut::leaf_tensor(gf, i); - ggml_tensor* param_leaf = leaf; - if (param_leaf != nullptr && params_tensor_set_.find(param_leaf) == params_tensor_set_.end()) { - param_leaf = param_leaf->view_src; - } + ggml_tensor* param_leaf = canonical_param_tensor(leaf); if (param_leaf != nullptr && - params_tensor_set_.find(param_leaf) != params_tensor_set_.end() && seen_params.insert(param_leaf).second) { used_params.push_back(param_leaf); } @@ -2101,11 +2118,17 @@ struct GGMLRunner { ggml_backend_t current = runtime_backend; 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); + ggml_tensor* node = ggml_graph_node(gf, i); + auto node_assignment = graph_cut_layer_split_node_assignments_.find(node); + if (node_assignment != graph_cut_layer_split_node_assignments_.end()) { + current = node_assignment->second; + } for (int s = 0; s < GGML_MAX_SRC; s++) { ggml_backend_t weight_backend = backend_for_weight(node->src[s]); if (weight_backend != nullptr) { - current = weight_backend; + if (node_assignment == graph_cut_layer_split_node_assignments_.end()) { + current = weight_backend; + } } } if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || @@ -2435,6 +2458,123 @@ struct GGMLRunner { return true; } + bool resolve_graph_cut_layer_split_plan(ggml_cgraph* gf, + GraphCutPlan* plan_out) { + GGML_ASSERT(plan_out != nullptr); + GGML_ASSERT(gf != nullptr); + *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend, + gf, + &graph_cut_plan_cache_, + 0, + params_tensor_set_, + get_desc().c_str()); + return true; + } + + bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf) { + graph_cut_layer_split_node_assignments_.clear(); + if (!graph_cut_layer_split_enabled) { + return true; + } + if (!is_multi_device()) { + LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str()); + return false; + } + + GraphCutPlan plan; + if (!resolve_graph_cut_layer_split_plan(gf, &plan)) { + return false; + } + if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) { + auto manager = weight_manager.lock(); + if (manager == nullptr) { + LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str()); + return false; + } + std::vector graph_params = collect_used_param_tensors(gf); + if (!graph_params.empty() && + !manager->assign_compute_backend(graph_params, runtime_backend)) { + LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s", + get_desc().c_str(), + sd::layer_split_backend_device_display_name(runtime_backend).c_str()); + return false; + } + for (ggml_tensor* param : graph_params) { + if (param != nullptr) { + graph_cut_layer_split_assignments_[param] = runtime_backend; + } + } + 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) { + graph_cut_layer_split_node_assignments_[node] = runtime_backend; + } + } + if (!graph_cut_layer_split_primary_notice_logged_) { + LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params", + get_desc().c_str(), + sd::layer_split_backend_device_display_name(runtime_backend).c_str(), + graph_params.size()); + graph_cut_layer_split_primary_notice_logged_ = true; + } else { + LOG_DEBUG("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params", + get_desc().c_str(), + sd::layer_split_backend_device_display_name(runtime_backend).c_str(), + graph_params.size()); + } + return true; + } + + std::vector split_backends; + split_backends.reserve(extra_runtime_backends.size() + 1); + split_backends.push_back(runtime_backend); + for (ggml_backend_t backend : extra_runtime_backends) { + if (backend != nullptr) { + split_backends.push_back(backend); + } + } + + auto manager = weight_manager.lock(); + if (manager == nullptr) { + LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str()); + return false; + } + + sd::GraphCutLayerSplitAssignment assignment; + auto canonicalize_param = [this](ggml_tensor* tensor) { + return canonical_param_tensor(tensor); + }; + if (!sd::partition_graph_cut_layer_split(get_desc().c_str(), + gf, + plan, + split_backends, + graph_cut_layer_split_backend_vram_limits_, + max_graph_vram_bytes, + graph_cut_layer_split_assignments_, + canonicalize_param, + &assignment)) { + return false; + } + + for (size_t i = 0; i < split_backends.size(); i++) { + if (assignment.tensors_by_backend[i].empty()) { + continue; + } + if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) { + LOG_ERROR("%s graph-cut layer split failed to assign params to %s", + get_desc().c_str(), + sd::layer_split_backend_device_display_name(split_backends[i]).c_str()); + return false; + } + } + + graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments); + sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment); + + return true; + } + struct PersistentExternalBinding { ggml_backend_buffer_t buffer = nullptr; void* data = nullptr; @@ -2972,6 +3112,11 @@ struct GGMLRunner { GGML_ASSERT(gf != nullptr); rebuild_params_tensor_set(); + if (!assign_graph_cut_layer_split_backends(gf)) { + free_compute_ctx(); + return std::nullopt; + } + if (can_attempt_graph_cut_segmented_compute()) { GraphCutPlan plan; if (!resolve_graph_cut_plan(gf, &plan)) { @@ -3025,6 +3170,22 @@ struct GGMLRunner { stream_layers_enabled = enabled; } + void set_graph_cut_layer_split_enabled(bool enabled) { + graph_cut_layer_split_enabled = enabled; + if (!enabled) { + graph_cut_layer_split_assignments_.clear(); + graph_cut_layer_split_node_assignments_.clear(); + graph_cut_layer_split_primary_notice_logged_ = false; + } + } + + void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) { + graph_cut_layer_split_backend_vram_limits_ = limits; + graph_cut_layer_split_assignments_.clear(); + graph_cut_layer_split_node_assignments_.clear(); + graph_cut_layer_split_primary_notice_logged_ = false; + } + void set_runtime_backends(const std::vector& backends) { extra_runtime_backends.clear(); for (ggml_backend_t backend : backends) { @@ -3036,6 +3197,9 @@ struct GGMLRunner { extra_runtime_backends.push_back(backend); } } + graph_cut_layer_split_assignments_.clear(); + graph_cut_layer_split_node_assignments_.clear(); + graph_cut_layer_split_primary_notice_logged_ = false; if (is_multi_device() && stream_layers_enabled) { LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring", get_desc().c_str()); diff --git a/src/core/layer_split_partition.cpp b/src/core/layer_split_partition.cpp index 951965e18..654b3056b 100644 --- a/src/core/layer_split_partition.cpp +++ b/src/core/layer_split_partition.cpp @@ -1,9 +1,11 @@ #include "core/layer_split_partition.h" #include -#include #include #include +#include +#include +#include #include "core/util.h" @@ -62,160 +64,194 @@ namespace sd { return name != nullptr ? name : "unknown"; } - static bool layer_split_backend_supports_tensor(ggml_backend_t backend, const ggml_tensor* tensor) { - return backend != nullptr && tensor != nullptr && ggml_backend_supports_op(backend, tensor); + static size_t graph_cut_layer_split_backend_vram_limit(const std::vector& backend_vram_limits, + size_t backend_index, + size_t primary_backend_vram_limit) { + if (backend_index < backend_vram_limits.size()) { + return backend_vram_limits[backend_index]; + } + return backend_index == 0 ? primary_backend_vram_limit : 0; } - static size_t layer_split_supported_target(const std::string& desc, - const std::string& tensor_name, - const ggml_tensor* tensor, - const std::vector& backends, - size_t preferred) { - if (tensor == nullptr || backends.empty()) { - return preferred; - } - size_t preferred_safe = std::min(preferred, backends.size() - 1); - if (layer_split_backend_supports_tensor(backends[preferred_safe], tensor)) { - return preferred_safe; - } + static std::vector graph_cut_layer_split_backend_capacities(const std::vector& backends, + const std::vector& backend_vram_limits, + size_t primary_backend_vram_limit) { + std::vector capacities(backends.size(), std::numeric_limits::max() / 4); + constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024; for (size_t i = 0; i < backends.size(); i++) { - if (layer_split_backend_supports_tensor(backends[i], tensor)) { - LOG_WARN("%s layer split: moving tensor '%s' from %s to %s because the preferred backend cannot run op=%s type=%s nbytes=%.2f MB", - desc.c_str(), - tensor_name.c_str(), - layer_split_backend_device_display_name(backends[preferred_safe]).c_str(), - layer_split_backend_device_display_name(backends[i]).c_str(), - ggml_op_name(tensor->op), - ggml_type_name(tensor->type), - ggml_nbytes(tensor) / (1024.0 * 1024.0)); - return i; + ggml_backend_dev_t dev = ggml_backend_get_device(backends[i]); + size_t free_bytes = 0, total_bytes = 0; + if (dev != nullptr) { + ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); + } + if (free_bytes > 0) { + capacities[i] = std::max((int64_t)free_bytes - compute_headroom_bytes, 0); + } + size_t limit_bytes = graph_cut_layer_split_backend_vram_limit(backend_vram_limits, + i, + primary_backend_vram_limit); + if (limit_bytes > 0) { + capacities[i] = std::min(capacities[i], (int64_t)limit_bytes); } } - LOG_WARN("%s layer split: tensor '%s' is not supported by any split backend: op=%s type=%s nbytes=%.2f MB", - desc.c_str(), - tensor_name.c_str(), - ggml_op_name(tensor->op), - ggml_type_name(tensor->type), - ggml_nbytes(tensor) / (1024.0 * 1024.0)); - return preferred_safe; + return capacities; } - std::vector> partition_layer_split_tensors( - const std::string& desc, - const std::map& tensors, - const std::map& split_tensors, - const std::vector& backends) { - std::vector> partitions(backends.size()); - if (backends.empty()) { - LOG_WARN("%s: no backend available for a layer split", desc.c_str()); - return partitions; - } + bool partition_graph_cut_layer_split(const char* desc, + ggml_cgraph* gf, + const sd::ggml_graph_cut::Plan& plan, + const std::vector& split_backends, + const std::vector& backend_vram_limits, + size_t primary_backend_vram_limit, + std::unordered_map& param_assignments, + const std::function& canonical_param_tensor, + GraphCutLayerSplitAssignment* assignment_out) { + GGML_ASSERT(gf != nullptr); + GGML_ASSERT(assignment_out != nullptr); + GGML_ASSERT(canonical_param_tensor != nullptr); + GGML_ASSERT(!split_backends.empty()); - std::map block_bytes; - std::map non_block_targets; - std::vector other_bytes_by_backend(backends.size(), 0); - int64_t total_block_bytes = 0; - int64_t total_other_bytes = 0; - int n_blocks = 0; - for (const auto& kv : tensors) { - int64_t bytes = (int64_t)ggml_nbytes(kv.second); - int idx = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1; - if (idx >= 0) { - block_bytes[idx] += bytes; - total_block_bytes += bytes; - n_blocks = std::max(n_blocks, idx + 1); - } else { - size_t target = layer_split_supported_target(desc, kv.first, kv.second, backends, 0); - non_block_targets[kv.first] = target; - other_bytes_by_backend[target] += bytes; - total_other_bytes += bytes; - } - } - if (n_blocks == 0) { - LOG_WARN("%s: no transformer blocks found for a layer split; keeping tensors on compatible backends starting from %s", - desc.c_str(), - layer_split_backend_device_display_name(backends[0]).c_str()); - for (const auto& kv : tensors) { - size_t target = 0; - auto target_it = non_block_targets.find(kv.first); - if (target_it != non_block_targets.end()) { - target = target_it->second; + GraphCutLayerSplitAssignment assignment; + assignment.segment_count = plan.segments.size(); + assignment.tensors_by_backend.resize(split_backends.size()); + assignment.bytes_by_backend.resize(split_backends.size(), 0); + assignment.first_segment_by_backend.resize(split_backends.size(), plan.segments.size()); + assignment.last_segment_by_backend.resize(split_backends.size(), 0); + + std::vector> segment_params(plan.segments.size()); + std::vector segment_param_bytes(plan.segments.size(), 0); + std::unordered_set seen_params; + for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) { + std::vector params = sd::ggml_graph_cut::param_tensors(gf, plan.segments[seg_idx]); + for (ggml_tensor* raw_param : params) { + ggml_tensor* param = canonical_param_tensor(raw_param); + if (param == nullptr || !seen_params.insert(param).second) { + continue; } - partitions[target][kv.first] = kv.second; + segment_params[seg_idx].push_back(param); + segment_param_bytes[seg_idx] += (int64_t)ggml_nbytes(param); } - return partitions; } - // Reserve compute headroom and subtract each device's actual non-block - // bytes from its block budget. - constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024; - std::vector device_weights(backends.size(), 1.0); - double weight_sum = 0.0; - for (size_t i = 0; i < backends.size(); i++) { - ggml_backend_dev_t dev = ggml_backend_get_device(backends[i]); - size_t free_bytes = 0, total_bytes = 0; - if (dev != nullptr) { - ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); - } - // Keep a small share even for tight devices instead of dropping them. - int64_t usable_bytes = std::max((int64_t)free_bytes - compute_headroom_bytes, - (int64_t)free_bytes / 8); - device_weights[i] = usable_bytes > 0 ? (double)usable_bytes : 1.0; - weight_sum += device_weights[i]; + int64_t total_param_bytes = 0; + for (int64_t bytes : segment_param_bytes) { + total_param_bytes += bytes; } - - std::vector block_budgets(backends.size(), 0); - const int64_t total_bytes = total_block_bytes + total_other_bytes; - for (size_t i = 0; i < backends.size(); i++) { - int64_t budget = (int64_t)((double)total_bytes * device_weights[i] / weight_sum); - budget = std::max(budget - other_bytes_by_backend[i], 0); - block_budgets[i] = budget; + if (total_param_bytes <= 0) { + LOG_ERROR("%s graph-cut layer split found no graph params to assign", desc); + return false; } - std::vector boundaries(backends.size(), n_blocks); - size_t current = 0; - int64_t used = 0; - for (int b = 0; b < n_blocks; b++) { - int64_t bytes = block_bytes.count(b) != 0 ? block_bytes[b] : 0; - if (current + 1 < backends.size() && used > 0 && used + bytes > block_budgets[current]) { - boundaries[current] = b; - current++; - used = 0; + std::vector backend_capacities = graph_cut_layer_split_backend_capacities(split_backends, + backend_vram_limits, + primary_backend_vram_limit); + + std::vector backend_by_segment(plan.segments.size(), split_backends[0]); + size_t current_backend = 0; + int64_t current_used = 0; + for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) { + int64_t bytes = segment_param_bytes[seg_idx]; + while (current_backend + 1 < split_backends.size() && + bytes > 0 && + current_used + bytes > backend_capacities[current_backend]) { + current_backend++; + current_used = 0; + } + if (bytes > 0 && current_used + bytes > backend_capacities[current_backend]) { + LOG_ERROR("%s graph-cut layer split: segment %zu needs %.1f MB on %s, but only %.1f MB is available under current VRAM limits", + desc, + seg_idx, + (current_used + bytes) / (1024.0 * 1024.0), + layer_split_backend_device_display_name(split_backends[current_backend]).c_str(), + backend_capacities[current_backend] / (1024.0 * 1024.0)); + return false; + } + current_used += bytes; + backend_by_segment[seg_idx] = split_backends[current_backend]; + + for (ggml_tensor* param : segment_params[seg_idx]) { + ggml_backend_t target_backend = split_backends[current_backend]; + auto assigned_it = param_assignments.find(param); + if (assigned_it == param_assignments.end()) { + param_assignments[param] = target_backend; + assignment.has_new_param_assignment = true; + } else { + target_backend = assigned_it->second; + } + + auto backend_it = std::find(split_backends.begin(), split_backends.end(), target_backend); + if (backend_it == split_backends.end()) { + LOG_ERROR("%s graph-cut layer split tensor '%s' is assigned to an unavailable backend", + desc, + ggml_get_name(param)); + return false; + } + size_t backend_idx = (size_t)std::distance(split_backends.begin(), backend_it); + assignment.first_segment_by_backend[backend_idx] = std::min(assignment.first_segment_by_backend[backend_idx], seg_idx); + assignment.last_segment_by_backend[backend_idx] = std::max(assignment.last_segment_by_backend[backend_idx], seg_idx + 1); + assignment.tensors_by_backend[backend_idx].push_back(param); + assignment.bytes_by_backend[backend_idx] += (int64_t)ggml_nbytes(param); } - used += bytes; } - for (const auto& kv : tensors) { - size_t target = 0; - int idx = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1; - if (idx >= 0) { - while (target < boundaries.size() && idx >= boundaries[target]) { - target++; + const int n_nodes = ggml_graph_n_nodes(gf); + for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) { + ggml_backend_t backend = backend_by_segment[seg_idx]; + const auto& segment = plan.segments[seg_idx]; + for (int node_index : segment.internal_node_indices) { + if (node_index < 0 || node_index >= n_nodes) { + continue; } - target = std::min(target, backends.size() - 1); - target = layer_split_supported_target(desc, kv.first, kv.second, backends, target); - } else { - auto target_it = non_block_targets.find(kv.first); - if (target_it != non_block_targets.end()) { - target = target_it->second; + ggml_tensor* node = ggml_graph_node(gf, node_index); + if (node != nullptr) { + assignment.node_assignments[node] = backend; + } + } + for (int node_index : segment.output_node_indices) { + if (node_index < 0 || node_index >= n_nodes) { + continue; + } + ggml_tensor* node = ggml_graph_node(gf, node_index); + if (node != nullptr) { + assignment.node_assignments[node] = backend; } } - partitions[target][kv.first] = kv.second; } - int range_start = 0; - for (size_t i = 0; i < backends.size(); i++) { - int range_end = boundaries[i]; - const char* non_block_suffix = other_bytes_by_backend[i] > 0 ? " + non-block tensors" : ""; - LOG_INFO("%s layer split: %s <- blocks [%d, %d)%s", - desc.c_str(), - layer_split_backend_device_display_name(backends[i]).c_str(), - range_start, - range_end, - non_block_suffix); - range_start = range_end; + *assignment_out = std::move(assignment); + return true; + } + + void log_graph_cut_layer_split_assignment(const char* desc, + const std::vector& split_backends, + const GraphCutLayerSplitAssignment& assignment) { + for (size_t i = 0; i < split_backends.size(); i++) { + if (i >= assignment.tensors_by_backend.size() || + assignment.tensors_by_backend[i].empty()) { + continue; + } + size_t first_segment = assignment.first_segment_by_backend[i] == assignment.segment_count + ? 0 + : assignment.first_segment_by_backend[i]; + size_t last_segment = assignment.last_segment_by_backend[i]; + if (assignment.has_new_param_assignment) { + LOG_INFO("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB", + desc, + layer_split_backend_device_display_name(split_backends[i]).c_str(), + first_segment, + last_segment, + assignment.tensors_by_backend[i].size(), + assignment.bytes_by_backend[i] / (1024.0 * 1024.0)); + } else { + LOG_DEBUG("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB", + desc, + layer_split_backend_device_display_name(split_backends[i]).c_str(), + first_segment, + last_segment, + assignment.tensors_by_backend[i].size(), + assignment.bytes_by_backend[i] / (1024.0 * 1024.0)); + } } - return partitions; } } // namespace sd diff --git a/src/core/layer_split_partition.h b/src/core/layer_split_partition.h index 5c2079388..9450b379e 100644 --- a/src/core/layer_split_partition.h +++ b/src/core/layer_split_partition.h @@ -1,23 +1,43 @@ #ifndef __SD_CORE_LAYER_SPLIT_PARTITION_H__ #define __SD_CORE_LAYER_SPLIT_PARTITION_H__ -#include +#include +#include #include +#include #include #include "ggml-backend.h" #include "ggml.h" +#include "core/ggml_graph_cut.h" + namespace sd { + struct GraphCutLayerSplitAssignment { + std::vector> tensors_by_backend; + std::vector bytes_by_backend; + std::vector first_segment_by_backend; + std::vector last_segment_by_backend; + std::unordered_map node_assignments; + size_t segment_count = 0; + bool has_new_param_assignment = false; + }; + std::string layer_split_backend_device_display_name(ggml_backend_t backend); int layer_split_tensor_block_index(const std::string& name); - - std::vector> partition_layer_split_tensors( - const std::string& desc, - const std::map& tensors, - const std::map& split_tensors, - const std::vector& backends); + bool partition_graph_cut_layer_split(const char* desc, + ggml_cgraph* gf, + const sd::ggml_graph_cut::Plan& plan, + const std::vector& split_backends, + const std::vector& backend_vram_limits, + size_t primary_backend_vram_limit, + std::unordered_map& param_assignments, + const std::function& canonical_param_tensor, + GraphCutLayerSplitAssignment* assignment_out); + void log_graph_cut_layer_split_assignment(const char* desc, + const std::vector& split_backends, + const GraphCutLayerSplitAssignment& assignment); } // namespace sd diff --git a/src/model_manager.cpp b/src/model_manager.cpp index c5bddcc9e..3a98bd545 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -134,7 +134,8 @@ bool ModelManager::register_param_tensors(const std::string& desc, ggml_backend_t compute_backend, ggml_backend_t params_backend, size_t* registered_tensor_size, - bool allow_split_buffer) { + bool allow_split_buffer, + bool params_follow_compute_backend) { if (desc.empty()) { LOG_ERROR("model manager tensor desc is empty"); return false; @@ -158,14 +159,15 @@ bool ModelManager::register_param_tensors(const std::string& desc, } ggml_set_name(tensor, name.c_str()); - auto state = std::make_unique(); - state->name = name; - state->tensor = tensor; - state->desc = desc; - state->residency_mode = residency_mode; - state->compute_backend = compute_backend; - state->params_backend = params_backend; - state->allow_split_buffer = allow_split_buffer; + auto state = std::make_unique(); + state->name = name; + state->tensor = tensor; + state->desc = desc; + state->residency_mode = residency_mode; + state->compute_backend = compute_backend; + state->params_backend = params_backend; + state->allow_split_buffer = allow_split_buffer; + state->params_follow_compute_backend = params_follow_compute_backend; new_states.push_back(std::move(state)); } @@ -919,6 +921,54 @@ bool ModelManager::resolve_required_tensor_states(const std::vector& tensors, + ggml_backend_t compute_backend) { + if (tensors.empty()) { + return true; + } + if (compute_backend == nullptr) { + LOG_ERROR("model manager cannot assign tensors to a null compute backend"); + return false; + } + + std::vector required_states; + if (!resolve_required_tensor_states(tensors, required_states)) { + return false; + } + + for (TensorState* state : required_states) { + if (state == nullptr || state->tensor == nullptr) { + continue; + } + + const bool params_follow_compute = state->params_follow_compute_backend || + state->residency_mode == ResidencyMode::Disk; + const bool compute_changes = state->compute_backend != compute_backend; + const bool params_changes = params_follow_compute && state->params_backend != compute_backend; + if (!compute_changes && !params_changes) { + continue; + } + + if (state->active_prepare_count > 0 || state->staged_to_compute_backend) { + LOG_ERROR("model manager cannot move active tensor '%s' to another compute backend", + state->name.c_str()); + return false; + } + if (params_changes && state->loaded_to_params_backend) { + LOG_ERROR("model manager cannot move loaded tensor '%s' to another params backend", + state->name.c_str()); + return false; + } + + state->compute_backend = compute_backend; + if (params_follow_compute) { + state->params_backend = compute_backend; + } + } + + return true; +} + bool ModelManager::prepare_params(const std::vector& tensors) { if (tensors.empty()) { return true; diff --git a/src/model_manager.h b/src/model_manager.h index 8ec7a7a17..d80032614 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -33,11 +33,12 @@ class ModelManager : public RunnerWeightManager { ggml_tensor* tensor = nullptr; std::string desc; - ResidencyMode residency_mode = ResidencyMode::ParamBackend; - ggml_backend_t compute_backend = nullptr; - ggml_backend_t params_backend = nullptr; - bool allow_split_buffer = false; - bool metadata_validated = false; + ResidencyMode residency_mode = ResidencyMode::ParamBackend; + ggml_backend_t compute_backend = nullptr; + ggml_backend_t params_backend = nullptr; + bool allow_split_buffer = false; + bool params_follow_compute_backend = false; + bool metadata_validated = false; int active_prepare_count = 0; @@ -129,8 +130,9 @@ class ModelManager : public RunnerWeightManager { ResidencyMode residency_mode, ggml_backend_t compute_backend, ggml_backend_t params_backend, - size_t* registered_tensor_size = nullptr, - bool allow_split_buffer = false); + size_t* registered_tensor_size = nullptr, + bool allow_split_buffer = false, + bool params_follow_compute_backend = false); template bool register_runner_params(const std::string& desc, @@ -170,6 +172,8 @@ class ModelManager : public RunnerWeightManager { bool validate_registered_tensors(); bool load_all_params_eagerly(); + bool assign_compute_backend(const std::vector& tensors, + ggml_backend_t compute_backend) override; bool prepare_params(const std::vector& tensors) override; void release_compute_backend_params(const std::vector& tensors) override; void release_params_backend_params(const std::vector& tensors) override; diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 466a88fa5..28628be7c 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -268,6 +268,15 @@ class StableDiffusionGGML { return max_vram_assignment.bytes_for_backend(backend_for(module)); } + std::vector layer_split_vram_limits_for_backends(const std::vector& backends) { + std::vector limits; + limits.reserve(backends.size()); + for (ggml_backend_t backend : backends) { + limits.push_back(max_vram_assignment.bytes_for_backend(backend)); + } + return limits; + } + bool ensure_backend_pair(SDBackendModule module) { if (backend_for(module) == nullptr) { return false; @@ -427,8 +436,9 @@ class StableDiffusionGGML { params_mem_size); } - // Register each layer-split partition with its compute backend; the - // ModelManager handles allocation, staging, and LoRA by backend. + // Register graph-cut layer-split tensors on the primary backend first. + // The first real graph assigns each param tensor to a runtime backend + // before weights are loaded or staged. template bool register_layer_split_runner_params(const std::string& desc, const std::shared_ptr& model, @@ -459,52 +469,29 @@ class StableDiffusionGGML { params_mem_size); } - std::map split_tensors; - if constexpr (std::is_base_of_v) { - model->get_layer_split_param_tensors(split_tensors); - } else { - split_tensors = group_tensors; - } - - auto partitions = sd::partition_layer_split_tensors(desc, group_tensors, split_tensors, module_backends); - bool is_split = false; - for (size_t i = 1; i < partitions.size(); i++) { - if (!partitions[i].empty()) { - is_split = true; - break; - } - } - if (!is_split) { - return model_manager->register_param_tensors(desc, - std::move(group_tensors), - residency_mode, - module_backends[0], - params_backend_for(module), - params_mem_size); - } - model->set_runtime_backends(module_backends); + model->set_graph_cut_layer_split_backend_vram_limits(layer_split_vram_limits_for_backends(module_backends)); + model->set_graph_cut_layer_split_enabled(true); const bool params_follow_runtime = backend_manager.params_backend_follows_runtime(module) || backend_manager.params_backend_is_disk(module); - for (size_t i = 0; i < module_backends.size(); i++) { - if (partitions[i].empty()) { - continue; - } - ggml_backend_t partition_params_backend = - params_follow_runtime ? module_backends[i] : params_backend_for(module); - if (partition_params_backend == nullptr) { - return false; - } - if (!model_manager->register_param_tensors(desc, - std::move(partitions[i]), - residency_mode, - module_backends[i], - partition_params_backend, - params_mem_size)) { - return false; - } + ggml_backend_t initial_params_backend = params_follow_runtime ? module_backends[0] : params_backend_for(module); + if (initial_params_backend == nullptr) { + return false; } - return true; + + LOG_INFO("%s graph-cut layer split: deferring %zu tensors across %zu runtime backends until first graph", + desc.c_str(), + group_tensors.size(), + module_backends.size()); + + return model_manager->register_param_tensors(desc, + std::move(group_tensors), + residency_mode, + module_backends[0], + initial_params_backend, + params_mem_size, + false, + params_follow_runtime); } bool init_backend() { @@ -529,6 +516,16 @@ class StableDiffusionGGML { return false; } + bool graph_cut_layer_split_active() { + for (SDBackendModule module : {SDBackendModule::DIFFUSION, SDBackendModule::TE}) { + if (backend_manager.split_mode(module) == SDSplitMode::LAYER && + backend_manager.runtime_backends(module).size() > 1) { + return true; + } + } + return false; + } + std::shared_ptr get_rng(rng_type_t rng_type) { if (rng_type == STD_DEFAULT_RNG) { return std::make_shared(); @@ -785,6 +782,10 @@ class StableDiffusionGGML { LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring"); stream_layers = false; } + if (eager_load && graph_cut_layer_split_active()) { + LOG_WARN("--eager-load is not supported with graph-cut layer split; weights will be prepared lazily"); + eager_load = false; + } std::map wtype_stat = model_loader.get_wtype_stat(); std::map conditioner_wtype_stat = model_loader.get_conditioner_wtype_stat(); diff --git a/src/weight_manager.h b/src/weight_manager.h index 28d6cf5c4..82f6d03e4 100644 --- a/src/weight_manager.h +++ b/src/weight_manager.h @@ -3,10 +3,14 @@ #include +#include "ggml-backend.h" + struct ggml_tensor; struct RunnerWeightManager { virtual ~RunnerWeightManager() = default; + virtual bool assign_compute_backend(const std::vector& tensors, + ggml_backend_t compute_backend) = 0; virtual bool prepare_params(const std::vector& tensors) = 0; virtual void release_compute_backend_params(const std::vector& tensors) = 0; virtual void release_params_backend_params(const std::vector& tensors) = 0;