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;