diff --git a/common/arg.cpp b/common/arg.cpp index dd92e4720b..d69dbe7169 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2720,6 +2720,19 @@ common_params_context common_params_parser_init(common_params & params, llama_ex else { throw std::invalid_argument("invalid value"); } } ).set_env("LLAMA_ARG_LOAD_MODE")); + add_opt(common_arg( + {"--tensor-read-lazy"}, "MODE", + "on-demand reading of tensors marked by the model arch (default: auto)\n" + "- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)\n" + "- auto: on, but only for tensors larger than 4 GiB\n" + "- off: always keep them resident", + [](common_params & params, const std::string & value) { + /**/ if (value == "on") { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_ON; } + else if (value == "auto") { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_AUTO; } + else if (value == "off") { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_OFF; } + else { throw std::invalid_argument("invalid value"); } + } + ).set_env("LLAMA_ARG_TENSOR_READ_LAZY")); add_opt(common_arg( {"--numa"}, "TYPE", "attempt optimizations that help on some NUMA systems\n" diff --git a/common/common.cpp b/common/common.cpp index 3d54bd6002..347e8e9fc4 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1688,6 +1688,7 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.main_gpu = params.main_gpu; mparams.split_mode = params.split_mode; mparams.load_mode = params.load_mode; + mparams.tensor_read_lazy = params.tensor_read_lazy; mparams.tensor_split = params.tensor_split; mparams.check_tensors = params.check_tensors; mparams.use_extra_bufts = !params.no_extra_bufts; diff --git a/common/common.h b/common/common.h index 9593e10d9a..b304b51728 100644 --- a/common/common.h +++ b/common/common.h @@ -476,6 +476,8 @@ struct common_params { enum llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER; // how to split the model across GPUs enum llama_load_mode load_mode = LLAMA_LOAD_MODE_AUTO; // how to load the model + enum llama_tensor_read_lazy tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_AUTO; // on-demand reading of tensors marked by the arch + common_cpu_params cpuparams; common_cpu_params cpuparams_batch; diff --git a/include/llama.h b/include/llama.h index 84313cfd74..b8020a7fce 100644 --- a/include/llama.h +++ b/include/llama.h @@ -214,6 +214,12 @@ extern "C" { LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode); LLAMA_API enum llama_load_mode llama_load_mode_from_str(const char * str); + enum llama_tensor_read_lazy { + LLAMA_TENSOR_READ_LAZY_OFF = 0, // always read the whole tensor up front + LLAMA_TENSOR_READ_LAZY_AUTO = 1, // lazy only for marked tensors larger than 4 GiB (requires mmap) + LLAMA_TENSOR_READ_LAZY_ON = 2, // read the rows of tensors marked by the arch on demand (requires mmap) + }; + enum llama_context_type { LLAMA_CONTEXT_TYPE_DEFAULT = 0, LLAMA_CONTEXT_TYPE_MTP = 1, @@ -315,6 +321,8 @@ extern "C" { enum llama_split_mode split_mode; // how to split the model across multiple GPUs enum llama_load_mode load_mode; // how to load the model + enum llama_tensor_read_lazy tensor_read_lazy; // on-demand reading of tensors marked by the arch + // the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE int32_t main_gpu; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index d4dd2b5d1a..ded2bca3cc 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1282,9 +1282,13 @@ struct ggml_tensor * llama_model_loader::create_tensor( return NULL; } - if ((flags & TENSOR_GET_ROW_LAZY) && use_mmap) { - const auto & w = require_weight(tn.str().c_str()); - lazy_tensor_ranges[w.idx].emplace_back(w.offs, w.offs + ggml_nbytes(cur)); + if ((flags & TENSOR_GET_ROW_LAZY) && use_mmap && tensor_read_lazy != LLAMA_TENSOR_READ_LAZY_OFF) { + // in auto mode, small tensors are cheap enough to keep resident + constexpr size_t auto_lazy_min_size = 4ull * 1024 * 1024 * 1024; + if (tensor_read_lazy == LLAMA_TENSOR_READ_LAZY_ON || ggml_nbytes(cur) > auto_lazy_min_size) { + const auto & w = require_weight(tn.str().c_str()); + lazy_tensor_ranges[w.idx].emplace_back(w.offs, w.offs + ggml_nbytes(cur)); + } } ggml_tensor t_meta = *cur; diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 7a4e608769..5b78ece484 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -83,6 +83,9 @@ struct llama_model_loader { bool no_alloc; bool load_mtp; + // set by the caller before the create_tensor() calls + enum llama_tensor_read_lazy tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_OFF; + llama_files files; llama_ftype ftype; llama_fver fver; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 47aae1b72a..578b44cea6 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2631,6 +2631,7 @@ llama_model_params llama_model_default_params() { /*.n_gpu_layers =*/ -1, /*.split_mode =*/ LLAMA_SPLIT_MODE_LAYER, /*.load_mode =*/ LLAMA_LOAD_MODE_AUTO, + /*.tensor_read_lazy =*/ LLAMA_TENSOR_READ_LAZY_AUTO, /*.main_gpu =*/ 0, /*.tensor_split =*/ nullptr, /*.progress_callback =*/ nullptr, diff --git a/src/llama.cpp b/src/llama.cpp index 1609fec88d..9c841ee352 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -318,6 +318,8 @@ static std::pair llama_model_load(struct gguf_context * meta llama_model_loader ml(metadata, set_tensor_data, set_tensor_data_ud, fname, splits, file, params.load_mode, params.check_tensors, params.no_alloc, params.load_mtp, params.kv_overrides, params.tensor_buft_overrides); + ml.tensor_read_lazy = params.tensor_read_lazy; + ml.print_info(); std::unique_ptr model_ptr(llama_model_create(ml, params));