diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index d940b1b616..49f3c4f8ea 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1497,7 +1497,6 @@ bool llama_model_loader::load_all_data( } GGML_ASSERT(size_data != 0 && "call init_mappings() first"); - std::vector> read_buf; std::vector>> validation_result; // 4 staging buffers for async uploads, each sized 1MB seems to be a good default for single NVMe drives. @@ -1598,7 +1597,25 @@ bool llama_model_loader::load_all_data( ggml_backend_name(upload_backend)); } + std::vector tensors; for (struct ggml_tensor * cur = ggml_get_first_tensor(ctx); cur != NULL; cur = ggml_get_next_tensor(ctx, cur)) { + tensors.push_back(cur); + } + + // without mmap, tensors in non-host buffers are staged through a temporary buffer sized like the tensor + // load them biggest-first so the largest staging buffer is allocated while the fewest weights are resident + if (!use_mmap) { + std::stable_sort(tensors.begin(), tensors.end(), [](const ggml_tensor * a, const ggml_tensor * b) { + const bool staged_a = a->buffer && !ggml_backend_buffer_is_host(a->buffer); + const bool staged_b = b->buffer && !ggml_backend_buffer_is_host(b->buffer); + if (staged_a != staged_b) { + return staged_a; + } + return staged_a && ggml_nbytes(a) > ggml_nbytes(b); + }); + } + + for (struct ggml_tensor * cur : tensors) { const auto * weight = get_weight(ggml_get_name(cur)); if (weight == nullptr) { // this can happen with split experts models @@ -1711,7 +1728,8 @@ bool llama_model_loader::load_all_data( buffer_idx %= n_buffers; } } else { - read_buf.resize(n_size); + // scoped to one tensor so only one staging buffer is alive at a time + std::vector> read_buf(n_size); file->seek(weight->offs, SEEK_SET); file->read_raw(read_buf.data(), n_size); ggml_backend_tensor_set(cur, read_buf.data(), 0, n_size); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 408e76e322..6344f2d8ae 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1816,6 +1816,14 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { return true; } + // without mmap, load non-host buffers first: their tensors go through a staging buffer, which is cheapest while the fewest weights are resident + if (!ml.use_mmap) { + std::stable_partition(ctx_buf_maps.begin(), ctx_buf_maps.end(), [](const auto & ctx_buf_map) { + const auto & buf_map = ctx_buf_map.second; + return !buf_map.empty() && !ggml_backend_buffer_is_host(buf_map.begin()->second); + }); + } + // load tensor data for (auto & [ctx, buf_map] : ctx_buf_maps) { if (!ml.load_all_data(ctx, buf_map, use_mlock ? &pimpl->mlock_mmaps : NULL, params.progress_callback, params.progress_callback_user_data)) {