args: refactor mlock/mmap/directio into load-mode (#20834)

* args: overhaul mmap/mlock/dio into single arg

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* docs: update docs with llama-gen-docs

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* chore: satisfy code quality

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* args: make the `+` sign an actual modifier now

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* chore: general code clean up + comments

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* arg: fix deprecated flags support

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* arg: quick sanity check

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* bench: sync llama-bench argument parsing

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* fix: bugfix variable behaviour + llama-bench lm column size

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* arg: inverse commands should do the opposite instead of doing nothing

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* bench: fix incorrect dash

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* bench: fix missing modifiers for deprecated flags

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* llama: switch back to thread_local

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* arg: switch back to single enum

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* docs: update arg docs

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* chore: fix missing `mlock` from llama_load_mode_from_str + cleanup llama-bench

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* llama: fix mlock not activating

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* arg: add deprecation warning when old and new flags are combined

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* arg: cont add comment for todo in the future

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* docs: sync with upstream

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

* docs: re-sync with upstream again

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>

---------

Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>
This commit is contained in:
Aaron Teo
2026-07-23 20:32:56 +08:00
committed by GitHub
parent da296d6e72
commit e6dd0e29a6
20 changed files with 315 additions and 207 deletions
+7 -20
View File
@@ -4,6 +4,7 @@
#include "ggml.h"
#include "gguf.h"
#include "llama-hparams.h"
#include "llama.h"
#include <algorithm>
#include <array>
@@ -522,8 +523,7 @@ llama_model_loader::llama_model_loader(
const std::string & fname,
std::vector<std::string> & splits,
FILE * file,
bool use_mmap,
bool use_direct_io,
llama_load_mode load_mode,
bool check_tensors,
bool no_alloc,
const llama_model_kv_override * param_overrides_p,
@@ -542,6 +542,9 @@ llama_model_loader::llama_model_loader(
tensor_buft_overrides = param_tensor_buft_overrides_p;
this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MLOCK;
this->use_direct_io = load_mode == LLAMA_LOAD_MODE_DIRECT_IO;
if (!fname.empty()) {
// Load the main GGUF
struct ggml_context * ctx = NULL;
@@ -562,20 +565,6 @@ llama_model_loader::llama_model_loader(
files.emplace_back(new llama_file(fname.c_str(), "rb", use_direct_io));
contexts.emplace_back(ctx);
if (use_mmap && use_direct_io) {
if (files.back()->has_direct_io()) {
LLAMA_LOG_WARN("%s: direct I/O is enabled, disabling mmap\n", __func__);
use_mmap = false;
} else {
LLAMA_LOG_WARN("%s: direct I/O is not available, using mmap\n", __func__);
use_direct_io = false;
// reopen file using std::fopen for mmap
files.pop_back();
files.emplace_back(new llama_file(fname.c_str(), "rb", false));
}
}
// Save tensors data offset of the main file.
// For subsidiary files, `meta` tensor data offset must not be used,
// so we build a unified tensors index for weights.
@@ -816,13 +805,11 @@ llama_model_loader::llama_model_loader(
}
}
if (!llama_mmap::SUPPORTED) {
if (this->use_mmap && !llama_mmap::SUPPORTED) {
LLAMA_LOG_WARN("%s: mmap is not supported on this platform\n", __func__);
use_mmap = false;
this->use_mmap = false;
}
this->use_mmap = use_mmap;
this->use_direct_io = use_direct_io;
this->check_tensors = check_tensors;
this->no_alloc = no_alloc;
}
+1 -2
View File
@@ -126,8 +126,7 @@ struct llama_model_loader {
const std::string & fname,
std::vector<std::string> & splits, // optional, only need if the split does not follow naming scheme
FILE * file,
bool use_mmap,
bool use_direct_io,
llama_load_mode load_mode,
bool check_tensors,
bool no_alloc,
const llama_model_kv_override * param_overrides_p,
+5 -6
View File
@@ -16,6 +16,7 @@
#include "llama-memory-hybrid-iswa.h"
#include "llama-memory-recurrent.h"
#include "llama.h"
#include "models/models.h"
#include "ggml.h"
@@ -1243,7 +1244,7 @@ void llama_model_base::load_vocab(llama_model_loader & ml) {
bool llama_model_base::load_tensors(llama_model_loader & ml) {
const auto & split_mode = params.split_mode;
const auto & use_mlock = params.use_mlock;
const bool use_mlock = params.load_mode == LLAMA_LOAD_MODE_MLOCK;
const auto & tensor_split = params.tensor_split;
const int n_layer_all = hparams.n_layer_all;
@@ -1253,8 +1254,8 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
this->ml = &ml; // to be used by create_tensor() and load_arch_tensors()
LLAMA_LOG_INFO("%s: loading model tensors, this can take a while... (mmap = %s, direct_io = %s)\n",
__func__, ml.use_mmap ? "true" : "false", ml.use_direct_io ? "true" : "false");
LLAMA_LOG_INFO("%s: loading model tensors, this can take a while... (load_mode = %s)\n",
__func__, llama_load_mode_name(params.load_mode));
// build a list of buffer types for the CPU and GPU devices
pimpl->cpu_buft_list = make_cpu_buft_list(devices, params.use_extra_bufts, params.no_host);
@@ -2318,15 +2319,13 @@ llama_model_params llama_model_default_params() {
/*.tensor_buft_overrides =*/ nullptr,
/*.n_gpu_layers =*/ -1,
/*.split_mode =*/ LLAMA_SPLIT_MODE_LAYER,
/*.load_mode =*/ LLAMA_LOAD_MODE_MMAP,
/*.main_gpu =*/ 0,
/*.tensor_split =*/ nullptr,
/*.progress_callback =*/ nullptr,
/*.progress_callback_user_data =*/ nullptr,
/*.kv_overrides =*/ nullptr,
/*.vocab_only =*/ false,
/*.use_mmap =*/ true,
/*.use_direct_io =*/ false,
/*.use_mlock =*/ false,
/*.check_tensors =*/ false,
/*.use_extra_bufts =*/ true,
/*.no_host =*/ false,
+4 -3
View File
@@ -2,6 +2,7 @@
#include "llama-model.h"
#include "llama-model-loader.h"
#include "llama-ext.h"
#include "llama.h"
#include <algorithm>
#include <cmath>
@@ -876,15 +877,15 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
// mmap consistently increases speed on Linux, and also increases speed on Windows with
// hot cache. It may cause a slowdown on macOS, possibly related to free memory.
#if defined(__linux__) || defined(_WIN32)
constexpr bool use_mmap = true;
constexpr llama_load_mode load_mode = LLAMA_LOAD_MODE_MMAP;
#else
constexpr bool use_mmap = false;
constexpr llama_load_mode load_mode = LLAMA_LOAD_MODE_NONE;
#endif
const llama_model_kv_override * kv_overrides = params->kv_overrides;
std::vector<std::string> splits = {};
llama_model_loader ml(/*metadata*/ nullptr, /*set_tensor_data*/ nullptr, /*set_tensor_data_ud*/ nullptr,
fname_inp, splits, /*file*/ nullptr, use_mmap, /*use_direct_io*/ false, /*check_tensors*/ true, /*no_alloc*/ false, kv_overrides, nullptr);
fname_inp, splits, /*file*/ nullptr, /*load_mode*/ load_mode, /*check_tensors*/ true, /*no_alloc*/ false, kv_overrides, nullptr);
ml.init_mappings(false); // no prefetching
auto mparams = llama_model_default_params();
+24 -2
View File
@@ -46,6 +46,28 @@ const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_ty
GGML_ABORT("fatal error");
}
const char * llama_load_mode_name(enum llama_load_mode load_mode) {
switch (load_mode) {
case LLAMA_LOAD_MODE_NONE:
return "none";
case LLAMA_LOAD_MODE_MMAP:
return "mmap";
case LLAMA_LOAD_MODE_MLOCK:
return "mlock";
case LLAMA_LOAD_MODE_DIRECT_IO:
return "dio";
}
GGML_ABORT("fatal error");
}
enum llama_load_mode llama_load_mode_from_str(const char * str) {
if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; }
if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; }
if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; }
if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; }
throw std::invalid_argument(std::string("unknown load mode: ") + str);
}
struct llama_sampler_chain_params llama_sampler_chain_default_params() {
struct llama_sampler_chain_params result = {
/*.no_perf =*/ true,
@@ -279,7 +301,7 @@ static bool llama_prepare_model_devices(const llama_model_params & params, llama
static std::pair<int, llama_model *> llama_model_load(struct gguf_context * metadata, llama_model_set_tensor_data_t set_tensor_data, void * set_tensor_data_ud,
const std::string & fname, std::vector<std::string> & splits, FILE * file, llama_model_params & params) {
try {
llama_model_loader ml(metadata, set_tensor_data, set_tensor_data_ud, fname, splits, file, params.use_mmap, params.use_direct_io,
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.kv_overrides, params.tensor_buft_overrides);
ml.print_info();
@@ -412,7 +434,7 @@ struct llama_model * llama_model_init_from_user(
GGML_ASSERT(metadata != nullptr);
std::string path_model;
std::vector<std::string> splits = {};
params.use_mmap = false;
params.load_mode = LLAMA_LOAD_MODE_NONE;
params.use_extra_bufts = false;
return llama_model_load_from_file_impl(metadata, set_tensor_data, set_tensor_data_ud, path_model, splits, /*file*/ nullptr, params);
}