fix: guard GPU memory capacity and propagate encoding failures (#1958)

This commit is contained in:
leejet 2026-09-13 21:35:46 +08:00 committed by GitHub
parent 44dd13716d
commit 9a977388a8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 104 additions and 29 deletions

View File

@ -294,6 +294,8 @@ endif()
if(MSVC) if(MSVC)
target_compile_options(${SD_LIB} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/bigobj>) target_compile_options(${SD_LIB} PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/bigobj>)
# ggml backends can throw C++ exceptions through their C API.
target_compile_options(${SD_LIB} PRIVATE $<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CXX_COMPILER_ID:MSVC>>:/EHsc->)
endif() endif()
if(APPLE) if(APPLE)

View File

@ -154,6 +154,11 @@ GiB", and with no budget set each device's free memory minus a 512 MiB margin
is used. These resolved GPU budgets, including the safety margin, also drive is used. These resolved GPU budgets, including the safety margin, also drive
the runner's graph-cut capacity checks. the runner's graph-cut capacity checks.
Runtime capacity checks also leave 512 MiB of currently free device memory for
backend scratch buffers and pipelines, including with explicit backend assignments.
They cap stale free-memory reports by the device's total memory minus tracked
resident allocations and reject reports that exceed the device's total memory.
Components are considered in `diffusion`, `te`, `vae` order so that repeatedly Components are considered in `diffusion`, `te`, `vae` order so that repeatedly
used diffusion weights have priority. Each component's weights use the first used diffusion weights have priority. Each component's weights use the first
storage location with enough remaining budget: storage location with enough remaining budget:

View File

@ -1254,7 +1254,10 @@ struct FluxCLIPEmbedder : public Conditioner {
true, true,
clip_skip, clip_skip,
false); false);
GGML_ASSERT(!pooled.empty()); if (pooled.empty()) {
LOG_ERROR("Flux CLIP-L encoding failed");
return {};
}
} else { } else {
pooled = sd::Tensor<float>::zeros({768}); pooled = sd::Tensor<float>::zeros({768});
} }
@ -1273,7 +1276,10 @@ struct FluxCLIPEmbedder : public Conditioner {
input_ids, input_ids,
sd::Tensor<float>(), sd::Tensor<float>(),
false); false);
GGML_ASSERT(!chunk_hidden_states.empty()); if (chunk_hidden_states.empty()) {
LOG_ERROR("Flux T5 encoding failed at chunk %d/%zu", chunk_idx + 1, chunk_count);
return {};
}
chunk_hidden_states = ::apply_token_weights(std::move(chunk_hidden_states), chunk_weights); chunk_hidden_states = ::apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
if (zero_out_masked) { if (zero_out_masked) {
chunk_hidden_states.fill_(0.0f); chunk_hidden_states.fill_(0.0f);

View File

@ -2,12 +2,14 @@
#include <algorithm> #include <algorithm>
#include <cstring> #include <cstring>
#include <exception>
#include <map> #include <map>
#include <unordered_map> #include <unordered_map>
#include <unordered_set> #include <unordered_set>
#include "core/ggml_extend_backend.h" #include "core/ggml_extend_backend.h"
#include "core/ggml_graph_cut.h" #include "core/ggml_graph_cut.h"
#include "core/util.h"
#include "ggml-cpu.h" #include "ggml-cpu.h"
#include "ggml/src/ggml-impl.h" #include "ggml/src/ggml-impl.h"
@ -228,11 +230,23 @@ namespace sd {
} }
} }
void ComputeWorkspace::segment_end() { bool ComputeWorkspace::segment_end() noexcept {
if (active_) { if (!active_) {
synchronize(); return true;
active_ = false;
} }
// Outer cleanup guards must not retry a failed backend submission.
active_ = false;
try {
synchronize();
return true;
} catch (const std::exception& error) {
LOG_ERROR("%s workspace synchronization failed during segment cleanup: %s",
ggml_backend_name(backend_), error.what());
} catch (...) {
LOG_ERROR("%s workspace synchronization failed during segment cleanup: unknown exception",
ggml_backend_name(backend_));
}
return false;
} }
bool ComputeWorkspace::release() { bool ComputeWorkspace::release() {

View File

@ -51,7 +51,7 @@ namespace sd {
const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend, const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
const AssignNodes& assign_nodes); const AssignNodes& assign_nodes);
void synchronize() const; void synchronize() const;
void segment_end(); bool segment_end() noexcept;
bool release(); bool release();
bool active() const { return active_; } bool active() const { return active_; }
ggml_backend_sched_t scheduler() const { return scheduler_; } ggml_backend_sched_t scheduler() const { return scheduler_; }

View File

@ -660,7 +660,13 @@ void SDBackendAssignment::set_module(SDBackendModule module, const std::string&
} }
void SDBackendHandleDeleter::operator()(ggml_backend_t backend) const { void SDBackendHandleDeleter::operator()(ggml_backend_t backend) const {
try {
ggml_backend_free(backend); ggml_backend_free(backend);
} catch (const std::exception& error) {
LOG_ERROR("backend cleanup failed: %s", error.what());
} catch (...) {
LOG_ERROR("backend cleanup failed: unknown exception");
}
} }
SDBackendManager::~SDBackendManager() { SDBackendManager::~SDBackendManager() {

View File

@ -1,4 +1,5 @@
#include <algorithm> #include <algorithm>
#include <exception>
#include <map> #include <map>
#include <utility> #include <utility>
@ -642,7 +643,14 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
params_tensor_set_.insert(parameter); params_tensor_set_.insert(parameter);
} }
} }
auto output = execute_graph(graph, n_threads, no_return, read_outputs); std::optional<sd::Tensor<float>> output;
try {
output = execute_graph(graph, n_threads, no_return, read_outputs);
} catch (const std::exception& error) {
LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(),
ggml_backend_name(runtime_backend), error.what());
return std::nullopt;
}
success = output.has_value(); success = output.has_value();
if (success) { if (success) {
cache_.graph_end(true); cache_.graph_end(true);
@ -956,6 +964,9 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
} }
} }
} }
if (!workspace_.segment_end()) {
return fail_segment("workspace synchronization");
}
// Final outputs and their callbacks may still be views of consumed cuts. // Final outputs and their callbacks may still be views of consumed cuts.
cut_cache_.prune(segment.future_cut_names); cut_cache_.prune(segment.future_cut_names);
} }

View File

@ -1586,16 +1586,33 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
} }
auto add = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; }; auto add = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; };
const size_t missing = compute_backend_alloc_size(states, true); const size_t missing = compute_backend_alloc_size(states, true);
result.required_device_bytes = add(request.pending_allocation_bytes, missing); // Backend scratch buffers and pipelines are not included in graph measurements.
constexpr size_t safety_margin = 512ULL * 1024ULL * 1024ULL;
result.required_device_bytes = add(add(request.pending_allocation_bytes, missing), safety_margin);
result.required_budget_bytes = add(request.runtime_peak_bytes(), missing); result.required_budget_bytes = add(request.runtime_peak_bytes(), missing);
auto device = ggml_backend_get_device(request.compute_backend); auto available_device_bytes = [&](ggml_backend_t backend) {
if (device != nullptr) { auto device = ggml_backend_get_device(backend);
if (device == nullptr) {
return SIZE_MAX;
}
size_t free_bytes = 0, total_bytes = 0; size_t free_bytes = 0, total_bytes = 0;
ggml_backend_dev_memory(device, &free_bytes, &total_bytes); ggml_backend_dev_memory(device, &free_bytes, &total_bytes);
if (free_bytes != 0 || total_bytes != 0) { if (free_bytes == 0 && total_bytes == 0) {
result.available_device_bytes = free_bytes; return SIZE_MAX;
} }
// Vulkan's heap budget subtraction can underflow when usage exceeds the budget.
if (total_bytes > 0 && free_bytes > total_bytes) {
return size_t{0};
} }
const size_t resident = add(compute_backend_resident_bytes(backend),
add(other_runtime_resident_bytes(request.owner_id, backend),
request.runtime_resident_bytes));
if (total_bytes > 0) {
free_bytes = std::min(free_bytes, resident < total_bytes ? total_bytes - resident : 0);
}
return free_bytes;
};
result.available_device_bytes = available_device_bytes(request.compute_backend);
if (request.max_backend_bytes > 0) { if (request.max_backend_bytes > 0) {
const size_t resident = add(compute_backend_resident_bytes(request.compute_backend), const size_t resident = add(compute_backend_resident_bytes(request.compute_backend),
other_runtime_resident_bytes(request.owner_id, request.compute_backend)); other_runtime_resident_bytes(request.owner_id, request.compute_backend));
@ -1619,11 +1636,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
// GGML exposes only a split buffer's total size, not per-device allocations. // GGML exposes only a split buffer's total size, not per-device allocations.
// Charge that upper bound on every participant instead of undercounting a shard. // Charge that upper bound on every participant instead of undercounting a shard.
for (const auto& entry : split_devices) { for (const auto& entry : split_devices) {
size_t free_bytes = 0, total_bytes = 0; result.available_device_bytes = std::min(result.available_device_bytes, available_device_bytes(entry.first));
ggml_backend_dev_memory(ggml_backend_get_device(entry.first), &free_bytes, &total_bytes);
if (free_bytes != 0 || total_bytes != 0) {
result.available_device_bytes = std::min(result.available_device_bytes, free_bytes);
}
if (entry.second > 0) { if (entry.second > 0) {
const size_t resident = add(compute_backend_resident_bytes(entry.first), const size_t resident = add(compute_backend_resident_bytes(entry.first),
other_runtime_resident_bytes(request.owner_id, entry.first)); other_runtime_resident_bytes(request.owner_id, entry.first));
@ -1740,11 +1753,17 @@ bool ModelManager::ensure_compute_backend_capacity(
} }
const auto capacity = check_capacity(request, required_states); const auto capacity = check_capacity(request, required_states);
LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %.2f MB device / %.2f MB budget", const std::string available_device = capacity.available_device_bytes == SIZE_MAX
? "unknown"
: sd_format("%.2f MB", capacity.available_device_bytes / (1024.0 * 1024.0));
const std::string available_budget = capacity.available_budget_bytes == SIZE_MAX
? "unlimited"
: sd_format("%.2f MB", capacity.available_budget_bytes / (1024.0 * 1024.0));
LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %s device / %s budget",
ggml_backend_name(compute_backend), ggml_backend_name(compute_backend),
capacity.required_device_bytes / (1024.0 * 1024.0), capacity.required_device_bytes / (1024.0 * 1024.0),
capacity.required_budget_bytes / (1024.0 * 1024.0), capacity.required_budget_bytes / (1024.0 * 1024.0),
capacity.available_device_bytes / (1024.0 * 1024.0), available_device.c_str(),
capacity.available_budget_bytes / (1024.0 * 1024.0)); available_budget.c_str());
return false; return false;
} }

View File

@ -31,7 +31,7 @@ public:
}; };
private: private:
static constexpr size_t MAX_RESIDENCY_BLOCK_BYTES = 64ULL * 1024ULL * 1024ULL; static constexpr size_t MAX_RESIDENCY_BLOCK_BYTES = 1024ULL * 1024ULL * 1024ULL;
struct TensorState { struct TensorState {
std::string name; std::string name;

View File

@ -438,6 +438,10 @@ namespace sd::pipeline {
condition_params.zero_out_masked = false; condition_params.zero_out_masked = false;
auto cond = sd->cond_stage_model->get_learned_condition(sd->n_threads, auto cond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
condition_params); condition_params);
if (cond.empty()) {
LOG_ERROR("failed to encode prompt");
return std::nullopt;
}
if (cond.c_concat.empty() && ref_image_params.pass_to_dit) { if (cond.c_concat.empty() && ref_image_params.pass_to_dit) {
cond.c_concat = latents->concat_latent; // TODO: optimize cond.c_concat = latents->concat_latent; // TODO: optimize
} }
@ -469,6 +473,10 @@ namespace sd::pipeline {
condition_params.zero_out_masked = zero_out_masked; condition_params.zero_out_masked = zero_out_masked;
uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
condition_params); condition_params);
if (uncond.empty()) {
LOG_ERROR("failed to encode negative prompt");
return std::nullopt;
}
} }
if (uncond.c_concat.empty() && ref_image_params.pass_to_dit) { if (uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
uncond.c_concat = latents->concat_latent; // TODO: optimize uncond.c_concat = latents->concat_latent; // TODO: optimize
@ -494,6 +502,10 @@ namespace sd::pipeline {
} }
img_uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, img_uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
condition_params); condition_params);
if (img_uncond.empty()) {
LOG_ERROR("failed to encode image guidance prompt");
return std::nullopt;
}
if (img_uncond.c_concat.empty() && ref_image_params.pass_to_dit) { if (img_uncond.c_concat.empty() && ref_image_params.pass_to_dit) {
img_uncond.c_concat = latents->img_uncond_concat_latent; // TODO: optimize img_uncond.c_concat = latents->img_uncond_concat_latent; // TODO: optimize
} }