diff --git a/docs/backend.md b/docs/backend.md index 29ac8031b..805143a3a 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -51,6 +51,40 @@ Module names are case-insensitive. Hyphens and underscores in module names are i sd-cli -m model.safetensors -p "a cat" --backend all=cuda0,te=cpu ``` +## Multiple devices per module (layer split) + +A `--backend` module assignment can list several devices separated by `&`: + +```shell +sd-cli -m model.safetensors -p "a cat" --backend "diffusion=cuda0&cuda1" +``` + +The module's transformer blocks are then distributed across the listed devices +in contiguous ranges sized proportionally to each device's free memory (minus a +compute-buffer headroom of about 2 GiB per device), and the +module's graphs are executed with a `ggml_backend_sched` that runs each block +on the device holding its weights, copying the residual stream at the range +boundaries. The first device in the list is the module's main device: it also +holds the non-block tensors (embeddings, final norms, small sub-runners such as +CLIP models or projectors) and the graph inputs/outputs. + +Layer split is supported for the `diffusion` and `te` modules. For `te` it +applies to the dominant text encoder (`t5xxl` or the LLM); other modules accept +only a single device. If the module has no recognizable transformer blocks, the +assignment falls back to the first listed device. + +`--params-backend` accepts no device lists. If the module has no explicit +params assignment, each block range's parameters are loaded directly to (and, +with `--params-backend diffusion=disk`, released directly from) its own device; +an explicit assignment such as `te=cpu` keeps the parameters on that backend +and stages each range to its device on demand. + +Layer split cannot be combined with `--max-vram` graph-cut segmentation or +`--stream-layers` for the split module; those are single-device mechanisms and +are disabled for it. + +Use `--list-devices` to see the device names available on the system. + ## Modules | Module | Purpose | Accepted names | diff --git a/examples/common/common.cpp b/examples/common/common.cpp index ac1cf32f6..7256e25b3 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -663,6 +663,18 @@ ArgOptions SDContextParams::get_options() { "but it usually offers faster inference speed and, in some cases, lower memory usage. " "The at_runtime mode, on the other hand, is exactly the opposite.", on_lora_apply_mode_arg}, + {"", + "--list-devices", + "list available ggml backend devices (one 'namedescription' per line) and exit; " + "the names are the device names accepted by --backend and --params-backend", + [](int /*argc*/, const char** /*argv*/, int /*index*/) { + size_t device_list_size = sd_list_devices(nullptr, 0); + std::vector devices(device_list_size + 1); + sd_list_devices(devices.data(), devices.size()); + fputs(devices.data(), stdout); + std::exit(0); + return 0; + }}, }; return options; diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 26bb61ef8..1acb0a22d 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -534,6 +534,12 @@ SD_API void disable_imatrix_collection(void); SD_API const char* sd_commit(void); SD_API const char* sd_version(void); +// List available ggml backend devices, one `namedescription` per line. +// The names are the device names accepted by the --backend / --params-backend +// assignment specs. Returns the number of bytes required, excluding the null +// terminator. Passing nullptr or buffer_size 0 only queries the required size. +SD_API size_t sd_list_devices(char* buffer, size_t buffer_size); + // for C API, caller needs to call free_sd_images to free the memory after use // This helps avoid CRT problems on Windows when memory is allocated in the library but freed in the caller, which may use a different CRT. SD_API void free_sd_images(sd_image_t* result_images, int num_images); diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index d63303a82..da5aa74f0 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -116,6 +116,8 @@ struct Conditioner { virtual void get_param_tensors(std::map& tensors) = 0; 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 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) {} virtual void runner_done() {} @@ -635,6 +637,18 @@ struct SD3CLIPEmbedder : public Conditioner { } } + void set_runtime_backends(const std::vector& backends) override { + if (t5) { + t5->set_runtime_backends(backends); + } + } + + void get_layer_split_param_tensors(std::map& tensors) override { + if (t5) { + t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + } + } + void set_flash_attention_enabled(bool enabled) override { if (clip_l) { clip_l->set_flash_attention_enabled(enabled); @@ -994,6 +1008,18 @@ struct FluxCLIPEmbedder : public Conditioner { } } + void set_runtime_backends(const std::vector& backends) override { + if (t5) { + t5->set_runtime_backends(backends); + } + } + + void get_layer_split_param_tensors(std::map& tensors) override { + if (t5) { + t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + } + } + void set_flash_attention_enabled(bool enabled) override { if (clip_l) { clip_l->set_flash_attention_enabled(enabled); @@ -1226,6 +1252,18 @@ struct T5CLIPEmbedder : public Conditioner { } } + void set_runtime_backends(const std::vector& backends) override { + if (t5) { + t5->set_runtime_backends(backends); + } + } + + void get_layer_split_param_tensors(std::map& tensors) override { + if (t5) { + t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + } + } + void set_flash_attention_enabled(bool enabled) override { if (t5) { t5->set_flash_attention_enabled(enabled); @@ -1418,6 +1456,18 @@ struct MiniT2IConditioner : public Conditioner { } } + void set_runtime_backends(const std::vector& backends) override { + if (t5) { + t5->set_runtime_backends(backends); + } + } + + void get_layer_split_param_tensors(std::map& tensors) override { + if (t5) { + t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + } + } + void set_flash_attention_enabled(bool enabled) override { if (t5) { t5->set_flash_attention_enabled(enabled); @@ -1502,6 +1552,14 @@ struct AnimaConditioner : public Conditioner { llm->set_stream_layers_enabled(enabled); } + void set_runtime_backends(const std::vector& backends) override { + llm->set_runtime_backends(backends); + } + + void get_layer_split_param_tensors(std::map& tensors) override { + llm->get_param_tensors(tensors, "text_encoders.llm"); + } + void set_flash_attention_enabled(bool enabled) override { llm->set_flash_attention_enabled(enabled); } @@ -1647,6 +1705,14 @@ struct LLMEmbedder : public Conditioner { llm->set_stream_layers_enabled(enabled); } + void set_runtime_backends(const std::vector& backends) override { + llm->set_runtime_backends(backends); + } + + void get_layer_split_param_tensors(std::map& tensors) override { + llm->get_param_tensors(tensors, "text_encoders.llm"); + } + void set_flash_attention_enabled(bool enabled) override { llm->set_flash_attention_enabled(enabled); } @@ -2316,6 +2382,14 @@ struct LTXAVEmbedder : public Conditioner { projector->set_max_graph_vram_bytes(max_vram_bytes); } + void set_runtime_backends(const std::vector& backends) override { + llm->set_runtime_backends(backends); + } + + void get_layer_split_param_tensors(std::map& tensors) override { + llm->get_param_tensors(tensors, "text_encoders.llm"); + } + void set_weight_adapter(const std::shared_ptr& adapter) override { llm->set_weight_adapter(adapter); projector->set_weight_adapter(adapter); diff --git a/src/core/ggml_extend.hpp b/src/core/ggml_extend.hpp index c4138a235..1dcf423f1 100644 --- a/src/core/ggml_extend.hpp +++ b/src/core/ggml_extend.hpp @@ -1746,6 +1746,11 @@ struct GGMLRunner { bool stream_layers_enabled = false; size_t observed_max_effective_budget_ = 0; + std::vector extra_runtime_backends; // borrowed (SDBackendManager-owned) + ggml_backend_sched_t sched = nullptr; // owned, multi-device only + ggml_backend_t cpu_fallback_backend = nullptr; // owned, sched requires a trailing CPU backend + bool multi_device_eval_callback_warned = false; + std::shared_ptr weight_adapter = nullptr; std::weak_ptr weight_manager; std::unordered_set kept_compute_param_tensor_set; @@ -2013,7 +2018,121 @@ struct GGMLRunner { return true; } + // Pass explicit buffer types: synthesized defaults can make CUDA devices + // report supporting each other's buffers and skip a required copy. + bool ensure_sched(ggml_cgraph* gf) { + if (sched != nullptr) { + return true; + } + std::vector backends; + backends.reserve(extra_runtime_backends.size() + 2); + backends.push_back(runtime_backend); + for (ggml_backend_t backend : extra_runtime_backends) { + backends.push_back(backend); + } + if (cpu_fallback_backend == nullptr && !sd_backend_is_cpu(runtime_backend)) { + cpu_fallback_backend = sd_backend_cpu_init(); + } + if (cpu_fallback_backend != nullptr) { + backends.push_back(cpu_fallback_backend); + } + + std::vector bufts; + bufts.reserve(backends.size()); + ggml_backend_dev_t main_dev = ggml_backend_get_device(runtime_backend); + for (ggml_backend_t backend : backends) { + ggml_backend_buffer_type_t buft = nullptr; + if (backend == cpu_fallback_backend && main_dev != nullptr) { + buft = ggml_backend_dev_host_buffer_type(main_dev); + } + if (buft == nullptr) { + buft = ggml_backend_get_default_buffer_type(backend); + } + bufts.push_back(buft); + } + + size_t graph_size = MAX_GRAPH_SIZE; + if (gf != nullptr) { + graph_size = std::max(graph_size, (size_t)ggml_graph_n_nodes(gf)); + } + sched = ggml_backend_sched_new(backends.data(), + bufts.data(), + (int)backends.size(), + graph_size, + /*parallel=*/false, + /*op_offload=*/false); + if (sched == nullptr) { + LOG_ERROR("%s: failed to create backend sched", get_desc().c_str()); + return false; + } + return true; + } + + ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const { + if (tensor == nullptr || tensor->buffer == nullptr) { + return nullptr; + } + if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS || + ggml_backend_buffer_is_host(tensor->buffer)) { + return nullptr; + } + ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); + if (dev == nullptr) { + return nullptr; + } + if (ggml_backend_get_device(runtime_backend) == dev) { + return runtime_backend; + } + for (ggml_backend_t backend : extra_runtime_backends) { + if (ggml_backend_get_device(backend) == dev) { + return backend; + } + } + return nullptr; + } + + // Weightless ops have no scheduler anchor, so pin them to the most recent + // weight device. Views must stay unpinned or cross-device copies can be + // skipped for their consumers. + void pin_multi_device_nodes(ggml_cgraph* gf) { + if (sched == nullptr || gf == nullptr) { + return; + } + 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); + 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->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || + node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) { + continue; + } + if (ggml_backend_supports_op(current, node)) { + ggml_backend_sched_set_tensor_backend(sched, node, current); + } + } + } + + bool is_multi_device() const { + return !extra_runtime_backends.empty(); + } + bool alloc_compute_buffer(ggml_cgraph* gf) { + if (is_multi_device()) { + // 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 + // later ggml_backend_sched_alloc_graph would split the already + // 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. + return ensure_sched(gf); + } if (compute_allocr != nullptr) { return true; } @@ -2229,12 +2348,14 @@ struct GGMLRunner { plan.valid && max_graph_vram_bytes > 0 && plan.segments.size() > 1 && - !sd_backend_is_cpu(runtime_backend); + !sd_backend_is_cpu(runtime_backend) && + !is_multi_device(); } bool can_attempt_graph_cut_segmented_compute() const { return max_graph_vram_bytes > 0 && - !sd_backend_is_cpu(runtime_backend); + !sd_backend_is_cpu(runtime_backend) && + !is_multi_device(); } bool resolve_graph_cut_plan(ggml_cgraph* gf, @@ -2490,7 +2611,14 @@ struct GGMLRunner { }; ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer); - if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) { + if (is_multi_device()) { + 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)) { + LOG_ERROR("%s sched alloc compute graph failed", get_desc().c_str()); + return std::nullopt; + } + } else if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) { LOG_ERROR("%s alloc compute graph failed", get_desc().c_str()); return std::nullopt; } @@ -2499,11 +2627,27 @@ struct GGMLRunner { if (sd_backend_is_cpu(runtime_backend)) { sd_backend_cpu_set_n_threads(runtime_backend, n_threads); } + if (cpu_fallback_backend != nullptr) { + sd_backend_cpu_set_n_threads(cpu_fallback_backend, n_threads); + } - ggml_status status = sd_backend_graph_compute_with_eval_callback(runtime_backend, - gf, - sd_get_backend_eval_callback(), - sd_get_backend_eval_callback_data()); + ggml_status status; + if (is_multi_device()) { + 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", + get_desc().c_str()); + multi_device_eval_callback_warned = true; + } + status = ggml_backend_sched_graph_compute(sched, gf); + if (status == GGML_STATUS_SUCCESS) { + ggml_backend_sched_synchronize(sched); + } + } else { + status = sd_backend_graph_compute_with_eval_callback(runtime_backend, + gf, + sd_get_backend_eval_callback(), + sd_get_backend_eval_callback_data()); + } if (status != GGML_STATUS_SUCCESS) { LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); return std::nullopt; @@ -2680,6 +2824,10 @@ struct GGMLRunner { free_params_ctx(); free_compute_ctx(); free_cache_ctx_and_buffer(); + if (cpu_fallback_backend != nullptr) { + ggml_backend_free(cpu_fallback_backend); + cpu_fallback_backend = nullptr; + } } virtual GGMLRunnerContext get_context() { @@ -2720,10 +2868,20 @@ struct GGMLRunner { ggml_gallocr_free(compute_allocr); compute_allocr = nullptr; } + if (sched != nullptr) { + ggml_backend_sched_free(sched); + sched = nullptr; + } } // 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); + } backend_tensor_data_map[tensor] = data; } @@ -2859,8 +3017,31 @@ struct GGMLRunner { } void set_stream_layers_enabled(bool enabled) { + if (enabled && is_multi_device()) { + LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring", + get_desc().c_str()); + return; + } stream_layers_enabled = enabled; } + + void set_runtime_backends(const std::vector& backends) { + extra_runtime_backends.clear(); + for (ggml_backend_t backend : backends) { + if (backend == nullptr || backend == runtime_backend) { + continue; + } + if (std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), backend) == + extra_runtime_backends.end()) { + extra_runtime_backends.push_back(backend); + } + } + if (is_multi_device() && stream_layers_enabled) { + LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring", + get_desc().c_str()); + stream_layers_enabled = false; + } + } }; class GGMLBlock { diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index f29bdb696..032ae52ab 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -669,8 +669,47 @@ void SDBackendManager::reset() { params_assignment_ = {}; } +static std::vector split_device_list(const std::string& value) { + std::vector names; + for (const std::string& raw : split_copy(value, '&')) { + const std::string name = trim_copy(raw); + if (!name.empty()) { + names.push_back(name); + } + } + return names; +} + +static std::string primary_device_name(const std::string& value) { + std::vector names = split_device_list(value); + return names.empty() ? std::string() : names.front(); +} + ggml_backend_t SDBackendManager::runtime_backend(SDBackendModule module) { - return init_cached_backend(runtime_assignment_.get(module)); + return init_cached_backend(primary_device_name(runtime_assignment_.get(module))); +} + +std::vector SDBackendManager::runtime_backends(SDBackendModule module) { + std::vector backends; + for (const std::string& name : split_device_list(runtime_assignment_.get(module))) { + ggml_backend_t backend = init_cached_backend(name); + if (backend == nullptr) { + LOG_ERROR("failed to initialize backend '%s' for module %s", + name.c_str(), + sd_backend_module_name(module)); + continue; + } + if (std::find(backends.begin(), backends.end(), backend) == backends.end()) { + backends.push_back(backend); + } + } + if (backends.empty()) { + ggml_backend_t backend = runtime_backend(module); + if (backend != nullptr) { + backends.push_back(backend); + } + } + return backends; } ggml_backend_t SDBackendManager::params_backend(SDBackendModule module) { @@ -696,6 +735,10 @@ bool SDBackendManager::params_backend_is_disk(SDBackendModule module) const { return is_disk_backend_token(params_assignment_.get(module)); } +bool SDBackendManager::params_backend_follows_runtime(SDBackendModule module) const { + return params_assignment_.get(module).empty(); +} + bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule module) { ggml_backend_t backend = runtime_backend(module); if (backend == nullptr) { @@ -729,7 +772,7 @@ bool SDBackendManager::init(const char* backend_spec, } bool SDBackendManager::validate(std::string* error) const { - auto validate_runtime_name = [&](const std::string& name) -> bool { + auto validate_single_runtime_name = [&](const std::string& name) -> bool { if (is_default_backend_token(name)) { return true; } @@ -747,11 +790,41 @@ bool SDBackendManager::validate(std::string* error) const { } return false; }; + auto validate_runtime_name = [&](const std::string& name) -> bool { + if (name.find('&') == std::string::npos) { + return validate_single_runtime_name(name); + } + std::vector names = split_device_list(name); + if (names.empty()) { + if (error != nullptr) { + *error = "invalid backend device list '" + name + "'"; + } + return false; + } + for (const std::string& entry : names) { + if (is_default_backend_token(entry)) { + if (error != nullptr) { + *error = "default backend token is not allowed in a device list '" + name + "'"; + } + return false; + } + if (!validate_single_runtime_name(entry)) { + return false; + } + } + return true; + }; auto validate_params_name = [&](const std::string& name) -> bool { if (is_disk_backend_token(name)) { return true; } - return validate_runtime_name(name); + if (name.find('&') != std::string::npos) { + if (error != nullptr) { + *error = "params_backend does not accept device lists ('" + name + "')"; + } + return false; + } + return validate_single_runtime_name(name); }; if (!validate_runtime_name(runtime_assignment_.default_name) || diff --git a/src/core/ggml_extend_backend.h b/src/core/ggml_extend_backend.h index 19b71d432..92cc6b693 100644 --- a/src/core/ggml_extend_backend.h +++ b/src/core/ggml_extend_backend.h @@ -6,6 +6,7 @@ #include #include #include +#include #include "ggml-backend.h" #include "ggml.h" @@ -58,9 +59,12 @@ class SDBackendManager { ggml_backend_t runtime_backend(SDBackendModule module); ggml_backend_t params_backend(SDBackendModule module); + std::vector runtime_backends(SDBackendModule module); + bool runtime_backend_is_cpu(SDBackendModule module); bool params_backend_is_cpu(SDBackendModule module); bool params_backend_is_disk(SDBackendModule module) const; + bool params_backend_follows_runtime(SDBackendModule module) const; bool runtime_backend_supports_host_buffer(SDBackendModule module); private: diff --git a/src/core/layer_split_partition.cpp b/src/core/layer_split_partition.cpp new file mode 100644 index 000000000..322b5a316 --- /dev/null +++ b/src/core/layer_split_partition.cpp @@ -0,0 +1,194 @@ +#include "core/layer_split_partition.h" + +#include +#include +#include +#include + +#include "core/util.h" + +namespace sd { + + static int tensor_block_index(const std::string& name) { + static const char* block_keywords[] = {"transformer_blocks.", "joint_blocks.", "double_blocks.", + "single_blocks.", "blocks.", "block.", "layers."}; + for (const char* keyword : block_keywords) { + size_t pos = name.find(keyword); + if (pos == std::string::npos) { + continue; + } + pos += std::strlen(keyword); + size_t end = pos; + while (end < name.size() && name[end] >= '0' && name[end] <= '9') { + end++; + } + if (end > pos && (end == name.size() || name[end] == '.')) { + return std::atoi(name.substr(pos, end - pos).c_str()); + } + } + return -1; + } + + std::string layer_split_backend_device_display_name(ggml_backend_t backend) { + ggml_backend_dev_t dev = ggml_backend_get_device(backend); + const char* name = dev != nullptr ? ggml_backend_dev_name(dev) : ggml_backend_name(backend); + 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 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; + } + 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; + } + } + 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; + } + + 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; + } + + 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 ? 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; + } + partitions[target][kv.first] = kv.second; + } + 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]; + } + + 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; + } + + 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; + } + used += bytes; + } + + for (const auto& kv : tensors) { + size_t target = 0; + int idx = split_tensors.count(kv.first) != 0 ? tensor_block_index(kv.first) : -1; + if (idx >= 0) { + while (target < boundaries.size() && idx >= boundaries[target]) { + target++; + } + 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; + } + } + 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; + } + return partitions; + } + +} // namespace sd diff --git a/src/core/layer_split_partition.h b/src/core/layer_split_partition.h new file mode 100644 index 000000000..61d8167ce --- /dev/null +++ b/src/core/layer_split_partition.h @@ -0,0 +1,23 @@ +#ifndef __SD_CORE_LAYER_SPLIT_PARTITION_H__ +#define __SD_CORE_LAYER_SPLIT_PARTITION_H__ + +#include +#include +#include + +#include "ggml-backend.h" +#include "ggml.h" + +namespace sd { + + std::string layer_split_backend_device_display_name(ggml_backend_t backend); + + std::vector> partition_layer_split_tensors( + const std::string& desc, + const std::map& tensors, + const std::map& split_tensors, + const std::vector& backends); + +} // namespace sd + +#endif // __SD_CORE_LAYER_SPLIT_PARTITION_H__ diff --git a/src/core/util.cpp b/src/core/util.cpp index 6d2479f9f..ff53c65fc 100644 --- a/src/core/util.cpp +++ b/src/core/util.cpp @@ -4,6 +4,8 @@ #include #include #include +#include +#include #include #include #include @@ -25,6 +27,7 @@ #include #endif +#include "ggml-backend.h" #include "ggml.h" #include "stable-diffusion.h" @@ -997,3 +1000,26 @@ std::vector> split_quotation_attention( } return result; } + +size_t sd_list_devices(char* buffer, size_t buffer_size) { + if (ggml_backend_dev_count() == 0) { + // dynamic-backend builds discover their backend modules at runtime + ggml_backend_load_all(); + } + + std::ostringstream oss; + for (size_t i = 0; i < ggml_backend_dev_count(); i++) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + const char* name = ggml_backend_dev_name(dev); + const char* desc = ggml_backend_dev_description(dev); + oss << (name ? name : "") << '\t' << (desc ? desc : "") << '\n'; + } + + std::string devices = oss.str(); + if (buffer != nullptr && buffer_size > 0) { + size_t copy_size = std::min(devices.size(), buffer_size - 1); + memcpy(buffer, devices.data(), copy_size); + buffer[copy_size] = '\0'; + } + return devices.size(); +} diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index eb592f414..e239fe1a7 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -2,11 +2,14 @@ #include #include #include +#include #include +#include #include #include "core/ggml_extend.hpp" #include "core/ggml_graph_cut.h" +#include "core/layer_split_partition.h" #include "core/rng.hpp" #include "core/rng_mt19937.hpp" @@ -170,6 +173,13 @@ static float get_cache_reuse_threshold(const sd_cache_params_t& params) { /*=============================================== StableDiffusionGGML ================================================*/ +template +struct has_set_runtime_backends : std::false_type {}; +template +struct has_set_runtime_backends().set_runtime_backends( + std::declval&>()))>> : std::true_type {}; + static_assert(std::atomic::is_always_lock_free, "sd_cancel_mode_t must be lock-free"); @@ -275,14 +285,114 @@ class StableDiffusionGGML { if (model_manager == nullptr) { return true; } + ModelManager::ResidencyMode residency_mode = + backend_manager.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend; + + std::vector module_backends = backend_manager.runtime_backends(module); + if (module_backends.size() > 1) { + if constexpr (has_set_runtime_backends::value) { + if (module == SDBackendModule::DIFFUSION || module == SDBackendModule::TE) { + return register_layer_split_runner_params(desc, + model, + module, + module_backends, + std::move(group_tensors), + residency_mode, + params_mem_size); + } + } + LOG_WARN("%s module does not support multiple runtime backends; using %s", + sd_backend_module_name(module), + sd::layer_split_backend_device_display_name(module_backends[0]).c_str()); + } return model_manager->register_param_tensors(desc, std::move(group_tensors), - backend_manager.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend, + residency_mode, backend_for(module), params_backend_for(module), params_mem_size); } + // Register each layer-split partition with its compute backend; the + // ModelManager handles allocation, staging, and LoRA by backend. + template + bool register_layer_split_runner_params(const std::string& desc, + const std::shared_ptr& model, + SDBackendModule module, + const std::vector& module_backends, + std::map group_tensors, + ModelManager::ResidencyMode residency_mode, + size_t* params_mem_size) { + bool has_cpu_device = false; + for (ggml_backend_t backend : module_backends) { + has_cpu_device = has_cpu_device || sd_backend_is_cpu(backend); + } + if (has_cpu_device) { + // The scheduler reserves the CPU slot for its fallback backend, and + // CPU weight participation is what --params-backend =cpu is + // for; a CPU device in a split list is almost certainly a mistake. + LOG_WARN( + "%s: layer split across a CPU device is not supported; using %s " + "(use --params-backend %s=cpu to keep weights in RAM)", + desc.c_str(), + sd::layer_split_backend_device_display_name(module_backends[0]).c_str(), + sd_backend_module_name(module)); + return model_manager->register_param_tensors(desc, + std::move(group_tensors), + residency_mode, + module_backends[0], + params_backend_for(module), + 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); + 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; + } + } + return true; + } + bool init_backend() { std::string error; if (!backend_manager.init(backend_spec.c_str(),