feat: enable single-GPU auto-fit with tiered parameter placement (#1942)

This commit is contained in:
leejet 2026-09-07 00:18:42 +08:00 committed by GitHub
parent dbb611264e
commit 80bac2d5fc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 408 additions and 330 deletions

View File

@ -126,32 +126,70 @@ Direct ("immediately") LoRA application cannot patch row-split tensors; with
explicit `--lora-apply-mode immediately` skips the split tensors with a explicit `--lora-apply-mode immediately` skips the split tensors with a
warning. warning.
## Automatic placement (`--auto-fit`) ## Automatic placement (`--auto-fit on|off`)
`--auto-fit` derives the `diffusion` / `te` / `vae` placements from the model `--auto-fit` requires `on` or `off` and defaults to `on` when omitted.
metadata and the per-device memory budgets, then feeds them into the same Explicit `--backend` or `--params-backend` assignments disable auto-fit,
backend assignment mechanism described above (the chosen specs are printed). regardless of argument order, even with `--auto-fit on`.
`--backend` and `--params-backend` are ignored while auto-fit is enabled.
When enabled, auto-fit uses one GPU for `diffusion` / `te` / `vae` computation. It chooses
the GPU with the largest available memory budget (the first device on a tie),
then derives parameter placements from the model metadata and the remaining
memory budgets. The chosen backend specifications are printed.
```shell ```shell
sd-cli -m model.safetensors -p "a cat" --auto-fit sd-cli -m model.safetensors -p "a cat" --auto-fit on
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 on --max-vram cuda0=8,cuda1=14
sd-cli -m model.safetensors -p "a cat" --auto-fit --split-mode row sd-cli -m model.safetensors -p "a cat" --auto-fit off
``` ```
Budgets reuse `--max-vram`: a positive per-device value caps what auto-fit 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 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 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 is used. These resolved GPU budgets, including the safety margin, also drive
modules that end up on a single device.) the runner's graph-cut capacity checks.
When everything fits resident, components are simply spread across the Components are considered in `diffusion`, `te`, `vae` order so that repeatedly
available GPUs. When it does not, auto-fit switches to time-share mode: the used diffusion weights have priority. Each component's weights use the first
heavy components get `disk` params residency (loaded for their phase, freed storage location with enough remaining budget:
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 1. The main GPU, leaving estimated space for computation and weight staging.
by default). Components that fit nowhere fall back to the CPU. If a VAE decode 2. CPU RAM, reserving the larger of 2 GiB or 10% of available RAM for other work.
still runs out of memory, tiling is enabled and the decode retried once. 3. Another GPU, choosing the one with the largest remaining budget that fits.
4. Disk, reloading weights on demand.
GPU cache space follows the same component priority. Before a lower-priority
component can become permanently resident, the planner leaves room for the full
weights and estimated compute space of higher-priority offloaded components.
If offloaded diffusion already needs the entire main GPU budget, TE and VAE also
use offloaded parameters. Their GPU copies can then be released after their
phases, leaving more room to reuse diffusion weights across sampling steps.
CPU parameter residency allows GPU weight caching; it does not force every
weight to be copied again at every step.
RAM and GPU budgets are shared across components. Each component uses a single
parameter backend; several other GPUs' capacities are not combined to store
one component. If available RAM cannot be queried, RAM residency is skipped.
Other GPUs store weights only: weights are copied to the main GPU for execution.
Auto-fit does not select multi-GPU layer/row computation, so `--split-mode` does
not change its placements. Use explicit backend assignments for multi-GPU
computation.
For example, a diffusion model whose full weights exceed the main GPU's budget
can use `--backend diffusion=cuda0 --params-backend diffusion=cpu` when RAM is
sufficient. Automatic graph segmentation can then load the required weights
for each segment and reclaim idle GPU copies. `--disable-segmented-compute`
still disables segmentation.
Initial compute reserves are estimates (2 GiB for diffusion and text encoders,
1 GiB for VAE); higher-priority placements also leave staging space for the
largest weight tensor of each lower-priority offloaded component. Actual segment
weights, compute buffers and caches must
still fit the runner's capacity checks. Offloading weights does not guarantee
that every resolution or frame count will fit, and auto-fit does not change a
component to CPU computation solely because its full weights exceed VRAM.
If a VAE decode fails, auto-fit retries with spatial tiling; supported video
decoders try temporal tiling first and can then add spatial tiling.
## Modules ## Modules
@ -203,7 +241,7 @@ sd-cli -m model.safetensors -p "a cat" --backend cuda0 --params-backend disk
This runs all modules on `cuda0`, reloads parameters from the model file as needed, and releases those parameter buffers after use. This runs all modules on `cuda0`, reloads parameters from the model file as needed, and releases those parameter buffers after use.
`disk` is never selected implicitly. If `--params-backend` is not set, parameters use the runtime backend. Outside `--auto-fit`, `disk` is never selected implicitly. If `--params-backend` is not set, parameters use the runtime backend.
Per-module assignments can be mixed: Per-module assignments can be mixed:
@ -252,4 +290,7 @@ The example CLI/server still accepts these older CPU placement flags as compatib
Because this default is inserted first, later explicit `--params-backend` entries can still override it, for example `--offload-to-cpu --params-backend te=disk` keeps non-TE parameters on CPU and reloads TE parameters from disk. Because this default is inserted first, later explicit `--params-backend` entries can still override it, for example `--offload-to-cpu --params-backend te=disk` keeps non-TE parameters on CPU and reloads TE parameters from disk.
Library callers should set `backend` and `params_backend` directly. The old CPU/offload fields are no longer part of the C API. Explicit `--backend` and `--params-backend` assignments are preferred for new commands. Library callers should set `backend` and `params_backend` directly. `sd_ctx_params_init()`
enables `auto_fit` by default; nonempty `backend` or `params_backend` assignments disable it.
The old CPU/offload fields are no longer part of the C API. Explicit `--backend` and
`--params-backend` assignments are preferred for new commands.

View File

@ -548,12 +548,6 @@ ArgOptions SDContextParams::get_options() {
"--eager-load", "--eager-load",
"load all params into the params backend at model-load time instead of lazily on first use (defaults to false)", "load all params into the params backend at model-load time instead of lazily on first use (defaults to false)",
true, &eager_load}, 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-sdxl-vae-conv-scale",
"force use of conv scale on sdxl vae", "force use of conv scale on sdxl vae",
@ -596,6 +590,23 @@ ArgOptions SDContextParams::get_options() {
true, &vae_conv_direct}, true, &vae_conv_direct},
}; };
auto on_auto_fit_arg = [&](int argc, const char** argv, int index) {
if (++index >= argc) {
LOG_ERROR("--auto-fit requires 'on' or 'off'");
return -1;
}
const std::string arg = argv[index];
if (arg == "on") {
auto_fit = true;
} else if (arg == "off") {
auto_fit = false;
} else {
LOG_ERROR("invalid --auto-fit value '%s'; expected 'on' or 'off'", argv[index]);
return -1;
}
return 1;
};
auto on_type_arg = [&](int argc, const char** argv, int index) { auto on_type_arg = [&](int argc, const char** argv, int index) {
if (++index >= argc) { if (++index >= argc) {
return -1; return -1;
@ -667,6 +678,12 @@ ArgOptions SDContextParams::get_options() {
}; };
options.manual_options = { options.manual_options = {
{"",
"--auto-fit",
"on|off (default: on). Use one GPU for diffusion/te/vae computation and place weights on that GPU, "
"RAM, another GPU, or disk in that order, according to available memory (--max-vram limits GPU budgets). "
"Disabled by explicit --backend or --params-backend; uses automatic graph segmentation when needed",
on_auto_fit_arg},
{"", {"",
"--type", "--type",
"weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, q4_K). " "weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, q4_K). "

View File

@ -158,7 +158,7 @@ struct SDContextParams {
std::string params_backend; std::string params_backend;
std::string split_mode; std::string split_mode;
std::string model_args; std::string model_args;
bool auto_fit = false; bool auto_fit = true;
std::string rpc_servers; std::string rpc_servers;
std::string effective_backend; std::string effective_backend;
std::string effective_params_backend; std::string effective_params_backend;

View File

@ -2,364 +2,384 @@
#include <algorithm> #include <algorithm>
#include <cctype> #include <cctype>
#include <cstddef>
#include <cstdint> #include <cstdint>
#include <fstream>
#include <utility> #include <utility>
#include <vector> #include <vector>
#if defined(_WIN32)
#ifndef NOMINMAX
#define NOMINMAX
#endif
#include <windows.h>
#elif defined(__APPLE__)
#include <mach/mach.h>
#endif
#include "core/ggml_extend_backend.h" #include "core/ggml_extend_backend.h"
#include "core/util.h" #include "core/util.h"
#include "ggml-backend.h" #include "ggml-backend.h"
namespace sd::backend_fit { namespace sd::backend_fit {
namespace {
constexpr int64_t MiB = 1024ll * 1024; static constexpr int64_t MiB = 1024ll * 1024;
enum class ComponentKind { enum class ComponentKind {
DIT = 0, DIT,
VAE = 1, CONDITIONER,
CONDITIONER = 2, VAE,
}; };
struct Component { struct Component {
ComponentKind kind;
const char* name;
int64_t params_bytes = 0;
int64_t reserve_bytes = 0;
int64_t staging_bytes = 0;
};
struct Device {
std::string name;
std::string description;
int64_t free_bytes = 0;
int64_t budget_bytes = 0;
};
enum class ParamsLocation {
MAIN_GPU,
CPU,
OTHER_GPU,
DISK,
};
struct Decision {
ParamsLocation params_location = ParamsLocation::DISK;
size_t params_device = SIZE_MAX;
};
struct Plan {
bool valid = false;
size_t main_device = SIZE_MAX;
std::vector<Decision> decisions;
};
static 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;
}
static std::vector<Component> estimate_components(ModelLoader& loader, ggml_type override_wtype) {
int64_t bytes[3] = {0, 0, 0};
int64_t largest_tensor[3] = {0, 0, 0};
for (const auto& [name, stored_tensor] : loader.get_tensor_storage_map()) {
TensorStorage ts = stored_tensor;
ComponentKind kind; ComponentKind kind;
const char* name; if (is_unused_tensor(ts.name) || !classify_tensor(ts.name, kind)) {
int64_t params_bytes = 0; continue;
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<size_t> device_idxs;
};
struct Plan {
bool valid = false;
bool time_share = false;
std::vector<Decision> 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.") || if (ts.expected_type != GGML_TYPE_COUNT) {
name.rfind("vae.", 0) == 0 || ts.type = ts.expected_type;
name.rfind("tae.", 0) == 0) { } else if (override_wtype != GGML_TYPE_COUNT && loader.tensor_should_be_converted(ts, override_wtype)) {
out = ComponentKind::VAE; ts.type = override_wtype;
return true;
} }
if (contains("text_encoders") || const int64_t tensor_bytes = (int64_t)ts.nbytes() + 64;
contains("cond_stage_model") || bytes[int(kind)] += tensor_bytes;
contains("te.text_model.") || largest_tensor[int(kind)] = std::max(largest_tensor[int(kind)], tensor_bytes);
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<Component> estimate_components(ModelLoader& loader, ggml_type override_wtype) { return {
const auto& storage = loader.get_tensor_storage_map(); {ComponentKind::DIT, "DiT", bytes[int(ComponentKind::DIT)], 2048 * MiB, largest_tensor[int(ComponentKind::DIT)]},
{ComponentKind::CONDITIONER, "Conditioner", bytes[int(ComponentKind::CONDITIONER)], 2048 * MiB, largest_tensor[int(ComponentKind::CONDITIONER)]},
{ComponentKind::VAE, "VAE", bytes[int(ComponentKind::VAE)], 1024 * MiB, largest_tensor[int(ComponentKind::VAE)]},
};
}
int64_t bytes[3] = {0, 0, 0}; static std::string budget_key(std::string name) {
for (const auto& [name, ts_const] : storage) { std::transform(name.begin(), name.end(), name.begin(), [](unsigned char c) { return (char)std::tolower(c); });
TensorStorage ts = ts_const; return name;
if (is_unused_tensor(ts.name)) { }
continue;
} static std::vector<Device> enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) {
ComponentKind kind; std::vector<Device> out;
if (!classify_tensor(ts.name, kind)) { for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
continue; ggml_backend_dev_t dev = ggml_backend_dev_get(i);
} if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) {
if (override_wtype != GGML_TYPE_COUNT && continue;
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;
} }
Device device;
device.name = ggml_backend_dev_name(dev);
device.description = ggml_backend_dev_description(dev);
size_t free_bytes = 0, total_bytes = 0;
ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
device.free_bytes = (int64_t)free_bytes;
std::vector<Component> out; float gib = budgets.default_gib;
out.push_back({ComponentKind::DIT, "DiT", bytes[int(ComponentKind::DIT)], 2048 * MiB, true}); auto it = budgets.backend_gib.find(budget_key(device.name));
out.push_back({ComponentKind::VAE, "VAE", bytes[int(ComponentKind::VAE)], 1024 * MiB, false}); if (it != budgets.backend_gib.end()) {
out.push_back({ComponentKind::CONDITIONER, "Conditioner", bytes[int(ComponentKind::CONDITIONER)], 2048 * MiB, true}); gib = it->second;
return out; }
if (gib > 0.f) {
device.budget_bytes = (int64_t)std::min(gib * 1024.0 * MiB, (double)device.free_bytes);
} else if (gib < 0.f) {
device.budget_bytes = (int64_t)std::max<double>(device.free_bytes + gib * 1024.0 * MiB, 0);
} else {
device.budget_bytes = std::max<int64_t>(device.free_bytes - 512 * MiB, 0);
}
out.push_back(std::move(device));
} }
return out;
}
std::vector<Device> enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) { static int64_t available_ram_bytes() {
std::vector<Device> out; #if defined(_WIN32)
for (size_t i = 0; i < ggml_backend_dev_count(); i++) { MEMORYSTATUSEX status{};
ggml_backend_dev_t dev = ggml_backend_dev_get(i); status.dwLength = sizeof(status);
if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) { if (GlobalMemoryStatusEx(&status)) {
continue; return (int64_t)status.ullAvailPhys;
}
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>((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<int64_t>(d.budget_bytes, 0);
out.push_back(d);
}
return out;
} }
#elif defined(__linux__)
Plan compute_plan(const std::vector<Component>& components, const std::vector<Device>& devices) { std::ifstream meminfo("/proc/meminfo");
Plan plan; std::string key, unit;
if (devices.empty()) { int64_t kib = 0;
return plan; while (meminfo >> key >> kib >> unit) {
if (key == "MemAvailable:" && unit == "kB" && kib >= 0) {
return kib * 1024;
} }
}
#elif defined(__APPLE__)
const mach_port_t host = mach_host_self();
vm_size_t page_size = 0;
vm_statistics64_data_t stats{};
mach_msg_type_number_t count = HOST_VM_INFO64_COUNT;
const bool ok = host_page_size(host, &page_size) == KERN_SUCCESS &&
host_statistics64(host, HOST_VM_INFO64, (host_info64_t)&stats, &count) == KERN_SUCCESS;
mach_port_deallocate(mach_task_self(), host);
if (ok) {
return ((int64_t)stats.free_count + stats.inactive_count) * page_size;
}
#endif
return -1;
}
std::vector<size_t> order(components.size()); static Plan compute_plan(const std::vector<Component>& components,
for (size_t i = 0; i < order.size(); i++) { const std::vector<Device>& devices,
order[i] = i; int64_t ram_budget_bytes) {
Plan plan;
for (size_t di = 0; di < devices.size(); ++di) {
if (devices[di].budget_bytes > 0 &&
(plan.main_device == SIZE_MAX || devices[di].budget_bytes > devices[plan.main_device].budget_bytes)) {
plan.main_device = di;
} }
std::sort(order.begin(), order.end(), [&](size_t a, size_t b) { }
return components[a].params_bytes > components[b].params_bytes; if (plan.main_device == SIZE_MAX) {
});
{
std::vector<int64_t> params_sum(devices.size(), 0);
std::vector<int64_t> max_reserve(devices.size(), 0);
std::vector<Decision> 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<int64_t>(d.budget_bytes - comp.reserve_bytes, 0);
}
if (comp.params_bytes <= capacity) {
std::vector<size_t> 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; return plan;
} }
void print_plan(const Plan& plan, std::vector<size_t> order(components.size());
const std::vector<Component>& components, for (size_t ci = 0; ci < components.size(); ++ci) {
const std::vector<Device>& devices) { order[ci] = ci;
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());
}
} }
std::stable_sort(order.begin(), order.end(), [&](size_t a, size_t b) {
return components[a].kind < components[b].kind;
});
void append_assignment(std::string& spec, const char* key, const std::string& value) { std::vector<int64_t> remaining;
if (!spec.empty()) { for (const Device& device : devices) {
spec += ","; remaining.push_back(std::max<int64_t>(device.budget_bytes, 0));
}
spec += key;
spec += "=";
spec += value;
} }
ram_budget_bytes = std::max<int64_t>(ram_budget_bytes, 0);
plan.decisions.resize(components.size());
void append_component_decision(const std::vector<Component>& components, for (size_t ci : order) {
const std::vector<Device>& devices, const Component& comp = components[ci];
const Plan& plan, Decision& decision = plan.decisions[ci];
ComponentKind kind, if (comp.params_bytes == 0) {
const char* module_key, continue;
std::string& runtime_spec, }
std::string& params_spec) {
for (size_t ci = 0; ci < components.size(); ci++) { // Higher-priority offloaded weights need GPU cache space across graph runs.
if (components[ci].kind != kind || components[ci].params_bytes == 0) { int64_t headroom = 0;
for (size_t other = 0; other < components.size(); ++other) {
if (components[other].params_bytes == 0) {
continue; continue;
} }
const Decision& decision = plan.decisions[ci]; const bool resident = other == ci || plan.decisions[other].params_location == ParamsLocation::MAIN_GPU;
if (decision.on_cpu) { const int64_t cached_weights = components[other].kind < comp.kind
append_assignment(runtime_spec, module_key, "cpu"); ? components[other].params_bytes
return; : components[other].staging_bytes;
headroom = std::max(headroom, components[other].reserve_bytes +
(resident ? 0 : cached_weights));
}
int64_t& main_remaining = remaining[plan.main_device];
if (headroom <= main_remaining && comp.params_bytes <= main_remaining - headroom) {
decision.params_location = ParamsLocation::MAIN_GPU;
decision.params_device = plan.main_device;
main_remaining -= comp.params_bytes;
continue;
}
if (comp.params_bytes <= ram_budget_bytes) {
decision.params_location = ParamsLocation::CPU;
ram_budget_bytes -= comp.params_bytes;
continue;
}
size_t best = SIZE_MAX;
for (size_t di = 0; di < devices.size(); ++di) {
if (di != plan.main_device && comp.params_bytes <= remaining[di] &&
(best == SIZE_MAX || remaining[di] > remaining[best])) {
best = di;
} }
if (decision.device_idxs.empty()) { }
return; if (best != SIZE_MAX) {
} decision.params_location = ParamsLocation::OTHER_GPU;
std::string device_list; decision.params_device = best;
for (size_t k = 0; k < decision.device_idxs.size(); k++) { remaining[best] -= comp.params_bytes;
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;
} }
} }
plan.valid = true;
return plan;
}
} // namespace static std::string params_backend_name(const Decision& decision, const std::vector<Device>& devices) {
switch (decision.params_location) {
case ParamsLocation::MAIN_GPU:
case ParamsLocation::OTHER_GPU:
return devices[decision.params_device].name;
case ParamsLocation::CPU:
return "cpu";
case ParamsLocation::DISK:
return "disk";
}
return "disk";
}
static void print_plan(const Plan& plan,
const std::vector<Component>& components,
const std::vector<Device>& devices,
int64_t free_ram,
int64_t ram_budget) {
LOG_INFO("auto-fit plan (single-GPU compute on %s):", devices[plan.main_device].name.c_str());
LOG_INFO(" devices:");
for (const Device& device : devices) {
LOG_INFO(" %-12s %-32s free %6lld MiB, budget %6lld MiB",
device.name.c_str(), device.description.c_str(),
(long long)(device.free_bytes / MiB), (long long)(device.budget_bytes / MiB));
}
if (free_ram < 0) {
LOG_WARN("auto-fit: available RAM is unknown; skipping CPU parameter residency");
} else {
LOG_INFO(" RAM free %6lld MiB, params budget %6lld MiB",
(long long)(free_ram / MiB), (long long)(ram_budget / MiB));
}
LOG_INFO(" main-GPU weight cache priority: diffusion > te > vae");
LOG_INFO(" components (params: main GPU -> RAM -> other GPU -> disk):");
for (size_t ci = 0; ci < components.size(); ++ci) {
const Component& comp = components[ci];
if (comp.params_bytes == 0) {
continue;
}
const std::string params = params_backend_name(plan.decisions[ci], devices);
LOG_INFO(" %-12s params %6lld MiB, compute reserve %5lld MiB -> compute %s, params %s",
comp.name, (long long)(comp.params_bytes / MiB), (long long)(comp.reserve_bytes / MiB),
devices[plan.main_device].name.c_str(), params.c_str());
}
}
static void append_assignment(std::string& spec, const char* key, const std::string& value) {
if (!spec.empty()) {
spec += ",";
}
spec += key;
spec += "=";
spec += value;
}
static const char* module_key(ComponentKind kind) {
switch (kind) {
case ComponentKind::DIT:
return "diffusion";
case ComponentKind::CONDITIONER:
return "te";
case ComponentKind::VAE:
return "vae";
}
return "";
}
bool derive_backend_specs(ModelLoader& loader, bool derive_backend_specs(ModelLoader& loader,
ggml_type override_wtype, ggml_type override_wtype,
sd::ggml_graph_cut::MaxVramAssignment& budgets, sd::ggml_graph_cut::MaxVramAssignment& budgets,
std::string& runtime_spec, std::string& runtime_spec,
std::string& params_spec) { std::string& params_spec) {
if (!runtime_spec.empty() || !params_spec.empty()) { std::string error;
LOG_WARN("--auto-fit is enabled; ignoring --backend / --params-backend"); if (!budgets.canonicalize_backend_keys(&error)) {
LOG_ERROR("%s", error.c_str());
return false;
} }
{ const auto components = estimate_components(loader, override_wtype);
std::string error; const auto devices = enumerate_gpu_devices(budgets);
if (!budgets.canonicalize_backend_keys(&error)) { const int64_t free_ram = available_ram_bytes();
LOG_ERROR("%s", error.c_str()); const int64_t ram_budget = std::max<int64_t>(free_ram - std::max<int64_t>(2048 * MiB, free_ram / 10), 0);
return false; const auto plan = compute_plan(components, devices, ram_budget);
} runtime_spec.clear();
} params_spec.clear();
auto components = estimate_components(loader, override_wtype);
auto devices = enumerate_gpu_devices(budgets);
auto plan = compute_plan(components, devices);
if (!plan.valid) { if (!plan.valid) {
LOG_WARN("auto-fit: no usable GPU devices; using the default backend"); if (devices.empty()) {
runtime_spec.clear(); LOG_WARN("auto-fit: no GPU devices; using the default backend");
params_spec.clear(); } else {
LOG_WARN("auto-fit: no GPU memory budget available; using CPU");
runtime_spec = "cpu";
}
return true; return true;
} }
print_plan(plan, components, devices); print_plan(plan, components, devices, free_ram, ram_budget);
for (size_t ci = 0; ci < components.size(); ++ci) {
if (components[ci].params_bytes == 0) {
continue;
}
const char* key = module_key(components[ci].kind);
append_assignment(runtime_spec, key, devices[plan.main_device].name);
if (plan.decisions[ci].params_location != ParamsLocation::MAIN_GPU) {
append_assignment(params_spec, key, params_backend_name(plan.decisions[ci], devices));
}
}
std::string derived_runtime_spec; // Keep the planner's safety margin when the runner resolves its device limits.
std::string derived_params_spec; for (const Device& device : devices) {
append_component_decision(components, devices, plan, ComponentKind::DIT, "diffusion", derived_runtime_spec, derived_params_spec); if (device.budget_bytes > 0) {
append_component_decision(components, devices, plan, ComponentKind::CONDITIONER, "te", derived_runtime_spec, derived_params_spec); budgets.backend_gib[budget_key(device.name)] = (float)(device.budget_bytes / (1024.0 * MiB));
append_component_decision(components, devices, plan, ComponentKind::VAE, "vae", derived_runtime_spec, derived_params_spec); }
}
runtime_spec = std::move(derived_runtime_spec); budgets.resolved_backend_bytes.clear();
params_spec = std::move(derived_params_spec);
LOG_INFO("auto-fit: --backend \"%s\"%s%s%s", LOG_INFO("auto-fit: --backend \"%s\"%s%s%s",
runtime_spec.empty() ? "(default)" : runtime_spec.c_str(), runtime_spec.empty() ? "(default)" : runtime_spec.c_str(),
params_spec.empty() ? "" : " --params-backend \"", params_spec.empty() ? "" : " --params-backend \"",
params_spec.c_str(), params_spec.c_str(), params_spec.empty() ? "" : "\"");
params_spec.empty() ? "" : "\"");
return true; return true;
} }

View File

@ -874,7 +874,7 @@ public:
backend_spec = SAFE_STR(sd_ctx_params->backend); backend_spec = SAFE_STR(sd_ctx_params->backend);
params_backend_spec = SAFE_STR(sd_ctx_params->params_backend); params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
split_mode_spec = SAFE_STR(sd_ctx_params->split_mode); split_mode_spec = SAFE_STR(sd_ctx_params->split_mode);
auto_fit_enabled = sd_ctx_params->auto_fit; auto_fit_enabled = sd_ctx_params->auto_fit && backend_spec.empty() && params_backend_spec.empty();
max_vram_assignment.reset(0.f); max_vram_assignment.reset(0.f);
{ {
std::string error; std::string error;
@ -3578,7 +3578,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
sd_ctx_params->backend = nullptr; sd_ctx_params->backend = nullptr;
sd_ctx_params->params_backend = nullptr; sd_ctx_params->params_backend = nullptr;
sd_ctx_params->split_mode = nullptr; sd_ctx_params->split_mode = nullptr;
sd_ctx_params->auto_fit = false; sd_ctx_params->auto_fit = true;
sd_ctx_params->rpc_servers = nullptr; sd_ctx_params->rpc_servers = nullptr;
sd_ctx_params->model_args = nullptr; sd_ctx_params->model_args = nullptr;
sd_ctx_params->pulid_weights_path = nullptr; sd_ctx_params->pulid_weights_path = nullptr;