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

feat: drive layer split from graph-cut segments (#1762)

Browse Source
This commit is contained in:
leejet
2026-07-07 23:16:52 +08:00
committed by GitHub
parent bb84971129
commit 9ef6e7398f
8 changed files with 615 additions and 210 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/core/layer_split_partition.cpp
+172 -136
View File
@@ -1,9 +1,11 @@
#include "core/layer_split_partition.h"
#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <limits>
#include <unordered_set>
#include <utility>
#include "core/util.h"
@@ -62,160 +64,194 @@ namespace sd {
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 graph_cut_layer_split_backend_vram_limit(const std::vector<size_t>& backend_vram_limits,
size_t backend_index,
size_t primary_backend_vram_limit) {
if (backend_index < backend_vram_limits.size()) {
return backend_vram_limits[backend_index];
}
return backend_index == 0 ? primary_backend_vram_limit : 0;
}
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.
static std::vector<int64_t> graph_cut_layer_split_backend_capacities(const std::vector<ggml_backend_t>& backends,
const std::vector<size_t>& backend_vram_limits,
size_t primary_backend_vram_limit) {
std::vector<int64_t> capacities(backends.size(), std::numeric_limits<int64_t>::max() / 4);
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];
if (free_bytes > 0) {
capacities[i] = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes, 0);
}
size_t limit_bytes = graph_cut_layer_split_backend_vram_limit(backend_vram_limits,
i,
primary_backend_vram_limit);
if (limit_bytes > 0) {
capacities[i] = std::min<int64_t>(capacities[i], (int64_t)limit_bytes);
}
}
return capacities;
}
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;
}
bool partition_graph_cut_layer_split(const char* desc,
ggml_cgraph* gf,
const sd::ggml_graph_cut::Plan& plan,
const std::vector<ggml_backend_t>& split_backends,
const std::vector<size_t>& backend_vram_limits,
size_t primary_backend_vram_limit,
std::unordered_map<const ggml_tensor*, ggml_backend_t>& param_assignments,
const std::function<ggml_tensor*(ggml_tensor*)>& canonical_param_tensor,
GraphCutLayerSplitAssignment* assignment_out) {
GGML_ASSERT(gf != nullptr);
GGML_ASSERT(assignment_out != nullptr);
GGML_ASSERT(canonical_param_tensor != nullptr);
GGML_ASSERT(!split_backends.empty());
GraphCutLayerSplitAssignment assignment;
assignment.segment_count = plan.segments.size();
assignment.tensors_by_backend.resize(split_backends.size());
assignment.bytes_by_backend.resize(split_backends.size(), 0);
assignment.first_segment_by_backend.resize(split_backends.size(), plan.segments.size());
assignment.last_segment_by_backend.resize(split_backends.size(), 0);
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;
std::vector<std::vector<ggml_tensor*>> segment_params(plan.segments.size());
std::vector<int64_t> segment_param_bytes(plan.segments.size(), 0);
std::unordered_set<ggml_tensor*> seen_params;
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
std::vector<ggml_tensor*> params = sd::ggml_graph_cut::param_tensors(gf, plan.segments[seg_idx]);
for (ggml_tensor* raw_param : params) {
ggml_tensor* param = canonical_param_tensor(raw_param);
if (param == nullptr || !seen_params.insert(param).second) {
continue;
}
segment_params[seg_idx].push_back(param);
segment_param_bytes[seg_idx] += (int64_t)ggml_nbytes(param);
}
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++;
int64_t total_param_bytes = 0;
for (int64_t bytes : segment_param_bytes) {
total_param_bytes += bytes;
}
if (total_param_bytes <= 0) {
LOG_ERROR("%s graph-cut layer split found no graph params to assign", desc);
return false;
}
std::vector<int64_t> backend_capacities = graph_cut_layer_split_backend_capacities(split_backends,
backend_vram_limits,
primary_backend_vram_limit);
std::vector<ggml_backend_t> backend_by_segment(plan.segments.size(), split_backends[0]);
size_t current_backend = 0;
int64_t current_used = 0;
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
int64_t bytes = segment_param_bytes[seg_idx];
while (current_backend + 1 < split_backends.size() &&
bytes > 0 &&
current_used + bytes > backend_capacities[current_backend]) {
current_backend++;
current_used = 0;
}
if (bytes > 0 && current_used + bytes > backend_capacities[current_backend]) {
LOG_ERROR("%s graph-cut layer split: segment %zu needs %.1f MB on %s, but only %.1f MB is available under current VRAM limits",
desc,
seg_idx,
(current_used + bytes) / (1024.0 * 1024.0),
layer_split_backend_device_display_name(split_backends[current_backend]).c_str(),
backend_capacities[current_backend] / (1024.0 * 1024.0));
return false;
}
current_used += bytes;
backend_by_segment[seg_idx] = split_backends[current_backend];
for (ggml_tensor* param : segment_params[seg_idx]) {
ggml_backend_t target_backend = split_backends[current_backend];
auto assigned_it = param_assignments.find(param);
if (assigned_it == param_assignments.end()) {
param_assignments[param] = target_backend;
assignment.has_new_param_assignment = true;
} else {
target_backend = assigned_it->second;
}
target = std::min(target, backends.size() - 1);
target = layer_split_supported_target(desc, kv.first, kv.second, backends, target);
auto backend_it = std::find(split_backends.begin(), split_backends.end(), target_backend);
if (backend_it == split_backends.end()) {
LOG_ERROR("%s graph-cut layer split tensor '%s' is assigned to an unavailable backend",
desc,
ggml_get_name(param));
return false;
}
size_t backend_idx = (size_t)std::distance(split_backends.begin(), backend_it);
assignment.first_segment_by_backend[backend_idx] = std::min(assignment.first_segment_by_backend[backend_idx], seg_idx);
assignment.last_segment_by_backend[backend_idx] = std::max(assignment.last_segment_by_backend[backend_idx], seg_idx + 1);
assignment.tensors_by_backend[backend_idx].push_back(param);
assignment.bytes_by_backend[backend_idx] += (int64_t)ggml_nbytes(param);
}
}
const int n_nodes = ggml_graph_n_nodes(gf);
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
ggml_backend_t backend = backend_by_segment[seg_idx];
const auto& segment = plan.segments[seg_idx];
for (int node_index : segment.internal_node_indices) {
if (node_index < 0 || node_index >= n_nodes) {
continue;
}
ggml_tensor* node = ggml_graph_node(gf, node_index);
if (node != nullptr) {
assignment.node_assignments[node] = backend;
}
}
for (int node_index : segment.output_node_indices) {
if (node_index < 0 || node_index >= n_nodes) {
continue;
}
ggml_tensor* node = ggml_graph_node(gf, node_index);
if (node != nullptr) {
assignment.node_assignments[node] = backend;
}
}
}
*assignment_out = std::move(assignment);
return true;
}
void log_graph_cut_layer_split_assignment(const char* desc,
const std::vector<ggml_backend_t>& split_backends,
const GraphCutLayerSplitAssignment& assignment) {
for (size_t i = 0; i < split_backends.size(); i++) {
if (i >= assignment.tensors_by_backend.size() ||
assignment.tensors_by_backend[i].empty()) {
continue;
}
size_t first_segment = assignment.first_segment_by_backend[i] == assignment.segment_count
? 0
: assignment.first_segment_by_backend[i];
size_t last_segment = assignment.last_segment_by_backend[i];
if (assignment.has_new_param_assignment) {
LOG_INFO("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
desc,
layer_split_backend_device_display_name(split_backends[i]).c_str(),
first_segment,
last_segment,
assignment.tensors_by_backend[i].size(),
assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
} else {
auto target_it = non_block_targets.find(kv.first);
if (target_it != non_block_targets.end()) {
target = target_it->second;
}
LOG_DEBUG("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
desc,
layer_split_backend_device_display_name(split_backends[i]).c_str(),
first_segment,
last_segment,
assignment.tensors_by_backend[i].size(),
assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
}
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