mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-09-23 22:47:41 -05:00
Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e790073e1c | ||
|
|
b11c95a41c | ||
|
|
68f3d6df9f | ||
|
|
7bcd189639 | ||
|
|
2574f59365 |
@@ -51,6 +51,97 @@ 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.
|
||||
|
||||
### 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.
|
||||
|
||||
## 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 |
|
||||
|
||||
@@ -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",
|
||||
@@ -501,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",
|
||||
@@ -663,6 +676,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 'name<TAB>description' 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<char> devices(device_list_size + 1);
|
||||
sd_list_devices(devices.data(), devices.size());
|
||||
fputs(devices.data(), stdout);
|
||||
std::exit(0);
|
||||
return 0;
|
||||
}},
|
||||
};
|
||||
|
||||
return options;
|
||||
@@ -818,6 +843,8 @@ 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"
|
||||
<< " 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"
|
||||
@@ -898,6 +925,8 @@ 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.auto_fit = auto_fit;
|
||||
sd_ctx_params.rpc_servers = rpc_servers.c_str();
|
||||
return sd_ctx_params;
|
||||
}
|
||||
|
||||
@@ -151,6 +151,8 @@ struct SDContextParams {
|
||||
bool eager_load = false;
|
||||
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;
|
||||
|
||||
@@ -227,6 +227,8 @@ 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"
|
||||
bool auto_fit;
|
||||
const char* rpc_servers;
|
||||
} sd_ctx_params_t;
|
||||
|
||||
@@ -534,6 +536,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 `name<TAB>description` 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);
|
||||
|
||||
@@ -116,6 +116,8 @@ public:
|
||||
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& 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<ggml_backend_t>& backends) {}
|
||||
virtual void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {}
|
||||
virtual void set_flash_attention_enabled(bool enabled) = 0;
|
||||
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
|
||||
virtual void runner_done() {}
|
||||
@@ -635,6 +637,18 @@ struct SD3CLIPEmbedder : public Conditioner {
|
||||
}
|
||||
}
|
||||
|
||||
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
|
||||
if (t5) {
|
||||
t5->set_runtime_backends(backends);
|
||||
}
|
||||
}
|
||||
|
||||
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& 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<ggml_backend_t>& backends) override {
|
||||
if (t5) {
|
||||
t5->set_runtime_backends(backends);
|
||||
}
|
||||
}
|
||||
|
||||
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& 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<ggml_backend_t>& backends) override {
|
||||
if (t5) {
|
||||
t5->set_runtime_backends(backends);
|
||||
}
|
||||
}
|
||||
|
||||
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& 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<ggml_backend_t>& backends) override {
|
||||
if (t5) {
|
||||
t5->set_runtime_backends(backends);
|
||||
}
|
||||
}
|
||||
|
||||
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& 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<ggml_backend_t>& backends) override {
|
||||
llm->set_runtime_backends(backends);
|
||||
}
|
||||
|
||||
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& 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<ggml_backend_t>& backends) override {
|
||||
llm->set_runtime_backends(backends);
|
||||
}
|
||||
|
||||
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& 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<ggml_backend_t>& backends) override {
|
||||
llm->set_runtime_backends(backends);
|
||||
}
|
||||
|
||||
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||
llm->get_param_tensors(tensors, "text_encoders.llm");
|
||||
}
|
||||
|
||||
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
|
||||
llm->set_weight_adapter(adapter);
|
||||
projector->set_weight_adapter(adapter);
|
||||
|
||||
@@ -0,0 +1,390 @@
|
||||
#include "backend_fit.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cctype>
|
||||
#include <cstdint>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#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<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.") ||
|
||||
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<Component> 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<Component> 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<Device> enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) {
|
||||
std::vector<Device> 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>((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;
|
||||
}
|
||||
|
||||
Plan compute_plan(const std::vector<Component>& components, const std::vector<Device>& devices) {
|
||||
Plan plan;
|
||||
if (devices.empty()) {
|
||||
return plan;
|
||||
}
|
||||
|
||||
std::vector<size_t> 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<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;
|
||||
}
|
||||
|
||||
void print_plan(const Plan& plan,
|
||||
const std::vector<Component>& components,
|
||||
const std::vector<Device>& 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<Component>& components,
|
||||
const std::vector<Device>& 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
|
||||
@@ -0,0 +1,23 @@
|
||||
#ifndef __SD_BACKEND_FIT_H__
|
||||
#define __SD_BACKEND_FIT_H__
|
||||
|
||||
#include <string>
|
||||
|
||||
#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__
|
||||
+192
-10
@@ -1388,6 +1388,9 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
|
||||
}
|
||||
|
||||
auto out = ggml_flash_attn_ext(ctx, q_in, k_in, v_in, mask_in, scale / kv_scale, 0, 0);
|
||||
if (!ggml_backend_supports_op(backend, out)) {
|
||||
return nullptr;
|
||||
}
|
||||
ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32);
|
||||
if (kv_scale != 1.0f) {
|
||||
out = ggml_ext_scale(ctx, out, 1.0f / kv_scale);
|
||||
@@ -1405,9 +1408,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
|
||||
|
||||
if (can_use_flash_attn) {
|
||||
kqv = build_kqv(q, k, v, mask);
|
||||
if (!ggml_backend_supports_op(backend, kqv)) {
|
||||
kqv = nullptr;
|
||||
} else {
|
||||
if (kqv != nullptr) {
|
||||
kqv = ggml_view_4d(ctx,
|
||||
kqv,
|
||||
d_head,
|
||||
@@ -1745,6 +1746,11 @@ protected:
|
||||
bool stream_layers_enabled = false;
|
||||
size_t observed_max_effective_budget_ = 0;
|
||||
|
||||
std::vector<ggml_backend_t> 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<WeightAdapter> weight_adapter = nullptr;
|
||||
std::weak_ptr<RunnerWeightManager> weight_manager;
|
||||
std::unordered_set<const ggml_tensor*> kept_compute_param_tensor_set;
|
||||
@@ -2012,7 +2018,121 @@ protected:
|
||||
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<ggml_backend_t> 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<ggml_backend_buffer_type_t> 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<size_t>(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;
|
||||
}
|
||||
@@ -2228,12 +2348,14 @@ protected:
|
||||
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,
|
||||
@@ -2489,7 +2611,14 @@ protected:
|
||||
};
|
||||
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;
|
||||
}
|
||||
@@ -2498,11 +2627,27 @@ protected:
|
||||
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;
|
||||
@@ -2679,6 +2824,10 @@ public:
|
||||
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() {
|
||||
@@ -2719,10 +2868,20 @@ public:
|
||||
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;
|
||||
}
|
||||
|
||||
@@ -2858,8 +3017,31 @@ public:
|
||||
}
|
||||
|
||||
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<ggml_backend_t>& 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 {
|
||||
|
||||
@@ -665,12 +665,52 @@ SDBackendManager::~SDBackendManager() {
|
||||
|
||||
void SDBackendManager::reset() {
|
||||
backends_.clear();
|
||||
runtime_assignment_ = {};
|
||||
params_assignment_ = {};
|
||||
runtime_assignment_ = {};
|
||||
params_assignment_ = {};
|
||||
split_mode_assignment_ = {};
|
||||
}
|
||||
|
||||
static std::vector<std::string> split_device_list(const std::string& value) {
|
||||
std::vector<std::string> 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<std::string> 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<ggml_backend_t> SDBackendManager::runtime_backends(SDBackendModule module) {
|
||||
std::vector<ggml_backend_t> 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 +736,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) {
|
||||
@@ -715,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();
|
||||
|
||||
@@ -724,12 +769,53 @@ 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<float>& 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<float> padded_split(std::max<size_t>(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_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,15 +833,56 @@ 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<std::string> 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);
|
||||
};
|
||||
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) {
|
||||
@@ -768,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;
|
||||
}
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#include "ggml-backend.h"
|
||||
#include "ggml.h"
|
||||
@@ -37,10 +38,16 @@ struct SDBackendHandleDeleter {
|
||||
|
||||
using SDBackendHandle = std::unique_ptr<struct ggml_backend, SDBackendHandleDeleter>;
|
||||
|
||||
enum class SDSplitMode {
|
||||
LAYER,
|
||||
ROW,
|
||||
};
|
||||
|
||||
class SDBackendManager {
|
||||
private:
|
||||
SDBackendAssignment runtime_assignment_;
|
||||
SDBackendAssignment params_assignment_;
|
||||
SDBackendAssignment split_mode_assignment_;
|
||||
std::unordered_map<std::string, SDBackendHandle> backends_;
|
||||
|
||||
public:
|
||||
@@ -52,15 +59,23 @@ public:
|
||||
|
||||
bool init(const char* backend_spec,
|
||||
const char* params_backend_spec,
|
||||
const char* split_mode_spec,
|
||||
std::string* error);
|
||||
void reset();
|
||||
|
||||
ggml_backend_t runtime_backend(SDBackendModule module);
|
||||
ggml_backend_t params_backend(SDBackendModule module);
|
||||
|
||||
std::vector<ggml_backend_t> 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<float>& tensor_split);
|
||||
|
||||
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:
|
||||
|
||||
@@ -0,0 +1,221 @@
|
||||
#include "core/layer_split_partition.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
|
||||
#include "core/util.h"
|
||||
|
||||
namespace sd {
|
||||
|
||||
static bool layer_split_path_segment_starts_at(const std::string& name, size_t pos) {
|
||||
return pos == 0 || name[pos - 1] == '.';
|
||||
}
|
||||
|
||||
static bool layer_split_has_path_segment(const std::string& name, const char* segment) {
|
||||
size_t pos = name.find(segment);
|
||||
while (pos != std::string::npos) {
|
||||
if (layer_split_path_segment_starts_at(name, pos)) {
|
||||
return true;
|
||||
}
|
||||
pos = name.find(segment, pos + 1);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
int layer_split_tensor_block_index(const std::string& name) {
|
||||
static const char* unet_block_segments[] = {"input_blocks.", "output_blocks.", "middle_block.",
|
||||
"down_blocks.", "up_blocks.", "mid_block."};
|
||||
for (const char* segment : unet_block_segments) {
|
||||
if (layer_split_has_path_segment(name, segment)) {
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
|
||||
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);
|
||||
while (pos != std::string::npos) {
|
||||
if (!layer_split_path_segment_starts_at(name, pos)) {
|
||||
pos = name.find(keyword, pos + 1);
|
||||
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());
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
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<ggml_backend_t>& 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<std::map<std::string, ggml_tensor*>> partition_layer_split_tensors(
|
||||
const std::string& desc,
|
||||
const std::map<std::string, ggml_tensor*>& tensors,
|
||||
const std::map<std::string, ggml_tensor*>& split_tensors,
|
||||
const std::vector<ggml_backend_t>& backends) {
|
||||
std::vector<std::map<std::string, ggml_tensor*>> partitions(backends.size());
|
||||
if (backends.empty()) {
|
||||
LOG_WARN("%s: no backend available for a layer split", desc.c_str());
|
||||
return partitions;
|
||||
}
|
||||
|
||||
std::map<int, int64_t> block_bytes;
|
||||
std::map<std::string, size_t> non_block_targets;
|
||||
std::vector<int64_t> other_bytes_by_backend(backends.size(), 0);
|
||||
int64_t total_block_bytes = 0;
|
||||
int64_t total_other_bytes = 0;
|
||||
int n_blocks = 0;
|
||||
for (const auto& kv : tensors) {
|
||||
int64_t bytes = (int64_t)ggml_nbytes(kv.second);
|
||||
int idx = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1;
|
||||
if (idx >= 0) {
|
||||
block_bytes[idx] += bytes;
|
||||
total_block_bytes += bytes;
|
||||
n_blocks = std::max(n_blocks, idx + 1);
|
||||
} else {
|
||||
size_t target = layer_split_supported_target(desc, kv.first, kv.second, backends, 0);
|
||||
non_block_targets[kv.first] = target;
|
||||
other_bytes_by_backend[target] += bytes;
|
||||
total_other_bytes += bytes;
|
||||
}
|
||||
}
|
||||
if (n_blocks == 0) {
|
||||
LOG_WARN("%s: no transformer blocks found for a layer split; keeping tensors on compatible backends starting from %s",
|
||||
desc.c_str(),
|
||||
layer_split_backend_device_display_name(backends[0]).c_str());
|
||||
for (const auto& kv : tensors) {
|
||||
size_t target = 0;
|
||||
auto target_it = non_block_targets.find(kv.first);
|
||||
if (target_it != non_block_targets.end()) {
|
||||
target = target_it->second;
|
||||
}
|
||||
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<double> 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>((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<int64_t> 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<int64_t>(budget - other_bytes_by_backend[i], 0);
|
||||
block_budgets[i] = budget;
|
||||
}
|
||||
|
||||
std::vector<int> 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 ? layer_split_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
|
||||
@@ -0,0 +1,24 @@
|
||||
#ifndef __SD_CORE_LAYER_SPLIT_PARTITION_H__
|
||||
#define __SD_CORE_LAYER_SPLIT_PARTITION_H__
|
||||
|
||||
#include <map>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "ggml-backend.h"
|
||||
#include "ggml.h"
|
||||
|
||||
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<std::map<std::string, ggml_tensor*>> partition_layer_split_tensors(
|
||||
const std::string& desc,
|
||||
const std::map<std::string, ggml_tensor*>& tensors,
|
||||
const std::map<std::string, ggml_tensor*>& split_tensors,
|
||||
const std::vector<ggml_backend_t>& backends);
|
||||
|
||||
} // namespace sd
|
||||
|
||||
#endif // __SD_CORE_LAYER_SPLIT_PARTITION_H__
|
||||
@@ -4,6 +4,8 @@
|
||||
#include <cmath>
|
||||
#include <codecvt>
|
||||
#include <cstdarg>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <exception>
|
||||
#include <fstream>
|
||||
#include <locale>
|
||||
@@ -25,6 +27,7 @@
|
||||
#include <unistd.h>
|
||||
#endif
|
||||
|
||||
#include "ggml-backend.h"
|
||||
#include "ggml.h"
|
||||
#include "stable-diffusion.h"
|
||||
|
||||
@@ -997,3 +1000,26 @@ std::vector<std::pair<std::string, float>> 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();
|
||||
}
|
||||
|
||||
@@ -27,6 +27,8 @@ struct MmapTensorStore {
|
||||
std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
|
||||
};
|
||||
|
||||
bool is_unused_tensor(const std::string& name);
|
||||
|
||||
class ModelLoader {
|
||||
protected:
|
||||
SDVersion version_ = VERSION_COUNT;
|
||||
|
||||
+61
-13
@@ -100,12 +100,41 @@ size_t estimate_tensors_size(const std::map<std::string, ggml_tensor*>& 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<std::string, ggml_tensor*> 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<TensorState>();
|
||||
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<TensorState>();
|
||||
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<TensorState*
|
||||
}
|
||||
|
||||
bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorState*>& states) {
|
||||
std::map<ggml_backend_t, std::vector<TensorState*>> states_by_compute_backend;
|
||||
std::map<std::pair<ggml_backend_t, ggml_backend_buffer_type_t>, std::vector<TensorState*>> 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::vector<TensorStat
|
||||
LOG_ERROR("model manager tensor '%s' is not loaded to params backend", state->name.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<TensorState*>& states = pair.second;
|
||||
if (states.empty()) {
|
||||
continue;
|
||||
@@ -285,7 +320,7 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorStat
|
||||
staged_tensors.push_back({state, staging_tensor});
|
||||
}
|
||||
|
||||
ggml_backend_buffer_t compute_buffer = ggml_backend_alloc_ctx_tensors(staging_ctx, compute_backend);
|
||||
ggml_backend_buffer_t compute_buffer = ggml_backend_alloc_ctx_tensors_from_buft(staging_ctx, staging_buft);
|
||||
if (compute_buffer == nullptr) {
|
||||
LOG_ERROR("model manager alloc compute params backend buffer failed, num_tensors = %zu",
|
||||
staged_tensors.size());
|
||||
@@ -350,6 +385,17 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& 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);
|
||||
|
||||
+9
-1
@@ -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<std::string, TensorState*> tensor_states_by_name_;
|
||||
std::vector<std::unique_ptr<ParamsStorageBlock>> params_storage_blocks_;
|
||||
std::vector<std::unique_ptr<ComputeStagingBlock>> compute_staging_blocks_;
|
||||
std::map<ggml_backend_t, ggml_backend_buffer_type_t> split_buffer_types_;
|
||||
bool warned_split_lora_skip_ = false;
|
||||
std::set<std::string> common_ignore_tensors_;
|
||||
std::vector<LoraSpec> loras_;
|
||||
SDVersion lora_version_ = VERSION_COUNT;
|
||||
@@ -91,6 +94,7 @@ private:
|
||||
bool stage_tensors_to_compute_backend(const std::vector<TensorState*>& 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<TensorState*>* 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<std::string> ignore_tensors);
|
||||
void set_loras(std::vector<LoraSpec> 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<std::string> 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 <typename Runner>
|
||||
bool register_runner_params(const std::string& desc,
|
||||
|
||||
+285
-18
@@ -2,11 +2,14 @@
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
#include <set>
|
||||
#include <type_traits>
|
||||
#include <unordered_set>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#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"
|
||||
@@ -17,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"
|
||||
@@ -170,6 +174,13 @@ static float get_cache_reuse_threshold(const sd_cache_params_t& params) {
|
||||
|
||||
/*=============================================== StableDiffusionGGML ================================================*/
|
||||
|
||||
template <typename T, typename = void>
|
||||
struct has_set_runtime_backends : std::false_type {};
|
||||
template <typename T>
|
||||
struct has_set_runtime_backends<T,
|
||||
std::void_t<decltype(std::declval<T&>().set_runtime_backends(
|
||||
std::declval<const std::vector<ggml_backend_t>&>()))>> : std::true_type {};
|
||||
|
||||
static_assert(std::atomic<sd_cancel_mode_t>::is_always_lock_free,
|
||||
"sd_cancel_mode_t must be lock-free");
|
||||
|
||||
@@ -208,6 +219,8 @@ public:
|
||||
bool eager_load = false;
|
||||
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;
|
||||
@@ -275,18 +288,230 @@ public:
|
||||
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<ggml_backend_t> module_backends = backend_manager.runtime_backends(module);
|
||||
if (module_backends.size() > 1) {
|
||||
if constexpr (has_set_runtime_backends<T>::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,
|
||||
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);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
bool register_row_split_runner_params(const std::string& desc,
|
||||
const std::shared_ptr<T>& model,
|
||||
SDBackendModule module,
|
||||
const std::vector<ggml_backend_t>& module_backends,
|
||||
std::map<std::string, ggml_tensor*> 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<float> 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>((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<std::string, ggml_tensor*> split_tensors;
|
||||
if constexpr (std::is_base_of_v<Conditioner, T>) {
|
||||
model->get_layer_split_param_tensors(split_tensors);
|
||||
} else {
|
||||
split_tensors = group_tensors;
|
||||
}
|
||||
|
||||
std::map<std::string, ggml_tensor*> row_split_map;
|
||||
std::map<std::string, ggml_tensor*> 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 <typename T>
|
||||
bool register_layer_split_runner_params(const std::string& desc,
|
||||
const std::shared_ptr<T>& model,
|
||||
SDBackendModule module,
|
||||
const std::vector<ggml_backend_t>& module_backends,
|
||||
std::map<std::string, ggml_tensor*> 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 <module>=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<std::string, ggml_tensor*> split_tensors;
|
||||
if constexpr (std::is_base_of_v<Conditioner, T>) {
|
||||
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(),
|
||||
params_backend_spec.c_str(),
|
||||
split_mode_spec.c_str(),
|
||||
&error)) {
|
||||
LOG_ERROR("backend config failed: %s", error.c_str());
|
||||
return false;
|
||||
@@ -294,6 +519,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<RNG> get_rng(rng_type_t rng_type) {
|
||||
if (rng_type == STD_DEFAULT_RNG) {
|
||||
return std::make_shared<STDDefaultRNG>();
|
||||
@@ -352,6 +587,8 @@ 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);
|
||||
auto_fit_enabled = sd_ctx_params->auto_fit;
|
||||
max_vram_assignment.reset(0.f);
|
||||
{
|
||||
std::string error;
|
||||
@@ -377,21 +614,6 @@ public:
|
||||
|
||||
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<ModelManager>();
|
||||
model_manager->set_n_threads(n_threads);
|
||||
model_manager->set_enable_mmap(enable_mmap);
|
||||
@@ -539,6 +761,31 @@ public:
|
||||
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<ggml_type, uint32_t> wtype_stat = model_loader.get_wtype_stat();
|
||||
std::map<ggml_type, uint32_t> conditioner_wtype_stat = model_loader.get_conditioner_wtype_stat();
|
||||
std::map<ggml_type, uint32_t> diffusion_model_wtype_stat = model_loader.get_diffusion_model_wtype_stat();
|
||||
@@ -580,12 +827,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;
|
||||
@@ -2450,7 +2702,16 @@ public:
|
||||
}
|
||||
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<LTXVideoVAE>(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<float> normalize_ltx_video_latents(const sd::Tensor<float>& x) {
|
||||
@@ -2806,6 +3067,8 @@ 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->auto_fit = false;
|
||||
sd_ctx_params->rpc_servers = nullptr;
|
||||
sd_ctx_params->pulid_weights_path = nullptr;
|
||||
}
|
||||
@@ -2845,6 +3108,8 @@ 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"
|
||||
"auto_fit: %s\n"
|
||||
"flash_attn: %s\n"
|
||||
"diffusion_flash_attn: %s\n"
|
||||
"circular_x: %s\n"
|
||||
@@ -2881,6 +3146,8 @@ 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->auto_fit),
|
||||
BOOL_STR(sd_ctx_params->flash_attn),
|
||||
BOOL_STR(sd_ctx_params->diffusion_flash_attn),
|
||||
BOOL_STR(sd_ctx_params->circular_x),
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user