fix: guard GPU memory capacity and propagate encoding failures

This commit is contained in:
leejet 2026-09-11 01:05:30 +08:00
parent 14eddb32b1
commit 58d3f6e240
5 changed files with 72 additions and 22 deletions

View File

@ -149,6 +149,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
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
used diffusion weights have priority. Each component's weights use the first
storage location with enough remaining budget:

View File

@ -1224,7 +1224,10 @@ struct FluxCLIPEmbedder : public Conditioner {
true,
clip_skip,
false);
GGML_ASSERT(!pooled.empty());
if (pooled.empty()) {
LOG_ERROR("Flux CLIP-L encoding failed");
return {};
}
} else {
pooled = sd::Tensor<float>::zeros({768});
}
@ -1243,7 +1246,10 @@ struct FluxCLIPEmbedder : public Conditioner {
input_ids,
sd::Tensor<float>(),
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);
if (zero_out_masked) {
chunk_hidden_states.fill_(0.0f);

View File

@ -1,4 +1,5 @@
#include <algorithm>
#include <exception>
#include <map>
#include <utility>
@ -624,7 +625,14 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
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();
if (success) {
cache_.graph_end(true);

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; };
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);
auto device = ggml_backend_get_device(request.compute_backend);
if (device != nullptr) {
auto available_device_bytes = [&](ggml_backend_t backend) {
auto device = ggml_backend_get_device(backend);
if (device == nullptr) {
return SIZE_MAX;
}
size_t free_bytes = 0, total_bytes = 0;
ggml_backend_dev_memory(device, &free_bytes, &total_bytes);
if (free_bytes != 0 || total_bytes != 0) {
result.available_device_bytes = free_bytes;
if (free_bytes == 0 && total_bytes == 0) {
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) {
const size_t resident = add(compute_backend_resident_bytes(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.
// Charge that upper bound on every participant instead of undercounting a shard.
for (const auto& entry : split_devices) {
size_t free_bytes = 0, total_bytes = 0;
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);
}
result.available_device_bytes = std::min(result.available_device_bytes, available_device_bytes(entry.first));
if (entry.second > 0) {
const size_t resident = add(compute_backend_resident_bytes(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);
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),
capacity.required_device_bytes / (1024.0 * 1024.0),
capacity.required_budget_bytes / (1024.0 * 1024.0),
capacity.available_device_bytes / (1024.0 * 1024.0),
capacity.available_budget_bytes / (1024.0 * 1024.0));
available_device.c_str(),
available_budget.c_str());
return false;
}

View File

@ -438,6 +438,10 @@ namespace sd::pipeline {
condition_params.zero_out_masked = false;
auto cond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
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) {
cond.c_concat = latents->concat_latent; // TODO: optimize
}
@ -466,6 +470,10 @@ namespace sd::pipeline {
condition_params.zero_out_masked = zero_out_masked;
uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
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) {
uncond.c_concat = latents->concat_latent; // TODO: optimize
@ -491,6 +499,10 @@ namespace sd::pipeline {
}
img_uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
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) {
img_uncond.c_concat = latents->img_uncond_concat_latent; // TODO: optimize
}