From 546369630ff32a3589bffdff803d2e4a69502353 Mon Sep 17 00:00:00 2001 From: Piotr Wilkin Date: Wed, 22 Jul 2026 16:15:55 +0200 Subject: [PATCH] llama : add --lazy-experts for MoE models larger than RAM Routed MoE expert tensors are the bulk of a large MoE checkpoint but only a few are touched per token. --lazy-experts maps the model without populating them: only the non-expert regions are prefetched (posix_madvise WILLNEED) and experts fault in from the page cache the first time they are routed to. This lets a model whose experts do not fit in RAM run off the page cache, which the OS then manages as an LRU of the hot expert set. On a 213 GB GLM-5.2 IQ1_M with 123 GB of RAM this is the difference between not running and running: ~198 GB of experts are left on demand and ~14.5 GB of non-expert weight is prefetched. Readahead is deliberately left enabled on the expert regions -- when the working set exceeds RAM the load is disk-bound, and a fault that pulls a larger contiguous chunk of the expert row is a win. LLAMA_LAZY_EXPERT_RANDOM forces MADV_RANDOM instead, which avoids dragging in neighbours and can be preferable when the model does fit. Two fixes fall out of this, both independent of the new flag: * llama_mmap::advise_range() aligns its range before calling (posix_)madvise. These require a page-aligned start, and tensor offsets are not aligned, so an unaligned call fails with EINVAL and the hint is silently dropped. Advisory hints round outward so the whole range is still covered; DONTNEED rounds inward so a neighbouring tensor's pages are never dropped. The error is also now reported correctly -- posix_madvise() returns the error number and does not set errno, so strerror(errno) printed an unrelated message. * When a tensor is uploaded from the mmap into a device buffer its mapped pages are dead weight. On a model larger than RAM they evict pages that are still needed and thrash the load, so they are dropped once the upload completes. Deliberately not included, having measured no benefit on this workload: pinning hot experts (neutral below ~24 GB of budget, and progressively worse above it, because pinned pages are not reclaimable and displace the page cache the rest of the model depends on); an explicit LRU/eviction cache with on-demand pread() for non-mmap loads (ties mmap warm, and its budget caps the cache below what mmap gets for free); and speculative prefetch of routed experts (on a working set larger than the cache it roughly doubles bytes read per token). The page cache is already the right policy here; the useful change is only to stop fighting it. Assisted-by: Claude Opus 4.8 --- common/arg.cpp | 9 +++++ common/common.cpp | 1 + common/common.h | 1 + include/llama.h | 2 + src/llama-mmap.cpp | 79 ++++++++++++++++++++++++++++++++++++++ src/llama-mmap.h | 10 +++++ src/llama-model-loader.cpp | 56 +++++++++++++++++++++++++-- src/llama-model-loader.h | 4 +- src/llama-model.cpp | 3 +- 9 files changed, 160 insertions(+), 5 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 1f5b5a8abfc..01a4d2cfbeb 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2500,6 +2500,15 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.use_mlock = true; } ).set_env("LLAMA_ARG_MLOCK")); + add_opt(common_arg( + {"--lazy-experts"}, + "do not populate routed MoE expert tensors when mapping the model; let each expert fault in\n" + "on demand the first time it is routed to. Lets a model whose experts do not fit in RAM run\n" + "off the page cache, at the cost of paging during generation. mmap only.", + [](common_params & params) { + params.lazy_experts = true; + } + ).set_env("LLAMA_ARG_LAZY_EXPERTS")); add_opt(common_arg( {"--mmap"}, {"--no-mmap"}, diff --git a/common/common.cpp b/common/common.cpp index 8f13217ab44..3da209f0e3f 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1565,6 +1565,7 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.check_tensors = params.check_tensors; mparams.use_extra_bufts = !params.no_extra_bufts; mparams.no_host = params.no_host; + mparams.lazy_experts = params.lazy_experts; if (params.kv_overrides.empty()) { mparams.kv_overrides = NULL; diff --git a/common/common.h b/common/common.h index bffc1767a7d..c93ea4d342c 100644 --- a/common/common.h +++ b/common/common.h @@ -583,6 +583,7 @@ struct common_params { bool no_op_offload = false; // globally disable offload host tensor operations to device bool no_extra_bufts = false; // disable extra buffer types (used for weight repacking) bool no_host = false; // bypass host buffer allowing extra buffers to be used + bool lazy_experts = false; // mmap: fault routed MoE expert tensors in on demand instead of up front bool single_turn = false; // single turn chat conversation diff --git a/include/llama.h b/include/llama.h index a311ac20235..a5f746b0d99 100644 --- a/include/llama.h +++ b/include/llama.h @@ -328,6 +328,8 @@ extern "C" { bool use_extra_bufts; // use extra buffer types (used for weight repacking) bool no_host; // bypass host buffer allowing extra buffers to be used bool no_alloc; // only load metadata and simulate memory allocations + bool lazy_experts; // mmap only: don't populate routed MoE expert tensors up front, + // let them fault in on demand the first time they are routed to }; struct llama_sampler_seq_config { diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index ed572da7fb5..58b1b3ab688 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -523,6 +523,56 @@ struct llama_mmap::impl { mapped_fragments = std::move(new_mapped_fragments); } + void advise_range(size_t offset, size_t len, advice a) { + if (len == 0 || offset >= size) { + return; + } + len = std::min(len, size - offset); + + // (posix_)madvise requires a page-aligned start address, and tensor offsets are not aligned, + // so snap the range to page boundaries. Advisory hints (WILLNEED/RANDOM) round OUTWARD so the + // whole requested range is still covered; the destructive DONTNEED rounds INWARD so it never + // drops a neighbouring tensor's pages. addr itself is page-aligned (mmap guarantees it). + const uintptr_t page = (uintptr_t) sysconf(_SC_PAGESIZE); + const uintptr_t lo = (uintptr_t) addr + offset; + const uintptr_t hi = lo + len; + uintptr_t astart, aend; + if (a == ADVICE_DONTNEED) { + astart = (lo + (page - 1)) & ~(page - 1); // round up + aend = hi & ~(page - 1); // round down + } else { + astart = lo & ~(page - 1); // round down + aend = (hi + (page - 1)) & ~(page - 1); // round up + const uintptr_t map_end = (uintptr_t) addr + size; + if (aend > map_end) { + aend = map_end; + } + } + if (aend <= astart) { + return; // nothing page-aligned to advise + } + void * const p = (void *) astart; + const size_t alen = (size_t) (aend - astart); + + // NB: posix_madvise() RETURNS the error number and does NOT set errno; madvise() (the Linux + // DONTNEED path) returns -1 and sets errno. Normalise to a single code for the message. + int err = 0; + switch (a) { + case ADVICE_WILLNEED: err = posix_madvise(p, alen, POSIX_MADV_WILLNEED); break; + case ADVICE_RANDOM: err = posix_madvise(p, alen, POSIX_MADV_RANDOM); break; + case ADVICE_DONTNEED: +#ifdef __linux__ + err = madvise(p, alen, MADV_DONTNEED) ? errno : 0; // on Linux this drops the clean file-backed pages +#else + err = posix_madvise(p, alen, POSIX_MADV_DONTNEED); +#endif + break; + } + if (err) { + LLAMA_LOG_WARN("warning: madvise(range, %d) failed: %s\n", (int) a, strerror(err)); + } + } + ~impl() { for (const auto & frag : mapped_fragments) { if (munmap((char *) addr + frag.first, frag.second - frag.first)) { @@ -582,6 +632,27 @@ struct llama_mmap::impl { GGML_UNUSED(last); } + void advise_range(size_t offset, size_t len, advice a) { + if (len == 0 || offset >= size || a != ADVICE_WILLNEED) { + return; // only WILLNEED (prefetch) is actionable here; RANDOM/DONTNEED are hints we skip + } + len = std::min(len, size - offset); +#if _WIN32_WINNT >= 0x602 + BOOL (WINAPI *pPrefetchVirtualMemory) (HANDLE, ULONG_PTR, PWIN32_MEMORY_RANGE_ENTRY, ULONG); + HMODULE hKernel32 = GetModuleHandleW(L"kernel32.dll"); + pPrefetchVirtualMemory = (decltype(pPrefetchVirtualMemory))(void *) GetProcAddress(hKernel32, "PrefetchVirtualMemory"); + if (pPrefetchVirtualMemory) { + WIN32_MEMORY_RANGE_ENTRY range; + range.VirtualAddress = (uint8_t *) addr + offset; + range.NumberOfBytes = (SIZE_T) len; + if (!pPrefetchVirtualMemory(GetCurrentProcess(), 1, &range, 0)) { + LLAMA_LOG_WARN("warning: PrefetchVirtualMemory(range) failed: %s\n", + llama_format_win_err(GetLastError()).c_str()); + } + } +#endif + } + ~impl() { if (hMapping) { if (addr) { @@ -611,6 +682,12 @@ struct llama_mmap::impl { throw std::runtime_error("mmap not supported"); } + + void advise_range(size_t offset, size_t len, advice a) { + GGML_UNUSED(offset); + GGML_UNUSED(len); + GGML_UNUSED(a); + } #endif void * addr; @@ -625,6 +702,8 @@ void * llama_mmap::addr() const { return pimpl->addr; } void llama_mmap::unmap_fragment(size_t first, size_t last) { pimpl->unmap_fragment(first, last); } +void llama_mmap::advise_range(size_t offset, size_t len, advice a) const { pimpl->advise_range(offset, len, a); } + #if defined(_POSIX_MEMLOCK_RANGE) || defined(_WIN32) const bool llama_mmap::SUPPORTED = true; #else diff --git a/src/llama-mmap.h b/src/llama-mmap.h index b7d5c61e95f..3d549aab8f7 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -41,6 +41,13 @@ struct llama_file { }; struct llama_mmap { + // access hint for advise_range(); maps to posix_madvise() / PrefetchVirtualMemory() where available + enum advice { + ADVICE_WILLNEED, // prefetch this range into RAM (readahead) + ADVICE_RANDOM, // random access: disable readahead so neighbours are not pulled in + ADVICE_DONTNEED, // hint that this range is no longer needed (allow eviction) + }; + llama_mmap(const llama_mmap &) = delete; llama_mmap(struct llama_file * file, size_t prefetch = (size_t) -1, bool numa = false); ~llama_mmap(); @@ -50,6 +57,9 @@ struct llama_mmap { void unmap_fragment(size_t first, size_t last); + // apply an access hint to a sub-range of the mapping ([offset, offset+len) clamped to the mapping) + void advise_range(size_t offset, size_t len, advice a) const; + static const bool SUPPORTED; private: diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 28f8bb7934b..50fbee06fea 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1339,11 +1339,21 @@ void llama_model_loader::done_getting_tensors(bool partial) const { } } -void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps) { +// routed MoE expert weight tensors (large, sparsely activated). Deliberately excludes shared +// experts (ffn_*_shexp, active every token) and the small ffn_norm_exps norm. +static bool is_lazy_expert_weight(const std::string & name) { + return name.find(".ffn_gate_exps.") != std::string::npos || + name.find(".ffn_up_exps.") != std::string::npos || + name.find(".ffn_down_exps.") != std::string::npos || + name.find(".ffn_gate_up_exps.") != std::string::npos; +} + +void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps, bool lazy_experts) { if (use_mmap) { mappings.reserve(files.size()); mmaps_used.reserve(files.size()); - for (const auto & file : files) { + for (uint16_t idx = 0; idx < files.size(); ++idx) { + const auto & file = files[idx]; bool is_numa = false; auto * dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); @@ -1355,7 +1365,40 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps } } - std::unique_ptr mapping = std::make_unique(file.get(), prefetch ? -1 : 0, is_numa); + // With lazy experts we do NOT MAP_POPULATE the whole file: expert regions must stay + // unloaded until routed to. Non-expert regions are prefetched explicitly below instead. + const bool lazy_this = lazy_experts && prefetch && !is_numa; + const size_t map_prefetch = (prefetch && !lazy_this) ? (size_t) -1 : 0; + std::unique_ptr mapping = std::make_unique(file.get(), map_prefetch, is_numa); + + if (lazy_this) { + size_t expert_bytes = 0; + size_t prefetch_bytes = 0; + for (const auto & [name, w] : weights_map) { + if (w.idx != idx) { + continue; + } + const size_t nbytes = ggml_nbytes(w.tensor); + if (is_lazy_expert_weight(name)) { + // Faulted in per expert, on demand. Readahead is deliberately left ON: each + // fault then pulls a larger contiguous chunk of the expert row in one go, + // which is much faster when the model is bigger than RAM and disk-bound. + // LLAMA_LAZY_EXPERT_RANDOM forces MADV_RANDOM (no readahead) instead, which + // avoids dragging neighbours in and can be preferable when the model fits. + static const bool force_random = getenv("LLAMA_LAZY_EXPERT_RANDOM") != nullptr; + if (force_random) { + mapping->advise_range(w.offs, nbytes, llama_mmap::ADVICE_RANDOM); + } + expert_bytes += nbytes; + } else { + mapping->advise_range(w.offs, nbytes, llama_mmap::ADVICE_WILLNEED); + prefetch_bytes += nbytes; + } + } + LLAMA_LOG_INFO("%s: lazy experts (file %u): %.1f MiB on-demand, %.1f MiB prefetched\n", + __func__, idx, expert_bytes / (1024.0 * 1024.0), prefetch_bytes / (1024.0 * 1024.0)); + } + mmaps_used.emplace_back(mapping->size(), 0); if (mlock_mmaps) { std::unique_ptr mlock_mmap(new llama_mlock()); @@ -1569,6 +1612,13 @@ bool llama_model_loader::load_all_data( mmap_used.second = std::max(mmap_used.second, weight->offs + n_size); } else { ggml_backend_tensor_set(cur, data, 0, n_size); + + // This tensor now lives in its device (or other non-mmap) buffer, so the bytes we + // just read through the mmap are dead weight. On a model larger than RAM, leaving + // them in the page cache evicts pages that are still needed and thrashes the load. + // Drop them now -- advisory and page-aligned inward, so a clean re-access simply + // re-faults from the file and a neighbour's pages are never touched. + mapping->advise_range(weight->offs, n_size, llama_mmap::ADVICE_DONTNEED); } } else { const auto & file = files.at(weight->idx); diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index c476026d3e5..11511fb25bc 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -186,7 +186,9 @@ struct llama_model_loader { void done_getting_tensors(bool partial = false) const; - void init_mappings(bool prefetch = true, llama_mlocks * mlock_mmaps = nullptr); + // lazy_experts: when true (and prefetching), routed MoE expert weight tensors are NOT populated + // up front; only non-expert regions are prefetched and experts fault in on demand. + void init_mappings(bool prefetch = true, llama_mlocks * mlock_mmaps = nullptr, bool lazy_experts = false); void get_mapping_range(size_t * first, size_t * last, void ** addr, int idx, ggml_context * ctx) const; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 7e5bab26c55..9c05294575c 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1517,7 +1517,7 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } } - ml.init_mappings(true, use_mlock ? &pimpl->mlock_mmaps : nullptr); + ml.init_mappings(true, use_mlock ? &pimpl->mlock_mmaps : nullptr, params.lazy_experts); pimpl->mappings.reserve(ml.mappings.size()); // create the backend buffers @@ -2331,6 +2331,7 @@ llama_model_params llama_model_default_params() { /*.use_extra_bufts =*/ true, /*.no_host =*/ false, /*.no_alloc =*/ false, + /*.lazy_experts =*/ false, }; return result;