From 68f3d6df9f1964e7f942dd242cc9c21b76fa273d Mon Sep 17 00:00:00 2001 From: "Piotr Wilkin (ilintar)" Date: Sat, 4 Jul 2026 10:31:18 +0200 Subject: [PATCH] feat: support for cross-device row split (#1735) --- docs/backend.md | 30 +++++++ examples/common/common.cpp | 9 ++ examples/common/common.h | 1 + include/stable-diffusion.h | 1 + src/core/ggml_extend_backend.cpp | 65 +++++++++++++- src/core/ggml_extend_backend.h | 11 +++ src/core/layer_split_partition.cpp | 6 +- src/core/layer_split_partition.h | 1 + src/model_manager.cpp | 74 +++++++++++++--- src/model_manager.h | 10 ++- src/stable-diffusion.cpp | 134 ++++++++++++++++++++++++++++- src/upscaler.cpp | 1 + 12 files changed, 322 insertions(+), 21 deletions(-) diff --git a/docs/backend.md b/docs/backend.md index 805143a3..7488ce39 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -85,6 +85,36 @@ are disabled for it. Use `--list-devices` to see the device names available on the system. +### Row split (`--split-mode row`) + +`--split-mode` selects how a multi-device module distributes its weights: +`layer` (the default, described above) or `row`. It accepts a single mode or +per-module assignments: + +```shell +sd-cli -m model.safetensors -p "a cat" --backend "diffusion=cuda0&cuda1" --split-mode row +sd-cli -m model.safetensors -p "a cat" --backend "diffusion=cuda0&cuda1,te=cuda0&cuda1" --split-mode diffusion=row,te=layer +``` + +In row mode the module keeps executing on its main (first listed) device, but +its transformer-block matmul weights are allocated in the backend's row-split +buffer type, which slices each weight's rows across the listed devices in +proportion to free memory and runs those matmuls on all devices in parallel. +Compared to a layer split this uses all GPUs within every layer (instead of +sequentially device by device) at the cost of a cross-device reduction per +matmul - usually the faster option when the devices have fast interconnect. + +Row split requires backend support for split buffers and is currently +available on CUDA only; on other backends (or when the listed devices belong +to different backend registries) the module falls back to a layer split. +Embeddings, normalization weights, biases and other non-block tensors stay in +regular buffers on the main device. + +Direct ("immediately") LoRA application cannot patch row-split tensors; with +`--split-mode row` the automatic LoRA mode selects runtime application, and an +explicit `--lora-apply-mode immediately` skips the split tensors with a +warning. + ## Modules | Module | Purpose | Accepted names | diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 7256e25b..bc80a06b 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -468,6 +468,13 @@ ArgOptions SDContextParams::get_options() { "parameter backend assignment, e.g. disk, cpu, or diffusion=disk,clip=cpu", (int)',', ¶ms_backend}, + {"", + "--split-mode", + "weight distribution for modules assigned multiple devices (--backend \"diffusion=cuda0&cuda1\"): " + "layer (whole transformer blocks per device, default) or row (matmul rows split across devices, CUDA only). " + "Accepts a single mode or per-module assignments, e.g. row or diffusion=row,te=layer", + (int)',', + &split_mode}, {"", "--rpc-servers", "comma-separated list of RPC servers to connect to for offloading, in the format host:port, e.g. localhost:50052,192.168.1.3:50052", @@ -830,6 +837,7 @@ std::string SDContextParams::to_string() const { << " eager_load: " << (eager_load ? "true" : "false") << ",\n" << " backend: \"" << backend << "\",\n" << " params_backend: \"" << params_backend << "\",\n" + << " split_mode: \"" << split_mode << "\",\n" << " enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n" << " control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n" << " clip_on_cpu: " << (clip_on_cpu ? "true" : "false") << ",\n" @@ -910,6 +918,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.eager_load = eager_load; sd_ctx_params.backend = effective_backend.c_str(); sd_ctx_params.params_backend = effective_params_backend.c_str(); + sd_ctx_params.split_mode = split_mode.c_str(); sd_ctx_params.rpc_servers = rpc_servers.c_str(); return sd_ctx_params; } diff --git a/examples/common/common.h b/examples/common/common.h index 941fa331..daa15e72 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -151,6 +151,7 @@ struct SDContextParams { bool eager_load = false; std::string backend; std::string params_backend; + std::string split_mode; std::string rpc_servers; std::string effective_backend; std::string effective_params_backend; diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 1acb0a22..17abd3c6 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -227,6 +227,7 @@ typedef struct { bool eager_load; // Load all params into the params backend at model-load time instead of lazily on first use const char* backend; const char* params_backend; + const char* split_mode; // weight distribution for multi-device modules: layer (default) or row, or per-module assignments e.g. "diffusion=row" const char* rpc_servers; } sd_ctx_params_t; diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index 032ae52a..c66bb63f 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -665,8 +665,9 @@ SDBackendManager::~SDBackendManager() { void SDBackendManager::reset() { backends_.clear(); - runtime_assignment_ = {}; - params_assignment_ = {}; + runtime_assignment_ = {}; + params_assignment_ = {}; + split_mode_assignment_ = {}; } static std::vector split_device_list(const std::string& value) { @@ -758,6 +759,7 @@ bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule modu bool SDBackendManager::init(const char* backend_spec, const char* params_backend_spec, + const char* split_mode_spec, std::string* error) { reset(); @@ -767,10 +769,51 @@ bool SDBackendManager::init(const char* backend_spec, if (!sd_parse_backend_assignment(SAFE_STR(params_backend_spec), ¶ms_assignment_, error)) { return false; } + if (!sd_parse_backend_assignment(SAFE_STR(split_mode_spec), &split_mode_assignment_, error)) { + return false; + } return validate(error); } +SDSplitMode SDBackendManager::split_mode(SDBackendModule module) const { + return lower_copy(trim_copy(split_mode_assignment_.get(module))) == "row" ? SDSplitMode::ROW + : SDSplitMode::LAYER; +} + +ggml_backend_buffer_type_t SDBackendManager::split_buffer_type(ggml_backend_t backend, + const std::vector& tensor_split) { + if (backend == nullptr) { + return nullptr; + } + ggml_backend_dev_t dev = ggml_backend_get_device(backend); + if (dev == nullptr) { + return nullptr; + } + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); + if (reg == nullptr) { + return nullptr; + } + auto fn = (ggml_backend_split_buffer_type_t)ggml_backend_reg_get_proc_address(reg, "ggml_backend_split_buffer_type"); + if (fn == nullptr) { + return nullptr; + } + int main_device = -1; + const size_t dev_count = ggml_backend_reg_dev_count(reg); + for (size_t i = 0; i < dev_count; ++i) { + if (ggml_backend_reg_dev_get(reg, i) == dev) { + main_device = (int)i; + break; + } + } + if (main_device < 0) { + return nullptr; + } + std::vector padded_split(std::max(tensor_split.size(), 64), 0.0f); + std::copy(tensor_split.begin(), tensor_split.end(), padded_split.begin()); + return fn(main_device, padded_split.data()); +} + bool SDBackendManager::validate(std::string* error) const { auto validate_single_runtime_name = [&](const std::string& name) -> bool { if (is_default_backend_token(name)) { @@ -826,9 +869,20 @@ bool SDBackendManager::validate(std::string* error) const { } return validate_single_runtime_name(name); }; + auto validate_split_mode_name = [&](const std::string& name) -> bool { + const std::string lower = lower_copy(trim_copy(name)); + if (lower.empty() || lower == "layer" || lower == "row") { + return true; + } + if (error != nullptr) { + *error = "invalid split mode '" + name + "' (expected layer or row)"; + } + return false; + }; if (!validate_runtime_name(runtime_assignment_.default_name) || - !validate_params_name(params_assignment_.default_name)) { + !validate_params_name(params_assignment_.default_name) || + !validate_split_mode_name(split_mode_assignment_.default_name)) { return false; } for (const auto& kv : runtime_assignment_.module_names) { @@ -841,6 +895,11 @@ bool SDBackendManager::validate(std::string* error) const { return false; } } + for (const auto& kv : split_mode_assignment_.module_names) { + if (!validate_split_mode_name(kv.second)) { + return false; + } + } return true; } diff --git a/src/core/ggml_extend_backend.h b/src/core/ggml_extend_backend.h index 92cc6b69..1f3bc8b3 100644 --- a/src/core/ggml_extend_backend.h +++ b/src/core/ggml_extend_backend.h @@ -38,10 +38,16 @@ struct SDBackendHandleDeleter { using SDBackendHandle = std::unique_ptr; +enum class SDSplitMode { + LAYER, + ROW, +}; + class SDBackendManager { private: SDBackendAssignment runtime_assignment_; SDBackendAssignment params_assignment_; + SDBackendAssignment split_mode_assignment_; std::unordered_map backends_; public: @@ -53,6 +59,7 @@ public: bool init(const char* backend_spec, const char* params_backend_spec, + const char* split_mode_spec, std::string* error); void reset(); @@ -61,6 +68,10 @@ public: std::vector runtime_backends(SDBackendModule module); + SDSplitMode split_mode(SDBackendModule module) const; + ggml_backend_buffer_type_t split_buffer_type(ggml_backend_t backend, + const std::vector& tensor_split); + bool runtime_backend_is_cpu(SDBackendModule module); bool params_backend_is_cpu(SDBackendModule module); bool params_backend_is_disk(SDBackendModule module) const; diff --git a/src/core/layer_split_partition.cpp b/src/core/layer_split_partition.cpp index 322b5a31..8e33eb98 100644 --- a/src/core/layer_split_partition.cpp +++ b/src/core/layer_split_partition.cpp @@ -9,7 +9,7 @@ namespace sd { - static int tensor_block_index(const std::string& name) { + int layer_split_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) { @@ -92,7 +92,7 @@ namespace sd { 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; + 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; @@ -160,7 +160,7 @@ namespace sd { for (const auto& kv : tensors) { size_t target = 0; - int idx = split_tensors.count(kv.first) != 0 ? tensor_block_index(kv.first) : -1; + 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++; diff --git a/src/core/layer_split_partition.h b/src/core/layer_split_partition.h index 61d8167c..5c207938 100644 --- a/src/core/layer_split_partition.h +++ b/src/core/layer_split_partition.h @@ -11,6 +11,7 @@ namespace sd { 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, diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 7095ec6a..c5bddcc9 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -100,12 +100,41 @@ size_t estimate_tensors_size(const std::map& tensors) return size; } +void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft) { + if (compute_backend == nullptr) { + return; + } + if (split_buft == nullptr) { + split_buffer_types_.erase(compute_backend); + return; + } + split_buffer_types_[compute_backend] = split_buft; +} + +bool ModelManager::tensor_shape_supports_split_buffer(const ggml_tensor* tensor) { + return tensor != nullptr && + tensor->view_src == nullptr && + ggml_is_contiguous(tensor) && + ggml_n_dims(tensor) == 2 && + tensor->ne[0] >= 256 && + tensor->ne[1] >= 256; +} + +ggml_backend_buffer_type_t ModelManager::split_buffer_type_for(const TensorState& state) const { + if (!state.allow_split_buffer || !tensor_shape_supports_split_buffer(state.tensor)) { + return nullptr; + } + auto it = split_buffer_types_.find(state.compute_backend); + return it != split_buffer_types_.end() ? it->second : nullptr; +} + bool ModelManager::register_param_tensors(const std::string& desc, std::map tensors, ResidencyMode residency_mode, ggml_backend_t compute_backend, ggml_backend_t params_backend, - size_t* registered_tensor_size) { + size_t* registered_tensor_size, + bool allow_split_buffer) { if (desc.empty()) { LOG_ERROR("model manager tensor desc is empty"); return false; @@ -129,13 +158,14 @@ 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; + 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; new_states.push_back(std::move(state)); } @@ -237,7 +267,7 @@ bool ModelManager::load_tensors_to_params_backend(const std::vector& states) { - std::map> states_by_compute_backend; + std::map, std::vector> states_by_staging_target; for (TensorState* state : states) { if (state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name)) { continue; @@ -257,11 +287,16 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vectorname.c_str()); return false; } - states_by_compute_backend[state->compute_backend].push_back(state); + ggml_backend_buffer_type_t staging_buft = split_buffer_type_for(*state); + if (staging_buft == nullptr) { + staging_buft = ggml_backend_get_default_buffer_type(state->compute_backend); + } + states_by_staging_target[{state->compute_backend, staging_buft}].push_back(state); } - for (const auto& pair : states_by_compute_backend) { - ggml_backend_t compute_backend = pair.first; + for (const auto& pair : states_by_staging_target) { + ggml_backend_t compute_backend = pair.first.first; + ggml_backend_buffer_type_t staging_buft = pair.first.second; const std::vector& states = pair.second; if (states.empty()) { continue; @@ -285,7 +320,7 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector& states LOG_ERROR("model manager compute backend is null for lora target tensor '%s'", state->name.c_str()); return false; } + if (state->tensor->buffer != nullptr && + ggml_backend_buffer_get_type(state->tensor->buffer) == split_buffer_type_for(*state)) { + if (!warned_split_lora_skip_) { + LOG_WARN( + "model manager skipping direct lora application to row-split tensors " + "(use --lora-apply-mode at_runtime with row split)"); + warned_split_lora_skip_ = true; + } + state->applied_lora_epoch = current_lora_epoch_; + continue; + } if (state->tensor->data == nullptr) { LOG_ERROR("model manager lora target tensor '%s' is not prepared", state->name.c_str()); return false; @@ -694,6 +740,8 @@ ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorStat if (compute_dev != nullptr) { params_buft = ggml_backend_dev_host_buffer_type(compute_dev); } + } else if (state.params_backend == state.compute_backend) { + params_buft = split_buffer_type_for(state); } if (params_buft == nullptr) { params_buft = ggml_backend_get_default_buffer_type(state.params_backend); diff --git a/src/model_manager.h b/src/model_manager.h index 9225e3ea..8ec7a7a1 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -36,6 +36,7 @@ private: 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; int active_prepare_count = 0; @@ -63,6 +64,8 @@ private: std::map tensor_states_by_name_; std::vector> params_storage_blocks_; std::vector> compute_staging_blocks_; + std::map split_buffer_types_; + bool warned_split_lora_skip_ = false; std::set common_ignore_tensors_; std::vector loras_; SDVersion lora_version_ = VERSION_COUNT; @@ -91,6 +94,7 @@ private: bool stage_tensors_to_compute_backend(const std::vector& states); ggml_backend_buffer_type_t params_buffer_type_for(const TensorState& state) const; + ggml_backend_buffer_type_t split_buffer_type_for(const TensorState& state) const; void release_compute_staging_blocks(bool force = false, const std::unordered_set* target_states = nullptr); void release_params_storage_blocks(bool force = false, @@ -114,6 +118,9 @@ public: void set_writable_mmap(bool writable_mmap) { writable_mmap_ = writable_mmap; } void set_common_ignore_tensors(std::set ignore_tensors); void set_loras(std::vector loras, SDVersion version); + void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft); + + static bool tensor_shape_supports_split_buffer(const ggml_tensor* tensor); std::set tensor_names() const; @@ -122,7 +129,8 @@ public: ResidencyMode residency_mode, ggml_backend_t compute_backend, ggml_backend_t params_backend, - size_t* registered_tensor_size = nullptr); + size_t* registered_tensor_size = nullptr, + bool allow_split_buffer = false); template bool register_runner_params(const std::string& desc, diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index e239fe1a..bdc2e925 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -218,6 +218,7 @@ public: bool eager_load = false; std::string backend_spec; std::string params_backend_spec; + std::string split_mode_spec; bool is_using_v_parameterization = false; bool is_using_edm_v_parameterization = false; @@ -292,6 +293,15 @@ public: if (module_backends.size() > 1) { if constexpr (has_set_runtime_backends::value) { if (module == SDBackendModule::DIFFUSION || module == SDBackendModule::TE) { + if (backend_manager.split_mode(module) == SDSplitMode::ROW) { + return register_row_split_runner_params(desc, + model, + module, + module_backends, + std::move(group_tensors), + residency_mode, + params_mem_size); + } return register_layer_split_runner_params(desc, model, module, @@ -313,6 +323,108 @@ public: params_mem_size); } + template + bool register_row_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) { + ggml_backend_t main_backend = module_backends[0]; + + auto fall_back_to_layer_split = [&](const char* reason) { + LOG_WARN("%s: row split unavailable (%s); falling back to layer split", desc.c_str(), reason); + return register_layer_split_runner_params(desc, + model, + module, + module_backends, + std::move(group_tensors), + residency_mode, + params_mem_size); + }; + + ggml_backend_dev_t main_dev = ggml_backend_get_device(main_backend); + ggml_backend_reg_t reg = main_dev != nullptr ? ggml_backend_dev_backend_reg(main_dev) : nullptr; + if (reg == nullptr) { + return fall_back_to_layer_split("no backend registry"); + } + const size_t reg_dev_count = ggml_backend_reg_dev_count(reg); + std::vector tensor_split(reg_dev_count, 0.0f); + constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024; + for (ggml_backend_t backend : module_backends) { + ggml_backend_dev_t dev = ggml_backend_get_device(backend); + int reg_index = -1; + for (size_t i = 0; i < reg_dev_count; i++) { + if (ggml_backend_reg_dev_get(reg, i) == dev) { + reg_index = (int)i; + break; + } + } + if (reg_index < 0) { + return fall_back_to_layer_split("devices span different backend registries"); + } + size_t free_bytes = 0, total_bytes = 0; + ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); + int64_t usable_bytes = std::max((int64_t)free_bytes - compute_headroom_bytes, + (int64_t)free_bytes / 8); + tensor_split[reg_index] = usable_bytes > 0 ? (float)((double)usable_bytes / (1024.0 * 1024.0)) : 1.0f; + } + + ggml_backend_buffer_type_t split_buft = backend_manager.split_buffer_type(main_backend, tensor_split); + if (split_buft == nullptr) { + return fall_back_to_layer_split("backend has no split buffer type"); + } + model_manager->set_split_buffer_type(main_backend, split_buft); + + std::map split_tensors; + if constexpr (std::is_base_of_v) { + model->get_layer_split_param_tensors(split_tensors); + } else { + split_tensors = group_tensors; + } + + std::map row_split_map; + std::map regular_map; + size_t row_split_bytes = 0; + for (const auto& kv : group_tensors) { + if (split_tensors.count(kv.first) != 0 && + sd::layer_split_tensor_block_index(kv.first) >= 0 && + ModelManager::tensor_shape_supports_split_buffer(kv.second)) { + row_split_map[kv.first] = kv.second; + row_split_bytes += ggml_nbytes(kv.second); + } else { + regular_map[kv.first] = kv.second; + } + } + if (row_split_map.empty()) { + return fall_back_to_layer_split("no row-splittable transformer block weights found"); + } + + LOG_INFO("%s row split: %zu tensors (%.1f MB) split across %zu devices (main %s)", + desc.c_str(), + row_split_map.size(), + row_split_bytes / (1024.f * 1024.f), + module_backends.size(), + sd::layer_split_backend_device_display_name(main_backend).c_str()); + + if (!model_manager->register_param_tensors(desc, + std::move(row_split_map), + residency_mode, + main_backend, + params_backend_for(module), + params_mem_size, + /*allow_split_buffer=*/true)) { + return false; + } + return model_manager->register_param_tensors(desc, + std::move(regular_map), + residency_mode, + main_backend, + 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 @@ -397,6 +509,7 @@ public: std::string error; if (!backend_manager.init(backend_spec.c_str(), params_backend_spec.c_str(), + split_mode_spec.c_str(), &error)) { LOG_ERROR("backend config failed: %s", error.c_str()); return false; @@ -404,6 +517,16 @@ public: return ensure_backend_pair(SDBackendModule::DIFFUSION); } + bool row_split_active() { + for (SDBackendModule module : {SDBackendModule::DIFFUSION, SDBackendModule::TE}) { + if (backend_manager.split_mode(module) == SDSplitMode::ROW && + 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(); @@ -462,6 +585,7 @@ public: eager_load = sd_ctx_params->eager_load; backend_spec = SAFE_STR(sd_ctx_params->backend); params_backend_spec = SAFE_STR(sd_ctx_params->params_backend); + split_mode_spec = SAFE_STR(sd_ctx_params->split_mode); max_vram_assignment.reset(0.f); { std::string error; @@ -690,12 +814,17 @@ public: // Avoid full-model LoRA merge buffers on constrained setups. const bool params_offloaded = params_backend_for(SDBackendModule::DIFFUSION) != backend_for(SDBackendModule::DIFFUSION); const bool streaming_constrained = stream_layers || params_offloaded; - if (have_quantized_weight || streaming_constrained) { + if (have_quantized_weight || streaming_constrained || row_split_active()) { apply_lora_immediately = false; } else { apply_lora_immediately = true; } } else if (sd_ctx_params->lora_apply_mode == LORA_APPLY_IMMEDIATELY) { + if (row_split_active()) { + LOG_WARN( + "row-split tensors do not support the immediately LoRA apply mode; " + "LoRAs will not be applied to them (use --lora-apply-mode at_runtime)"); + } apply_lora_immediately = true; } else { apply_lora_immediately = false; @@ -2916,6 +3045,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO; sd_ctx_params->backend = nullptr; sd_ctx_params->params_backend = nullptr; + sd_ctx_params->split_mode = nullptr; sd_ctx_params->rpc_servers = nullptr; sd_ctx_params->pulid_weights_path = nullptr; } @@ -2955,6 +3085,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "eager_load: %s\n" "backend: %s\n" "params_backend: %s\n" + "split_mode: %s\n" "flash_attn: %s\n" "diffusion_flash_attn: %s\n" "circular_x: %s\n" @@ -2991,6 +3122,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { BOOL_STR(sd_ctx_params->eager_load), SAFE_STR(sd_ctx_params->backend), SAFE_STR(sd_ctx_params->params_backend), + SAFE_STR(sd_ctx_params->split_mode), BOOL_STR(sd_ctx_params->flash_attn), BOOL_STR(sd_ctx_params->diffusion_flash_attn), BOOL_STR(sd_ctx_params->circular_x), diff --git a/src/upscaler.cpp b/src/upscaler.cpp index 88a8a633..dbb99af3 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -46,6 +46,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path, std::string error; if (!backend_manager.init(backend_spec.c_str(), params_backend_spec.c_str(), + /*split_mode_spec=*/nullptr, &error)) { LOG_ERROR("upscaler backend config failed: %s", error.c_str()); return false;