diff --git a/common/common.cpp b/common/common.cpp index 072190f7..44176139 100644 --- a/common/common.cpp +++ b/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); diff --git a/common/common.h b/common/common.h index dc360925..ed1123cb 100644 --- a/common/common.h +++ b/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) diff --git a/include/llama.h b/include/llama.h index 82f81680..160e9643 100644 --- a/include/llama.h +++ b/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 }; diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index 599a47f7..76aa6627 100644 --- a/src/llama-mmap.cpp +++ b/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) diff --git a/src/llama-mmap.h b/src/llama-mmap.h index b7ee0168..febbb32e 100644 --- a/src/llama-mmap.h +++ b/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; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 37d27ae9..738b13de 100644 --- a/src/llama-model-loader.cpp +++ b/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 std::enable_if::value, bool>::type llama_model_loader::get_arr_n(const std::string & key, T & result, const bool required) { diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index a07eefb8..c4465032 100644 --- a/src/llama-model-loader.h +++ b/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 diff --git a/src/llama.cpp b/src/llama.cpp index 47683357..d2f33b71 100644 --- a/src/llama.cpp +++ b/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, };