loader: add `--defer-ple` to keep per-layer token embedding out of resident memory (#2389)

* loader: disable expert deferral on huge-page mappings

* loader: add --defer-ple for the per-layer token embedding

common/common.cpp

@@ -2180,6 +2180,10 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa
         params.warmup = false;
         return true;
     }
+    if (arg == "--defer-ple") {
+        params.defer_ple = true;
+        return true;
+    }
     if (arg == "--prefetch-experts") {
         params.prefetch_experts = true;
         return true;
@@ -3306,6 +3310,7 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param
     options.push_back({ "*",           "-ncmoe, --n-cpu-moe N",          "keep MoE weights of the first N layers in CPU memory"});
     options.push_back({ "*",           "-thp,   --transparent-huge-pages", "use transparent huge pages on Linux"});
     options.push_back({ "*",           "       --defer-experts",        "defer expert mmap residency on Linux to reduce model load time"});
+    options.push_back({ "*",           "       --defer-ple",            "keep the per-layer token embedding on the file instead of resident in memory (Linux)"});
     options.push_back({ "*",           "       --prefetch-experts",     "stream mmap'd MoE expert weights into the page cache on Linux"});
     options.push_back({ "*",           "       --prefetch-experts-threads N",
                                                                         "number of expert prefetch workers, tune to drive speed/type (default: auto)"});
@@ -4282,6 +4287,7 @@ struct llama_model_params common_model_params_to_llama(const gpt_params & params
     mparams.mtp             = params.has_mtp || params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP);
     mparams.flash_attn      = params.flash_attn;
     mparams.defer_experts   = params.defer_experts;
+    mparams.defer_ple       = params.defer_ple;
     mparams.swa_compress    = params.swa_compress;
     if (params.kv_overrides.empty()) {
         mparams.kv_overrides = NULL;
@@ -5372,6 +5378,7 @@ void yaml_dump_non_result_info(FILE * stream, const gpt_params & params, const l
     fprintf(stream, "merge_qkv: %s # default: false\n", params.merge_qkv ? "true" : "false");
     fprintf(stream, "merge_up_gate_exps: %s # default: false\n", params.merge_up_gate_exps ? "true" : "false");
     fprintf(stream, "defer_experts: %s # default: false\n", params.defer_experts ? "true" : "false");
+    fprintf(stream, "defer_ple: %s # default: false\n", params.defer_ple ? "true" : "false");
     fprintf(stream, "prefetch_experts: %s # default: false\n", params.prefetch_experts ? "true" : "false");
     fprintf(stream, "prefetch_experts_threads: %d # default: 0 (auto)\n", params.prefetch_experts_threads);
     fprintf(stream, "max_extra_alloc: %d # default: 256\n", params.max_extra_alloc_MiB);

common/common.h

@@ -456,6 +456,7 @@ struct gpt_params {
     bool merge_qkv         = false; // if true, merge separate Q, K, V tensors into a single, contiguous tensor
     bool merge_up_gate_exps= false; // if true, merge ffn_up_exps and ffn_gate_exps into a single, contiguous tensor
     bool defer_experts     = false; // if true, defer expert mmap residency to speed up model loading (Linux only)
+    bool defer_ple         = false; // if true, keep the per-layer token embedding on the file (Linux only)
     bool prefetch_experts  = false; // if true, stream mmap'd MoE expert weights into the page cache (Linux only)
     int  prefetch_experts_threads = 0; // number of expert prefetch workers (<=0 = auto)
     bool k_cache_hadamard  = false; // if true, use Hadamard transform for the K-cache (only makes sense with quantized cache)

include/llama.h

@@ -433,6 +433,7 @@ extern "C" {
         bool dry_run;       // skip loading tensors
         bool flash_attn;
         bool defer_experts;    // defer expert mmap residency to speed up model loading (Linux only)
+        bool defer_ple;        // keep the per-layer token embedding on the file instead of resident in memory (Linux only)
         bool swa_compress;     // must match llama_context_params::swa_compress; the fit also assumes that context's n_ubatch
     };
 

src/llama-mmap.cpp

@@ -399,6 +399,20 @@ struct llama_mmap::impl {
 #endif
     }
 
+    void random_fragment(size_t first, size_t last) {
+        int page_size = mapped_page_size > 0 ? mapped_page_size : sysconf(_SC_PAGESIZE);
+        align_range(&first, &last, page_size);
+        size_t len = last - first;
+
+        if (len == 0) {
+            return;
+        }
+
+        if (int err = posix_madvise((uint8_t *) addr + first, len, POSIX_MADV_RANDOM)) {
+            LLAMA_LOG_WARN("warning: posix_madvise(..., POSIX_MADV_RANDOM) failed: %s\n", strerror(err));
+        }
+    }
+
     void unmap_fragment(size_t first, size_t last) {
         int page_size = mapped_page_size > 0 ? mapped_page_size : sysconf(_SC_PAGESIZE);
         align_range(&first, &last, page_size);
@@ -493,6 +507,11 @@ struct llama_mmap::impl {
         GGML_UNUSED(last);
     }
 
+    void random_fragment(size_t first, size_t last) {
+        GGML_UNUSED(first);
+        GGML_UNUSED(last);
+    }
+
     void unmap_fragment(size_t first, size_t last) {
         GGML_UNUSED(first);
         GGML_UNUSED(last);
@@ -520,6 +539,13 @@ struct llama_mmap::impl {
         throw std::runtime_error("mmap not supported");
     }
 
+    void random_fragment(size_t first, size_t last) {
+        GGML_UNUSED(first);
+        GGML_UNUSED(last);
+
+        throw std::runtime_error("mmap not supported");
+    }
+
     void unmap_fragment(size_t first, size_t last) {
         GGML_UNUSED(first);
         GGML_UNUSED(last);
@@ -528,6 +554,9 @@ struct llama_mmap::impl {
     }
 #endif
 
+    // set only when the huge-page mapping succeeded, which is the anonymous case
+    bool is_anonymous() const { return mapped_page_size > 0; }
+
     void * addr;
     size_t size;
     size_t mapped_page_size = 0;
@@ -539,8 +568,10 @@ llama_mmap::~llama_mmap() = default;
 
 size_t llama_mmap::size() const { return pimpl->size; }
 void * llama_mmap::addr() const { return pimpl->addr; }
+bool llama_mmap::is_anonymous() const { return pimpl->is_anonymous(); }
 
 void llama_mmap::dontneed_fragment(size_t first, size_t last) { pimpl->dontneed_fragment(first, last); }
+void llama_mmap::random_fragment(size_t first, size_t last) { pimpl->random_fragment(first, last); }
 void llama_mmap::unmap_fragment(size_t first, size_t last) { pimpl->unmap_fragment(first, last); }
 
 #if defined(_POSIX_MEMLOCK_RANGE) || defined(_WIN32)

src/llama-mmap.h

@@ -45,9 +45,12 @@ struct llama_mmap {
 
     size_t size() const;
     void * addr() const;
+    bool is_anonymous() const;
 
     void dontneed_fragment(size_t first, size_t last);
 
+    void random_fragment(size_t first, size_t last);
+
     void unmap_fragment(size_t first, size_t last);
 
     static const bool SUPPORTED;

src/llama-model-loader.cpp

@@ -662,6 +662,15 @@ bool llama_model_loader::should_defer_expert_mmaps() const {
     return defer_experts && use_mmap && !expert_tensor_index.empty();
 }
 
+bool llama_model_loader::has_anonymous_mapping() const {
+    for (const auto & mapping : mappings) {
+        if (mapping->is_anonymous()) {
+            return true;
+        }
+    }
+    return false;
+}
+
 void llama_model_loader::drop_mmap_expert_pages() const {
     if (!use_mmap || mappings.empty() || expert_tensor_index.file_ranges.empty()) {
         return;
@@ -676,6 +685,35 @@ void llama_model_loader::drop_mmap_expert_pages() const {
     }
 }
 
+void llama_model_loader::build_ple_tensor_index() {
+    ple_tensor_index = {};
+
+    const auto * weight = get_weight(LLM_TN(get_arch())(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").c_str());
+    if (weight == nullptr) {
+        return;
+    }
+
+    const size_t tensor_bytes = ggml_nbytes(weight->tensor);
+
+    ple_tensor_index.file_ranges.resize(files.size());
+    ple_tensor_index.file_ranges.at(weight->idx).push_back({ weight->offs, weight->offs + tensor_bytes });
+    ple_tensor_index.deferred_bytes = tensor_bytes;
+}
+
+bool llama_model_loader::should_defer_ple_mmaps() const {
+    return defer_ple && use_mmap && !ple_tensor_index.empty();
+}
+
+void llama_model_loader::apply_ple_mmap_policy() const {
+    for (size_t idx = 0; idx < ple_tensor_index.file_ranges.size(); ++idx) {
+        for (const auto & range : ple_tensor_index.file_ranges[idx]) {
+            // sparse row lookups through get_rows: readahead would evict more than it fetches
+            mappings[idx]->random_fragment(range.first, range.last);
+            mappings[idx]->dontneed_fragment(range.first, range.last);
+        }
+    }
+}
+
 template<typename T>
 typename std::enable_if<std::is_integral<T>::value, bool>::type
 llama_model_loader::get_arr_n(const std::string & key, T & result, const bool required) {

src/llama-model-loader.h

@@ -51,6 +51,7 @@ struct llama_model_loader {
     bool merge_qkv = false;
     bool merge_up_gate_exps = false;
     bool defer_experts = false;
+    bool defer_ple = false;
 
     llama_files files;
     llama_ftype ftype;
@@ -91,6 +92,7 @@ struct llama_model_loader {
     mutable bool arch_resolved = false;
     mutable llm_arch resolved_arch = LLM_ARCH_UNKNOWN;
     llama_expert_tensor_index expert_tensor_index;
+    llama_expert_tensor_index ple_tensor_index;
 
     llama_model_loader(const std::string & fname, int ncmoe, bool use_mmap, bool check_tensors, bool repack_tensors, bool use_thp,
             bool merge_qkv, bool merge_up_gate_exps, bool defer_experts,
@@ -183,8 +185,16 @@ struct llama_model_loader {
 
     bool should_defer_expert_mmaps() const;
 
+    bool has_anonymous_mapping() const;
+
     void drop_mmap_expert_pages() const;
 
+    void build_ple_tensor_index();
+
+    bool should_defer_ple_mmaps() const;
+
+    void apply_ple_mmap_policy() const;
+
     void get_mapping_range(size_t * first, size_t * last, void ** addr, int idx, ggml_context * ctx) const;
 
     // for backwards compatibility, does not support ggml-backend

src/llama.cpp

@@ -4745,10 +4745,36 @@ static bool llm_load_tensors(
         defer_expert_mmap = false;
     }
 
+    bool defer_ple_mmap = ml.should_defer_ple_mmaps();
+    if (defer_ple_mmap && use_mlock) {
+        LLAMA_LOG_WARN("%s: deferred per-layer token embedding disabled because mlock keeps mmap ranges resident\n", __func__);
+        defer_ple_mmap = false;
+    }
+    if (ml.defer_ple && !ml.use_mmap && !ml.ple_tensor_index.empty()) {
+        LLAMA_LOG_WARN("%s: --defer-ple had no effect: creating the tensors disabled mmap\n", __func__);
+    }
+
     ml.done_getting_tensors();
 
     // --dry-run skips MAP_POPULATE/WILLNEED — tensor data is never read.
-    ml.init_mappings(!defer_expert_mmap && !dry_run, use_mlock ? &model.mlock_mmaps : nullptr, ml.use_thp);
+    ml.init_mappings(!defer_expert_mmap && !defer_ple_mmap && !dry_run, use_mlock ? &model.mlock_mmaps : nullptr, ml.use_thp);
+
+    // dropping a range discards an anonymous huge-page mapping, so test the mapping and not the -thp flag
+    if (ml.has_anonymous_mapping()) {
+        if (defer_expert_mmap) {
+            LLAMA_LOG_WARN("%s: deferred expert loading disabled because the model is mapped on huge pages\n", __func__);
+            defer_expert_mmap = false;
+        }
+        if (defer_ple_mmap) {
+            LLAMA_LOG_WARN("%s: deferred per-layer token embedding disabled because the model is mapped on huge pages\n", __func__);
+            defer_ple_mmap = false;
+        }
+    }
+    if (defer_ple_mmap && !dry_run) {
+        LLAMA_LOG_INFO("%s: deferring %.2f GiB of per-layer token embedding to the file\n", __func__,
+                ml.ple_tensor_index.deferred_bytes / 1024.0 / 1024.0 / 1024.0);
+    }
+
     model.mappings.reserve(ml.mappings.size());
 
     // create the backend buffers
@@ -4889,6 +4915,10 @@ static bool llm_load_tensors(
         if (defer_expert_mmap) {
             ml.drop_mmap_expert_pages();
         }
+
+        if (defer_ple_mmap) {
+            ml.apply_ple_mmap_policy();
+        }
     }
 
     if (model.is_mla_model()) {
@@ -5010,6 +5040,8 @@ static int llama_model_load(const std::string & fname, llama_model & model, llam
 
         model.mtp = params.mtp;
 
+        ml.defer_ple = params.defer_ple;
+
         try {
             llm_load_arch(ml, model);
         } catch(const std::exception & e) {
@@ -5034,6 +5066,20 @@ static int llama_model_load(const std::string & fname, llama_model & model, llam
             ml.build_expert_tensor_index(model.hparams);
 #else
             LLAMA_LOG_WARN("%s: deferred expert loading is only supported on Linux; ignoring defer_experts\n", __func__);
+#endif
+        }
+        if (params.defer_ple) {
+#ifdef __linux__
+            if (!params.use_mmap) {
+                LLAMA_LOG_WARN("%s: --defer-ple had no effect: mmap is disabled\n", __func__);
+            } else {
+                ml.build_ple_tensor_index();
+                if (ml.ple_tensor_index.empty()) {
+                    LLAMA_LOG_WARN("%s: --defer-ple had no effect: no per-layer token embedding\n", __func__);
+                }
+            }
+#else
+            LLAMA_LOG_WARN("%s: deferred per-layer token embedding is only supported on Linux; ignoring defer_ple\n", __func__);
 #endif
         }
         try {
@@ -7888,6 +7934,7 @@ struct llama_model_params llama_model_default_params() {
         /*.dry_run                     =*/ false,
         /*.flash_attn                  =*/ true,
         /*.defer_experts               =*/ false,
+        /*.defer_ple                   =*/ false,
         /*.swa_compress                =*/ false,
     };