diff --git a/docs/backend.md b/docs/backend.md index 7488ce397..353255371 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -115,6 +115,33 @@ Direct ("immediately") LoRA application cannot patch row-split tensors; with explicit `--lora-apply-mode immediately` skips the split tensors with a warning. +## Automatic placement (`--auto-fit`) + +`--auto-fit` derives the `diffusion` / `te` / `vae` placements from the model +metadata and the per-device memory budgets, then feeds them into the same +backend assignment mechanism described above (the chosen specs are printed). +`--backend` and `--params-backend` are ignored while auto-fit is enabled. + +```shell +sd-cli -m model.safetensors -p "a cat" --auto-fit +sd-cli -m model.safetensors -p "a cat" --auto-fit --max-vram cuda0=8,cuda1=14 +sd-cli -m model.safetensors -p "a cat" --auto-fit --split-mode row +``` + +Budgets reuse `--max-vram`: a positive per-device value caps what auto-fit +plans with on that device, a negative value means "free memory minus that many +GiB", and with no budget set each device's free memory minus a 512 MiB margin +is used. (The same values still drive graph-cut segmented execution for +modules that end up on a single device.) + +When everything fits resident, components are simply spread across the +available GPUs. When it does not, auto-fit switches to time-share mode: the +heavy components get `disk` params residency (loaded for their phase, freed +after), and a component too large for any single device is split across all +GPUs with the layer/row split mechanism (`--split-mode` selects which, layer +by default). Components that fit nowhere fall back to the CPU. If a VAE decode +still runs out of memory, tiling is enabled and the decode retried once. + ## Modules | Module | Purpose | Accepted names | diff --git a/examples/common/common.cpp b/examples/common/common.cpp index bc80a06bb..2a785c0c0 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -508,6 +508,12 @@ ArgOptions SDContextParams::get_options() { "--eager-load", "load all params into the params backend at model-load time instead of lazily on first use (defaults to false)", true, &eager_load}, + {"", + "--auto-fit", + "pick the diffusion/te/vae device placements automatically from the model size and the per-device " + "memory budgets (--max-vram; defaults to free memory minus a small margin). Overrides --backend and " + "--params-backend; may split modules across GPUs (--split-mode still selects layer or row)", + true, &auto_fit}, {"", "--force-sdxl-vae-conv-scale", "force use of conv scale on sdxl vae", @@ -838,6 +844,7 @@ std::string SDContextParams::to_string() const { << " backend: \"" << backend << "\",\n" << " params_backend: \"" << params_backend << "\",\n" << " split_mode: \"" << split_mode << "\",\n" + << " auto_fit: " << (auto_fit ? "true" : "false") << ",\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" @@ -919,6 +926,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { 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.auto_fit = auto_fit; 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 daa15e72a..654871359 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -152,6 +152,7 @@ struct SDContextParams { std::string backend; std::string params_backend; std::string split_mode; + bool auto_fit = false; 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 17abd3c62..c905698e5 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -228,6 +228,7 @@ typedef struct { 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" + bool auto_fit; const char* rpc_servers; } sd_ctx_params_t; diff --git a/src/core/backend_fit.cpp b/src/core/backend_fit.cpp new file mode 100644 index 000000000..460194740 --- /dev/null +++ b/src/core/backend_fit.cpp @@ -0,0 +1,390 @@ +#include "backend_fit.h" + +#include +#include +#include +#include +#include + +#include "core/ggml_extend_backend.h" +#include "core/util.h" +#include "ggml-backend.h" + +namespace sd::backend_fit { + namespace { + + constexpr int64_t MiB = 1024ll * 1024; + + enum class ComponentKind { + DIT = 0, + VAE = 1, + CONDITIONER = 2, + }; + + struct Component { + ComponentKind kind; + const char* name; + int64_t params_bytes = 0; + int64_t reserve_bytes = 0; + bool splittable = false; + }; + + struct Device { + ggml_backend_dev_t dev = nullptr; + std::string name; + std::string description; + int64_t free_bytes = 0; + int64_t total_bytes = 0; + int64_t budget_bytes = 0; + }; + + struct Decision { + ComponentKind kind; + bool on_cpu = false; + std::vector device_idxs; + }; + + struct Plan { + bool valid = false; + bool time_share = false; + std::vector decisions; + }; + + bool classify_tensor(const std::string& name, ComponentKind& out) { + auto contains = [&](const char* s) { return name.find(s) != std::string::npos; }; + + if (contains("model.diffusion_model.") || contains("unet.")) { + out = ComponentKind::DIT; + return true; + } + if (contains("first_stage_model.") || + name.rfind("vae.", 0) == 0 || + name.rfind("tae.", 0) == 0) { + out = ComponentKind::VAE; + return true; + } + if (contains("text_encoders") || + contains("cond_stage_model") || + contains("te.text_model.") || + contains("conditioner") || + name.rfind("text_encoder.", 0) == 0 || + name.rfind("text_embedding_projection.", 0) == 0 || + contains(".aggregate_embed.")) { + out = ComponentKind::CONDITIONER; + return true; + } + return false; + } + + std::vector estimate_components(ModelLoader& loader, ggml_type override_wtype) { + const auto& storage = loader.get_tensor_storage_map(); + + int64_t bytes[3] = {0, 0, 0}; + for (const auto& [name, ts_const] : storage) { + TensorStorage ts = ts_const; + if (is_unused_tensor(ts.name)) { + continue; + } + ComponentKind kind; + if (!classify_tensor(ts.name, kind)) { + continue; + } + if (override_wtype != GGML_TYPE_COUNT && + loader.tensor_should_be_converted(ts, override_wtype)) { + ts.type = override_wtype; + } else if (ts.expected_type != GGML_TYPE_COUNT && ts.expected_type != ts.type) { + ts.type = ts.expected_type; + } + bytes[int(kind)] += (int64_t)ts.nbytes() + 64; + } + + std::vector out; + out.push_back({ComponentKind::DIT, "DiT", bytes[int(ComponentKind::DIT)], 2048 * MiB, true}); + out.push_back({ComponentKind::VAE, "VAE", bytes[int(ComponentKind::VAE)], 1024 * MiB, false}); + out.push_back({ComponentKind::CONDITIONER, "Conditioner", bytes[int(ComponentKind::CONDITIONER)], 2048 * MiB, true}); + return out; + } + + std::vector enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) { + std::vector out; + for (size_t i = 0; i < ggml_backend_dev_count(); i++) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) { + continue; + } + Device d; + d.dev = dev; + d.name = ggml_backend_dev_name(dev); + d.description = ggml_backend_dev_description(dev); + size_t free_bytes = 0, total_bytes = 0; + ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); + d.free_bytes = (int64_t)free_bytes; + d.total_bytes = (int64_t)total_bytes; + + std::string budget_key = d.name; + std::transform(budget_key.begin(), budget_key.end(), budget_key.begin(), + [](unsigned char c) { return (char)std::tolower(c); }); + float gib = budgets.default_gib; + auto it = budgets.backend_gib.find(budget_key); + if (it != budgets.backend_gib.end()) { + gib = it->second; + } + if (gib > 0.f) { + d.budget_bytes = std::min((int64_t)(gib * 1024.0 * 1024.0 * 1024.0), d.free_bytes); + } else if (gib < 0.f) { + d.budget_bytes = d.free_bytes + (int64_t)(gib * 1024.0 * 1024.0 * 1024.0); + } else { + d.budget_bytes = d.free_bytes - 512 * MiB; + } + d.budget_bytes = std::max(d.budget_bytes, 0); + out.push_back(d); + } + return out; + } + + Plan compute_plan(const std::vector& components, const std::vector& devices) { + Plan plan; + if (devices.empty()) { + return plan; + } + + std::vector order(components.size()); + for (size_t i = 0; i < order.size(); i++) { + order[i] = i; + } + std::sort(order.begin(), order.end(), [&](size_t a, size_t b) { + return components[a].params_bytes > components[b].params_bytes; + }); + + { + std::vector params_sum(devices.size(), 0); + std::vector max_reserve(devices.size(), 0); + std::vector decisions(components.size()); + bool ok = true; + for (size_t ci : order) { + const Component& comp = components[ci]; + decisions[ci].kind = comp.kind; + if (comp.params_bytes == 0) { + continue; + } + int best = -1; + for (size_t di = 0; di < devices.size(); di++) { + int64_t need = params_sum[di] + comp.params_bytes + std::max(max_reserve[di], comp.reserve_bytes); + if (need <= devices[di].budget_bytes && + (best < 0 || devices[di].budget_bytes - params_sum[di] > devices[best].budget_bytes - params_sum[best])) { + best = (int)di; + } + } + if (best < 0) { + ok = false; + break; + } + params_sum[best] += comp.params_bytes; + max_reserve[best] = std::max(max_reserve[best], comp.reserve_bytes); + decisions[ci].device_idxs.push_back((size_t)best); + } + if (ok) { + plan.valid = true; + plan.time_share = false; + plan.decisions = std::move(decisions); + return plan; + } + } + + plan.decisions.assign(components.size(), {}); + for (size_t ci : order) { + const Component& comp = components[ci]; + Decision& decision = plan.decisions[ci]; + decision.kind = comp.kind; + if (comp.params_bytes == 0) { + continue; + } + int best = -1; + for (size_t di = 0; di < devices.size(); di++) { + if (comp.params_bytes + comp.reserve_bytes <= devices[di].budget_bytes && + (best < 0 || devices[di].budget_bytes > devices[best].budget_bytes)) { + best = (int)di; + } + } + if (best >= 0) { + decision.device_idxs.push_back((size_t)best); + continue; + } + if (comp.splittable && devices.size() > 1) { + int64_t capacity = 0; + for (const Device& d : devices) { + capacity += std::max(d.budget_bytes - comp.reserve_bytes, 0); + } + if (comp.params_bytes <= capacity) { + std::vector idxs(devices.size()); + for (size_t i = 0; i < idxs.size(); i++) { + idxs[i] = i; + } + std::sort(idxs.begin(), idxs.end(), [&](size_t a, size_t b) { + return devices[a].budget_bytes > devices[b].budget_bytes; + }); + decision.device_idxs = std::move(idxs); + continue; + } + } + decision.on_cpu = true; + } + plan.valid = true; + plan.time_share = true; + return plan; + } + + void print_plan(const Plan& plan, + const std::vector& components, + const std::vector& devices) { + LOG_INFO("auto-fit plan%s:", plan.time_share ? " (time-share: params load per phase and free after)" : ""); + LOG_INFO(" devices:"); + for (const Device& d : devices) { + LOG_INFO(" %-12s %-32s free %6lld MiB, budget %6lld MiB", + d.name.c_str(), d.description.c_str(), + (long long)(d.free_bytes / MiB), (long long)(d.budget_bytes / MiB)); + } + LOG_INFO(" components:"); + for (size_t ci = 0; ci < components.size(); ci++) { + const Component& comp = components[ci]; + const Decision& decision = plan.decisions[ci]; + std::string target; + if (comp.params_bytes == 0) { + target = "(not present)"; + } else if (decision.on_cpu) { + target = "CPU"; + } else { + for (size_t k = 0; k < decision.device_idxs.size(); k++) { + if (k > 0) { + target += " & "; + } + target += devices[decision.device_idxs[k]].name; + } + if (decision.device_idxs.size() > 1) { + target += " (split)"; + } + } + LOG_INFO(" %-12s params %6lld MiB, compute reserve %5lld MiB -> %s", + comp.name, + (long long)(comp.params_bytes / MiB), + (long long)(comp.reserve_bytes / MiB), + target.c_str()); + } + } + + void append_assignment(std::string& spec, const char* key, const std::string& value) { + if (!spec.empty()) { + spec += ","; + } + spec += key; + spec += "="; + spec += value; + } + + void append_component_decision(const std::vector& components, + const std::vector& devices, + const Plan& plan, + ComponentKind kind, + const char* module_key, + std::string& runtime_spec, + std::string& params_spec) { + for (size_t ci = 0; ci < components.size(); ci++) { + if (components[ci].kind != kind || components[ci].params_bytes == 0) { + continue; + } + const Decision& decision = plan.decisions[ci]; + if (decision.on_cpu) { + append_assignment(runtime_spec, module_key, "cpu"); + return; + } + if (decision.device_idxs.empty()) { + return; + } + std::string device_list; + for (size_t k = 0; k < decision.device_idxs.size(); k++) { + if (k > 0) { + device_list += "&"; + } + device_list += devices[decision.device_idxs[k]].name; + } + append_assignment(runtime_spec, module_key, device_list); + if (plan.time_share) { + append_assignment(params_spec, module_key, "disk"); + } + return; + } + } + + } // namespace + + bool derive_backend_specs(ModelLoader& loader, + ggml_type override_wtype, + sd::ggml_graph_cut::MaxVramAssignment& budgets, + std::string& runtime_spec, + std::string& params_spec) { + if (!runtime_spec.empty() || !params_spec.empty()) { + LOG_WARN("--auto-fit is enabled; ignoring --backend / --params-backend"); + } + + { + std::string error; + if (!budgets.canonicalize_backend_keys(&error)) { + LOG_ERROR("%s", error.c_str()); + return false; + } + } + + auto components = estimate_components(loader, override_wtype); + auto devices = enumerate_gpu_devices(budgets); + auto plan = compute_plan(components, devices); + if (!plan.valid) { + LOG_WARN("auto-fit: no usable GPU devices; using the default backend"); + runtime_spec.clear(); + params_spec.clear(); + return true; + } + + print_plan(plan, components, devices); + + std::string derived_runtime_spec; + std::string derived_params_spec; + append_component_decision(components, devices, plan, ComponentKind::DIT, "diffusion", derived_runtime_spec, derived_params_spec); + append_component_decision(components, devices, plan, ComponentKind::CONDITIONER, "te", derived_runtime_spec, derived_params_spec); + append_component_decision(components, devices, plan, ComponentKind::VAE, "vae", derived_runtime_spec, derived_params_spec); + + runtime_spec = std::move(derived_runtime_spec); + params_spec = std::move(derived_params_spec); + + LOG_INFO("auto-fit: --backend \"%s\"%s%s%s", + runtime_spec.empty() ? "(default)" : runtime_spec.c_str(), + params_spec.empty() ? "" : " --params-backend \"", + params_spec.c_str(), + params_spec.empty() ? "" : "\""); + return true; + } + + bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) { + if (prefer_temporal_tiling) { + if (tiling_params.temporal_tiling) { + return false; + } + tiling_params.temporal_tiling = true; + } else { + if (tiling_params.enabled) { + return false; + } + tiling_params.enabled = true; + if (tiling_params.tile_size_x <= 0) { + tiling_params.tile_size_x = 256; + } + if (tiling_params.tile_size_y <= 0) { + tiling_params.tile_size_y = 256; + } + } + + LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling", + tiling_params.temporal_tiling ? "temporal" : "spatial"); + return true; + } + +} // namespace sd::backend_fit diff --git a/src/core/backend_fit.h b/src/core/backend_fit.h new file mode 100644 index 000000000..9ef298b3b --- /dev/null +++ b/src/core/backend_fit.h @@ -0,0 +1,23 @@ +#ifndef __SD_BACKEND_FIT_H__ +#define __SD_BACKEND_FIT_H__ + +#include + +#include "core/ggml_graph_cut.h" +#include "model_loader.h" +#include "stable-diffusion.h" + +namespace sd::backend_fit { + + bool derive_backend_specs(ModelLoader& loader, + ggml_type override_wtype, + sd::ggml_graph_cut::MaxVramAssignment& budgets, + std::string& runtime_spec, + std::string& params_spec); + + bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, + bool prefer_temporal_tiling); + +} // namespace sd::backend_fit + +#endif // __SD_BACKEND_FIT_H__ diff --git a/src/model_loader.h b/src/model_loader.h index 4dc700f20..529f3e890 100644 --- a/src/model_loader.h +++ b/src/model_loader.h @@ -27,6 +27,8 @@ struct MmapTensorStore { std::shared_ptr mmbuffer; }; +bool is_unused_tensor(const std::string& name); + class ModelLoader { protected: SDVersion version_ = VERSION_COUNT; diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index bdc2e9253..91c25ae79 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -20,6 +20,7 @@ #include "stable-diffusion.h" #include "conditioning/conditioner.hpp" +#include "core/backend_fit.h" #include "extensions/generation_extension.h" #include "model/adapter/lora.hpp" #include "model/diffusion/anima.hpp" @@ -219,6 +220,7 @@ class StableDiffusionGGML { std::string backend_spec; std::string params_backend_spec; std::string split_mode_spec; + bool auto_fit_enabled = false; bool is_using_v_parameterization = false; bool is_using_edm_v_parameterization = false; @@ -586,6 +588,7 @@ class StableDiffusionGGML { 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); + auto_fit_enabled = sd_ctx_params->auto_fit; max_vram_assignment.reset(0.f); { std::string error; @@ -611,21 +614,6 @@ class StableDiffusionGGML { ggml_log_set(ggml_log_callback_default, nullptr); - if (!init_backend()) { - return false; - } - { - std::string error; - if (!max_vram_assignment.canonicalize_backend_keys(&error)) { - LOG_ERROR("%s", error.c_str()); - return false; - } - } - if (stream_layers && !backend_manager.params_backend_is_cpu(SDBackendModule::DIFFUSION)) { - LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring"); - stream_layers = false; - } - model_manager = std::make_shared(); model_manager->set_n_threads(n_threads); model_manager->set_enable_mmap(enable_mmap); @@ -773,6 +761,31 @@ class StableDiffusionGGML { model_loader.set_wtype_override(wtype, tensor_type_rules); } + if (auto_fit_enabled) { + if (!sd::backend_fit::derive_backend_specs(model_loader, + wtype, + max_vram_assignment, + backend_spec, + params_backend_spec)) { + return false; + } + } + + if (!init_backend()) { + return false; + } + { + std::string error; + if (!max_vram_assignment.canonicalize_backend_keys(&error)) { + LOG_ERROR("%s", error.c_str()); + return false; + } + } + if (stream_layers && !backend_manager.params_backend_is_cpu(SDBackendModule::DIFFUSION)) { + LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring"); + stream_layers = false; + } + std::map wtype_stat = model_loader.get_wtype_stat(); std::map conditioner_wtype_stat = model_loader.get_conditioner_wtype_stat(); std::map diffusion_model_wtype_stat = model_loader.get_diffusion_model_wtype_stat(); @@ -2689,7 +2702,16 @@ class StableDiffusionGGML { } auto latents = first_stage_model->diffusion_to_vae_latents(x); first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling); - return first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); + auto decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); + if (decoded.empty() && auto_fit_enabled) { + bool prefer_temporal_tiling = decode_video && std::dynamic_pointer_cast(first_stage_model) != nullptr; + if (sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { + first_stage_model->free_compute_buffer(); + first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling); + decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); + } + } + return decoded; } sd::Tensor normalize_ltx_video_latents(const sd::Tensor& x) { @@ -3046,6 +3068,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { sd_ctx_params->backend = nullptr; sd_ctx_params->params_backend = nullptr; sd_ctx_params->split_mode = nullptr; + sd_ctx_params->auto_fit = false; sd_ctx_params->rpc_servers = nullptr; sd_ctx_params->pulid_weights_path = nullptr; } @@ -3086,6 +3109,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "backend: %s\n" "params_backend: %s\n" "split_mode: %s\n" + "auto_fit: %s\n" "flash_attn: %s\n" "diffusion_flash_attn: %s\n" "circular_x: %s\n" @@ -3123,6 +3147,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { 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->auto_fit), BOOL_STR(sd_ctx_params->flash_attn), BOOL_STR(sd_ctx_params->diffusion_flash_attn), BOOL_STR(sd_ctx_params->circular_x),