llama : add --lazy-experts for MoE models larger than RAM

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"},

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;

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
 

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 {

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

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:

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<llama_mmap> mapping = std::make_unique<llama_mmap>(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<llama_mmap> mapping = std::make_unique<llama_mmap>(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<llama_mlock> 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);

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;
 

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;