llama-model: fix insonsistent ctxs <-> bufs order (#16581)

This commit is contained in:
Johannes Gäßler
2025-10-17 17:41:09 +02:00
committed by GitHub
parent 41386cf365
commit 66b0dbcb2d

View File

@@ -421,11 +421,8 @@ struct llama_model::impl {
llama_mlocks mlock_bufs; llama_mlocks mlock_bufs;
llama_mlocks mlock_mmaps; llama_mlocks mlock_mmaps;
// contexts where the model tensors metadata is stored // contexts where the model tensors metadata is stored as well ass the corresponding buffers:
std::vector<ggml_context_ptr> ctxs; std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
// the model memory buffers for the tensor data
std::vector<ggml_backend_buffer_ptr> bufs;
buft_list_t cpu_buft_list; buft_list_t cpu_buft_list;
std::map<ggml_backend_dev_t, buft_list_t> gpu_buft_list; std::map<ggml_backend_dev_t, buft_list_t> gpu_buft_list;
@@ -2182,7 +2179,14 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
max_n_tensors += n_layer*2; // duplicated rope freq tensors max_n_tensors += n_layer*2; // duplicated rope freq tensors
const size_t ctx_size = ggml_tensor_overhead()*max_n_tensors; const size_t ctx_size = ggml_tensor_overhead()*max_n_tensors;
std::map<ggml_backend_buffer_type_t, ggml_context *> ctx_map; // define a comparator for the buft -> ctx map to ensure that the order is well-defined:
struct ggml_backend_buft_comparator {
bool operator()(const ggml_backend_buffer_type_t & lhs, const ggml_backend_buffer_type_t & rhs) const {
return ggml_backend_buft_name(lhs) < ggml_backend_buft_name(rhs);
}
};
std::map<ggml_backend_buffer_type_t, ggml_context_ptr, ggml_backend_buft_comparator> ctx_map;
auto ctx_for_buft = [&](ggml_backend_buffer_type_t buft) -> ggml_context * { auto ctx_for_buft = [&](ggml_backend_buffer_type_t buft) -> ggml_context * {
auto it = ctx_map.find(buft); auto it = ctx_map.find(buft);
if (it == ctx_map.end()) { if (it == ctx_map.end()) {
@@ -2197,12 +2201,11 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
throw std::runtime_error(format("failed to create ggml context")); throw std::runtime_error(format("failed to create ggml context"));
} }
ctx_map[buft] = ctx; ctx_map.emplace(buft, ctx);
pimpl->ctxs.emplace_back(ctx);
return ctx; return ctx;
} }
return it->second; return it->second.get();
}; };
const auto TENSOR_DUPLICATED = llama_model_loader::TENSOR_DUPLICATED; const auto TENSOR_DUPLICATED = llama_model_loader::TENSOR_DUPLICATED;
@@ -6037,16 +6040,15 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
pimpl->mappings.reserve(ml.mappings.size()); pimpl->mappings.reserve(ml.mappings.size());
// create the backend buffers // create the backend buffers
std::vector<std::pair<ggml_context *, llama_buf_map>> ctx_bufs; std::vector<std::pair<ggml_context *, llama_buf_map>> ctx_buf_maps;
ctx_bufs.reserve(ctx_map.size()); ctx_buf_maps.reserve(ctx_map.size());
// Ensure we have enough capacity for the maximum backend buffer we will potentially create // Ensure we have enough capacity for the maximum backend buffer we will potentially create
const size_t n_max_backend_buffer = ctx_map.size() * ml.files.size(); const size_t n_max_backend_buffer = ctx_map.size() * ml.files.size();
pimpl->bufs.reserve(n_max_backend_buffer); pimpl->ctxs_bufs.reserve(n_max_backend_buffer);
for (auto & it : ctx_map) { for (auto & [buft, ctx_ptr] : ctx_map) {
ggml_backend_buffer_type_t buft = it.first; ggml_context * ctx = ctx_ptr.get();
ggml_context * ctx = it.second;
// skip contexts without tensors // skip contexts without tensors
if (ggml_get_first_tensor(ctx) == nullptr) { if (ggml_get_first_tensor(ctx) == nullptr) {
@@ -6070,6 +6072,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
bool buffer_from_host_ptr_supported = props.caps.buffer_from_host_ptr; bool buffer_from_host_ptr_supported = props.caps.buffer_from_host_ptr;
bool is_default_buft = buft == ggml_backend_dev_buffer_type(dev); bool is_default_buft = buft == ggml_backend_dev_buffer_type(dev);
ggml_backend_buffer_t buf = nullptr;
if (ml.use_mmap && use_mmap_buffer && buffer_from_host_ptr_supported && is_default_buft) { if (ml.use_mmap && use_mmap_buffer && buffer_from_host_ptr_supported && is_default_buft) {
for (uint32_t idx = 0; idx < ml.files.size(); idx++) { for (uint32_t idx = 0; idx < ml.files.size(); idx++) {
// only the mmap region containing the tensors in the model is mapped to the backend buffer // only the mmap region containing the tensors in the model is mapped to the backend buffer
@@ -6082,20 +6085,18 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
continue; continue;
} }
const size_t max_size = ggml_get_max_tensor_size(ctx); const size_t max_size = ggml_get_max_tensor_size(ctx);
ggml_backend_buffer_t buf = ggml_backend_dev_buffer_from_host_ptr(dev, (char *) addr + first, last - first, max_size); buf = ggml_backend_dev_buffer_from_host_ptr(dev, (char *) addr + first, last - first, max_size);
if (buf == nullptr) { if (buf == nullptr) {
throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft))); throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft)));
} }
pimpl->bufs.emplace_back(buf);
buf_map.emplace(idx, buf); buf_map.emplace(idx, buf);
} }
} }
else { else {
ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft); buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft);
if (buf == nullptr) { if (buf == nullptr) {
throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft))); throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft)));
} }
pimpl->bufs.emplace_back(buf);
if (use_mlock && ggml_backend_buffer_is_host(buf)) { if (use_mlock && ggml_backend_buffer_is_host(buf)) {
pimpl->mlock_bufs.emplace_back(new llama_mlock); pimpl->mlock_bufs.emplace_back(new llama_mlock);
auto & mlock_buf = pimpl->mlock_bufs.back(); auto & mlock_buf = pimpl->mlock_bufs.back();
@@ -6106,10 +6107,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
buf_map.emplace(idx, buf); buf_map.emplace(idx, buf);
} }
} }
pimpl->ctxs_bufs.emplace_back(std::move(ctx_ptr), buf);
if (pimpl->bufs.empty()) {
throw std::runtime_error("failed to allocate buffer");
}
for (auto & buf : buf_map) { for (auto & buf : buf_map) {
// indicate that this buffer contains weights // indicate that this buffer contains weights
@@ -6117,7 +6115,7 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
ggml_backend_buffer_set_usage(buf.second, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); ggml_backend_buffer_set_usage(buf.second, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
} }
ctx_bufs.emplace_back(ctx, buf_map); ctx_buf_maps.emplace_back(ctx, buf_map);
} }
if (llama_supports_gpu_offload()) { if (llama_supports_gpu_offload()) {
@@ -6135,22 +6133,20 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
} }
// print memory requirements per buffer type // print memory requirements per buffer type
for (auto & buf : pimpl->bufs) { for (auto & [_, buf] : pimpl->ctxs_bufs) {
LLAMA_LOG_INFO("%s: %12s model buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf.get()), ggml_backend_buffer_get_size(buf.get()) / 1024.0 / 1024.0); LLAMA_LOG_INFO("%s: %12s model buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf.get()), ggml_backend_buffer_get_size(buf.get()) / 1024.0 / 1024.0);
} }
// populate tensors_by_name // populate tensors_by_name
for (auto & ctx : pimpl->ctxs) { for (auto & [ctx, _] : pimpl->ctxs_bufs) {
for (auto * cur = ggml_get_first_tensor(ctx.get()); cur != NULL; cur = ggml_get_next_tensor(ctx.get(), cur)) { for (auto * cur = ggml_get_first_tensor(ctx.get()); cur != NULL; cur = ggml_get_next_tensor(ctx.get(), cur)) {
tensors_by_name.emplace_back(ggml_get_name(cur), cur); tensors_by_name.emplace_back(ggml_get_name(cur), cur);
} }
} }
// load tensor data // load tensor data
for (auto & it : ctx_bufs) { for (auto & [ctx, buf_map] : ctx_buf_maps) {
ggml_context * ctx = it.first; if (!ml.load_all_data(ctx, buf_map, use_mlock ? &pimpl->mlock_mmaps : NULL, params.progress_callback, params.progress_callback_user_data)) {
auto & bufs = it.second;
if (!ml.load_all_data(ctx, bufs, use_mlock ? &pimpl->mlock_mmaps : NULL, params.progress_callback, params.progress_callback_user_data)) {
return false; return false;
} }
} }
@@ -6190,8 +6186,8 @@ size_t llama_model::n_devices() const {
std::map<ggml_backend_buffer_type_t, size_t> llama_model::memory_breakdown() const { std::map<ggml_backend_buffer_type_t, size_t> llama_model::memory_breakdown() const {
std::map<ggml_backend_buffer_type_t, size_t> ret; std::map<ggml_backend_buffer_type_t, size_t> ret;
for (const ggml_backend_buffer_ptr & buf_ptr : pimpl->bufs) { for (const auto & [_, buf] : pimpl->ctxs_bufs) {
ret[ggml_backend_buffer_get_type(buf_ptr.get())] += ggml_backend_buffer_get_size(buf_ptr.get()); ret[ggml_backend_buffer_get_type(buf.get())] += ggml_backend_buffer_get_size(buf.get());
} }
return ret; return ret;
} }