upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-08-03 08:40:40 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix: make parameter loading backend-aware (#1828)

Browse Source
This commit is contained in:
yzyyzyhhh
2026-07-29 22:18:44 +08:00
committed by GitHub
parent 53856e7ec8
commit 2993b7fb43
7 changed files with 162 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
src/conditioning/conditioner.hpp
View File

@@ -117,6 +117,7 @@ public:
virtual SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) = 0;
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {}
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) {}
@@ -1664,6 +1665,10 @@ struct AnimaConditioner : public Conditioner {
llm->get_param_tensors(tensors, "text_encoders.llm");
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
llm->get_param_tensor_ops(tensor_ops);
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
llm->set_max_graph_vram_bytes(max_vram_bytes);
}
@@ -1847,6 +1852,10 @@ struct LLMEmbedder : public Conditioner {
}
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
llm->get_param_tensor_ops(tensor_ops);
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
llm->set_max_graph_vram_bytes(max_vram_bytes);
if (byt5) {
@@ -2828,6 +2837,10 @@ struct LTXAVEmbedder : public Conditioner {
projector->get_param_tensors(tensors, "text_embedding_projection");
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
llm->get_param_tensor_ops(tensor_ops);
}
void set_flash_attention_enabled(bool enabled) override {
llm->set_flash_attention_enabled(enabled);
projector->set_flash_attention_enabled(enabled);

58
src/core/ggml_extend.hpp
View File

@@ -1753,7 +1753,7 @@ protected:
std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
std::vector<ggml_backend_t> extra_runtime_backends; // borrowed (SDBackendManager-owned)
ggml_backend_sched_t sched = nullptr; // owned, multi-device only
ggml_backend_sched_t sched = nullptr; // owned
ggml_backend_t cpu_fallback_backend = nullptr; // owned, sched requires a trailing CPU backend
bool multi_device_eval_callback_warned = false;
@@ -2147,8 +2147,22 @@ protected:
return !extra_runtime_backends.empty();
}
bool graph_requires_backend_fallback(ggml_cgraph* gf) const {
if (gf == nullptr || sd_backend_is_cpu(runtime_backend)) {
return false;
}
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);
if (node != nullptr && !ggml_backend_supports_op(runtime_backend, node)) {
return true;
}
}
return false;
}
bool alloc_compute_buffer(ggml_cgraph* gf) {
if (is_multi_device()) {
if (sched != nullptr || is_multi_device() || graph_requires_backend_fallback(gf)) {
// 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
@@ -2156,6 +2170,10 @@ protected:
// 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.
if (compute_allocr != nullptr) {
ggml_gallocr_free(compute_allocr);
compute_allocr = nullptr;
}
return ensure_sched(gf);
}
if (compute_allocr != nullptr) {
@@ -2753,7 +2771,7 @@ protected:
};
ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer);
if (is_multi_device()) {
if (sched != nullptr) {
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)) {
@@ -2774,9 +2792,9 @@ protected:
}
ggml_status status;
if (is_multi_device()) {
if (sched != nullptr) {
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",
LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring",
get_desc().c_str());
multi_device_eval_callback_warned = true;
}
@@ -3018,12 +3036,9 @@ public:
// 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);
}
// The scheduler only allocates standalone data tensors when they are
// marked as graph inputs. The flag is harmless for single-backend graphs.
ggml_set_input(tensor);
backend_tensor_data_map[tensor] = data;
}
@@ -3240,6 +3255,11 @@ protected:
virtual void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {}
virtual enum ggml_op param_usage_op(const std::string& name) const {
(void)name;
return GGML_OP_NONE;
}
public:
void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") {
if (prefix.size() > 0) {
@@ -3290,6 +3310,18 @@ public:
}
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
for (auto& pair : blocks) {
pair.second->get_param_tensor_ops(tensor_ops);
}
for (auto& pair : params) {
enum ggml_op op = param_usage_op(pair.first);
if (op != GGML_OP_NONE) {
tensor_ops[pair.second] = op;
}
}
}
virtual std::string get_desc() {
return "GGMLBlock";
}
@@ -3417,6 +3449,10 @@ protected:
params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings);
}
enum ggml_op param_usage_op(const std::string& name) const override {
return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE;
}
public:
Embedding(int64_t num_embeddings, int64_t embedding_dim)
: embedding_dim(embedding_dim),

4
src/model/te/llm.hpp
View File

@@ -1657,6 +1657,10 @@ namespace LLM {
model.get_param_tensors(tensors, prefix);
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
model.get_param_tensor_ops(tensor_ops);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* input_ids,
ggml_tensor* input_pos,

15
src/model_loader.cpp
View File

@@ -1053,7 +1053,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
std::atomic<size_t> tensor_idx(0);
std::atomic<bool> failed(false);
std::vector<std::thread> workers;
std::mutex rpc_backend_mutex;
std::mutex backend_tensor_set_mutex;
for (int i = 0; i < n_threads; ++i) {
workers.emplace_back([&, file_path, is_zip]() {
@@ -1214,17 +1214,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
if (dst_tensor->buffer != nullptr && !ggml_backend_buffer_is_host(dst_tensor->buffer)) {
t0 = ggml_time_ms();
// RPC backends require serialized access to prevent concurrency issues
const char* buffer_type_name = ggml_backend_buft_name(ggml_backend_buffer_get_type(dst_tensor->buffer));
bool is_rpc_buffer = buffer_type_name != nullptr &&
std::string(buffer_type_name).find("RPC") != std::string::npos;
if (is_rpc_buffer) {
std::lock_guard<std::mutex> lock(rpc_backend_mutex);
ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
} else {
ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
}
std::lock_guard<std::mutex> lock(backend_tensor_set_mutex);
ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
t1 = ggml_time_ms();
copy_to_backend_time_ms.fetch_add(t1 - t0);

67
src/model_manager.cpp
View File

@@ -53,6 +53,48 @@ static bool backend_supports_host_buffer(ggml_backend_t backend) {
return props.caps.buffer_from_host_ptr;
}
static bool device_supports_param_op(ggml_backend_dev_t device,
ggml_tensor* weight,
enum ggml_op op,
ggml_backend_buffer_type_t buft) {
if (op == GGML_OP_NONE) {
return true;
}
if (device == nullptr || weight == nullptr || buft == nullptr || weight->buffer != nullptr) {
return false;
}
ggml_init_params params;
params.mem_size = ggml_tensor_overhead() * 2;
params.mem_buffer = nullptr;
params.no_alloc = true;
ggml_context* ctx = ggml_init(params);
if (ctx == nullptr) {
return false;
}
ggml_tensor* op_tensor = nullptr;
if (op == GGML_OP_GET_ROWS) {
ggml_tensor* indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1);
op_tensor = ggml_get_rows(ctx, weight, indices);
}
if (op_tensor == nullptr) {
ggml_free(ctx);
return false;
}
weight->buffer = ggml_backend_buft_alloc_buffer(buft, 0);
if (weight->buffer == nullptr) {
ggml_free(ctx);
return false;
}
bool supported = ggml_backend_dev_supports_op(device, op_tensor);
ggml_backend_buffer_free(weight->buffer);
weight->buffer = nullptr;
ggml_free(ctx);
return supported;
}
ModelManager::~ModelManager() {
release_all();
}
@@ -135,7 +177,8 @@ bool ModelManager::register_param_tensors(const std::string& desc,
ggml_backend_t params_backend,
size_t* registered_tensor_size,
bool allow_split_buffer,
bool params_follow_compute_backend) {
bool params_follow_compute_backend,
const std::map<ggml_tensor*, enum ggml_op>* tensor_ops) {
if (desc.empty()) {
LOG_ERROR("model manager tensor desc is empty");
return false;
@@ -168,6 +211,12 @@ bool ModelManager::register_param_tensors(const std::string& desc,
state->params_backend = params_backend;
state->allow_split_buffer = allow_split_buffer;
state->params_follow_compute_backend = params_follow_compute_backend;
if (tensor_ops != nullptr) {
auto op_it = tensor_ops->find(tensor);
if (op_it != tensor_ops->end()) {
state->usage_op = op_it->second;
}
}
new_states.push_back(std::move(state));
}
@@ -844,6 +893,22 @@ ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorStat
if (params_buft == nullptr) {
params_buft = ggml_backend_get_default_buffer_type(state.params_backend);
}
if (state.usage_op != GGML_OP_NONE &&
state.compute_backend != nullptr) {
ggml_backend_dev_t compute_dev = ggml_backend_get_device(state.compute_backend);
if (device_supports_param_op(compute_dev, state.tensor, state.usage_op, params_buft)) {
return params_buft;
}
ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
params_buft = cpu_dev != nullptr ? ggml_backend_dev_buffer_type(cpu_dev) : nullptr;
if (!device_supports_param_op(cpu_dev, state.tensor, state.usage_op, params_buft)) {
LOG_ERROR("model manager has no compatible buffer for tensor '%s' used by %s",
state.name.c_str(),
ggml_op_name(state.usage_op));
return nullptr;
}
}
return params_buft;
}

4
src/model_manager.h
View File

@@ -39,6 +39,7 @@ private:
bool allow_split_buffer = false;
bool params_follow_compute_backend = false;
bool metadata_validated = false;
enum ggml_op usage_op = GGML_OP_NONE;
int active_prepare_count = 0;
@@ -132,7 +133,8 @@ public:
ggml_backend_t params_backend,
size_t* registered_tensor_size = nullptr,
bool allow_split_buffer = false,
bool params_follow_compute_backend = false);
bool params_follow_compute_backend = false,
const std::map<ggml_tensor*, enum ggml_op>* tensor_ops = nullptr);
bool unregister_param_tensors(const std::string& desc,
size_t* registered_tensor_size = nullptr);

31
src/stable-diffusion.cpp
View File

@@ -320,7 +320,11 @@ public:
return true;
}
std::map<std::string, ggml_tensor*> group_tensors;
std::map<ggml_tensor*, enum ggml_op> tensor_ops;
model->get_param_tensors(group_tensors);
if constexpr (std::is_base_of_v<Conditioner, T>) {
model->get_param_tensor_ops(tensor_ops);
}
if (model_manager == nullptr) {
return true;
}
@@ -337,6 +341,7 @@ public:
module,
module_backends,
std::move(group_tensors),
tensor_ops,
residency_mode,
params_mem_size);
}
@@ -345,6 +350,7 @@ public:
module,
module_backends,
std::move(group_tensors),
tensor_ops,
residency_mode,
params_mem_size);
}
@@ -358,7 +364,10 @@ public:
residency_mode,
backend_for(module),
params_backend_for(module),
params_mem_size);
params_mem_size,
false,
false,
&tensor_ops);
}
template <typename T>
@@ -367,6 +376,7 @@ public:
SDBackendModule module,
const std::vector<ggml_backend_t>& module_backends,
std::map<std::string, ggml_tensor*> group_tensors,
const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
ModelManager::ResidencyMode residency_mode,
size_t* params_mem_size) {
ggml_backend_t main_backend = module_backends[0];
@@ -378,6 +388,7 @@ public:
module,
module_backends,
std::move(group_tensors),
tensor_ops,
residency_mode,
params_mem_size);
};
@@ -452,7 +463,9 @@ public:
main_backend,
params_backend_for(module),
params_mem_size,
/*allow_split_buffer=*/true)) {
/*allow_split_buffer=*/true,
false,
&tensor_ops)) {
return false;
}
return model_manager->register_param_tensors(desc,
@@ -460,7 +473,10 @@ public:
residency_mode,
main_backend,
params_backend_for(module),
params_mem_size);
params_mem_size,
false,
false,
&tensor_ops);
}
// Register graph-cut layer-split tensors on the primary backend first.
@@ -472,6 +488,7 @@ public:
SDBackendModule module,
const std::vector<ggml_backend_t>& module_backends,
std::map<std::string, ggml_tensor*> group_tensors,
const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
ModelManager::ResidencyMode residency_mode,
size_t* params_mem_size) {
bool has_cpu_device = false;
@@ -493,7 +510,10 @@ public:
residency_mode,
module_backends[0],
params_backend_for(module),
params_mem_size);
params_mem_size,
false,
false,
&tensor_ops);
}
model->set_runtime_backends(module_backends);
@@ -518,7 +538,8 @@ public:
initial_params_backend,
params_mem_size,
false,
params_follow_runtime);
params_follow_runtime,
&tensor_ops);
}
bool unload_control_net() {