From 86f4f516e508bead7c4c9597fed57ac788792173 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Wed, 25 Mar 2026 07:29:29 +0100 Subject: [PATCH] Auto-fit offloaded tensors to available VRAM (MoE models) (#1501) * WIP: automatically fit model in available VRAM * WIP * This seems pretty solid --- common/common.cpp | 21 +++ common/common.h | 2 + include/llama.h | 2 + src/llama-arch.h | 44 +++---- src/llama-model.cpp | 7 +- src/llama.cpp | 304 +++++++++++++++++++++++++++++++++++++------- 6 files changed, 313 insertions(+), 67 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index f1770e2d..0c3b423e 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1500,6 +1500,21 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa //} return true; } + if (arg == "--fit") { + params.fit = true; + return true; + } + if (arg == "--fit-margin") { + CHECK_ARG; + int32_t margin = std::stoi(argv[i]); + if (margin < 0) { + fprintf(stderr, "error: Invalid value for --fit-margin: %d (must be >= 0)\n", margin); + invalid_param = true; + } else { + params.fit_margin = margin; + } + return true; + } if (arg == "--no-mmap") { params.use_mmap = false; return true; @@ -2479,6 +2494,8 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param options.push_back({ "*", " --run-time-repack", "repack tensors if interleaved variant is available"}); options.push_back({ "*", " --cpu-moe", "keep all MoE weights in CPU memory"}); options.push_back({ "*", " --n-cpu-moe N", "keep MoE weights of the first N layers in CPU memory"}); + options.push_back({ "*", " --fit-margin N", "safety margin in MiB when auto-fitting model offloading"}); + options.push_back({ "*", " --fit", "automatically determine which tensors to offload to the GPU(s)"}); options.push_back({ "*", " --numa TYPE", "attempt optimizations that help on some NUMA systems\n" " - distribute: spread execution evenly over all nodes\n" " - isolate: only spawn threads on CPUs on the node that execution started on\n" @@ -3311,6 +3328,8 @@ struct llama_model_params common_model_params_to_llama(const gpt_params & params mparams.main_gpu = params.main_gpu; mparams.max_gpu = params.max_gpu; mparams.ncmoe = params.ncmoe; + mparams.fit = params.fit; + mparams.fit_margin = params.fit_margin; mparams.type_k = kv_cache_type_from_str(params.cache_type_k); mparams.type_v = kv_cache_type_from_str(params.cache_type_v); mparams.max_ctx_size = params.n_ctx; @@ -4357,6 +4376,8 @@ void yaml_dump_non_result_info(FILE * stream, const gpt_params & params, const l fprintf(stream, "main_gpu: %d # default: 0\n", params.main_gpu); fprintf(stream, "max_gpu: %d # default: 0\n", params.max_gpu); fprintf(stream, "ncmoe: %d # default: 0\n", params.ncmoe); + fprintf(stream, "fit: %d # default: false\n", params.fit); + fprintf(stream, "fit_margin: %d # default: 0\n", params.fit_margin); fprintf(stream, "min_keep: %d # default: 0 (disabled)\n", sparams.min_keep); fprintf(stream, "mirostat: %d # default: 0 (disabled)\n", sparams.mirostat); fprintf(stream, "mirostat_ent: %f # default: 5.0\n", sparams.mirostat_tau); diff --git a/common/common.h b/common/common.h index 7d0ea30a..67f091a8 100644 --- a/common/common.h +++ b/common/common.h @@ -226,6 +226,8 @@ struct gpt_params { int32_t main_gpu = 0; // the GPU that is used for scratch and small tensors int32_t max_gpu = 0; // max number of GPUs to use at a time for split mode "graph" int32_t ncmoe = 0; // number of layers in which MoE tensors are left in VRAM + int32_t fit_margin = 0; // safety margin for auto-fit in MiB + bool fit = false; // automatically fit model (for now just using MoE tensor overrides) float tensor_split[128] = {0}; // how split tensors should be distributed across GPUs int32_t grp_attn_n = 1; // group-attention factor int32_t grp_attn_w = 512; // group-attention width diff --git a/include/llama.h b/include/llama.h index 186b16fb..d823c72e 100644 --- a/include/llama.h +++ b/include/llama.h @@ -378,6 +378,8 @@ extern "C" { int32_t n_seq_max; int32_t n_ubatch; int32_t amb; + int32_t fit_margin; + bool fit; // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() const float * tensor_split; diff --git a/src/llama-arch.h b/src/llama-arch.h index 6b56dfe1..3f621174 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -235,111 +235,111 @@ struct LLM_KV { }; enum llm_tensor { - LLM_TENSOR_TOKEN_EMBD, + LLM_TENSOR_TOKEN_EMBD, // 0 LLM_TENSOR_TOKEN_EMBD_NORM, LLM_TENSOR_TOKEN_TYPES, LLM_TENSOR_POS_EMBD, LLM_TENSOR_OUTPUT, - LLM_TENSOR_OUTPUT_NORM, + LLM_TENSOR_OUTPUT_NORM, // 5 LLM_TENSOR_ROPE_FREQS, LLM_TENSOR_ROPE_FACTORS_LONG, LLM_TENSOR_ROPE_FACTORS_SHORT, LLM_TENSOR_ATTN_Q, - LLM_TENSOR_ATTN_K, + LLM_TENSOR_ATTN_K, // 10 LLM_TENSOR_ATTN_V, LLM_TENSOR_ATTN_QKV, LLM_TENSOR_ATTN_OUT, LLM_TENSOR_ATTN_NORM, - LLM_TENSOR_ATTN_NORM_2, + LLM_TENSOR_ATTN_NORM_2, // 15 LLM_TENSOR_ATTN_OUT_NORM, LLM_TENSOR_ATTN_POST_NORM, LLM_TENSOR_ATTN_ROT_EMBD, LLM_TENSOR_ATTN_SINKS, - LLM_TENSOR_ATTN_GATE, + LLM_TENSOR_ATTN_GATE, // 20 LLM_TENSOR_FFN_GATE_INP, LLM_TENSOR_FFN_GATE_INP_SHEXP, LLM_TENSOR_FFN_NORM, LLM_TENSOR_FFN_POST_NORM, - LLM_TENSOR_FFN_GATE, + LLM_TENSOR_FFN_GATE, // 25 LLM_TENSOR_FFN_DOWN, LLM_TENSOR_FFN_UP, LLM_TENSOR_FFN_ACT, LLM_TENSOR_FFN_DOWN_EXP, // split experts for backward compatibility - LLM_TENSOR_FFN_GATE_EXP, + LLM_TENSOR_FFN_GATE_EXP, // 30 LLM_TENSOR_FFN_UP_EXP, LLM_TENSOR_FFN_NORM_EXPS, LLM_TENSOR_FFN_DOWN_EXPS, // merged experts LLM_TENSOR_FFN_GATE_EXPS, - LLM_TENSOR_FFN_UP_EXPS, + LLM_TENSOR_FFN_UP_EXPS, // 35 LLM_TENSOR_FFN_GATE_UP_EXPS, LLM_TENSOR_FFN_DOWN_SHEXP, LLM_TENSOR_FFN_GATE_SHEXP, LLM_TENSOR_FFN_UP_SHEXP, - LLM_TENSOR_FFN_EXP_PROBS_B, + LLM_TENSOR_FFN_EXP_PROBS_B, // 40 LLM_TENSOR_ATTN_Q_NORM, LLM_TENSOR_ATTN_K_NORM, LLM_TENSOR_LAYER_OUT_NORM, LLM_TENSOR_SSM_IN, - LLM_TENSOR_SSM_CONV1D, + LLM_TENSOR_SSM_CONV1D, // 45 LLM_TENSOR_SSM_X, LLM_TENSOR_SSM_DT, LLM_TENSOR_SSM_A, LLM_TENSOR_SSM_A_NOSCAN, - LLM_TENSOR_SSM_D, + LLM_TENSOR_SSM_D, // 45 LLM_TENSOR_SSM_NORM, LLM_TENSOR_SSM_OUT, LLM_TENSOR_SSM_BETA_ALPHA, LLM_TENSOR_SSM_ALPHA, - LLM_TENSOR_SSM_BETA, + LLM_TENSOR_SSM_BETA, // 50 LLM_TENSOR_ATTN_Q_A, LLM_TENSOR_ATTN_Q_B, LLM_TENSOR_ATTN_KV_A_MQA, LLM_TENSOR_ATTN_KQ_A_MQA, - LLM_TENSOR_ATTN_KV_B, + LLM_TENSOR_ATTN_KV_B, // 55 LLM_TENSOR_ATTN_K_B, LLM_TENSOR_ATTN_V_B, LLM_TENSOR_ATTN_Q_A_NORM, LLM_TENSOR_ATTN_KV_A_NORM, - LLM_TENSOR_ATTN_SUB_NORM, + LLM_TENSOR_ATTN_SUB_NORM, // 60 LLM_TENSOR_FFN_SUB_NORM, LLM_TENSOR_DEC_ATTN_NORM, LLM_TENSOR_DEC_ATTN_Q, LLM_TENSOR_DEC_ATTN_K, - LLM_TENSOR_DEC_ATTN_V, + LLM_TENSOR_DEC_ATTN_V, // 65 LLM_TENSOR_DEC_ATTN_OUT, LLM_TENSOR_DEC_ATTN_REL_B, LLM_TENSOR_DEC_CROSS_ATTN_NORM, LLM_TENSOR_DEC_CROSS_ATTN_Q, - LLM_TENSOR_DEC_CROSS_ATTN_K, + LLM_TENSOR_DEC_CROSS_ATTN_K, // 70 LLM_TENSOR_DEC_CROSS_ATTN_V, LLM_TENSOR_DEC_CROSS_ATTN_OUT, LLM_TENSOR_DEC_CROSS_ATTN_REL_B, LLM_TENSOR_DEC_FFN_NORM, - LLM_TENSOR_DEC_FFN_GATE, + LLM_TENSOR_DEC_FFN_GATE, // 75 LLM_TENSOR_DEC_FFN_DOWN, LLM_TENSOR_DEC_FFN_UP, LLM_TENSOR_DEC_OUTPUT_NORM, LLM_TENSOR_ENC_ATTN_NORM, - LLM_TENSOR_ENC_ATTN_Q, + LLM_TENSOR_ENC_ATTN_Q, // 80 LLM_TENSOR_ENC_ATTN_K, LLM_TENSOR_ENC_ATTN_V, LLM_TENSOR_ENC_ATTN_OUT, LLM_TENSOR_ENC_ATTN_REL_B, - LLM_TENSOR_ENC_FFN_NORM, + LLM_TENSOR_ENC_FFN_NORM, // 85 LLM_TENSOR_ENC_FFN_GATE, LLM_TENSOR_ENC_FFN_DOWN, LLM_TENSOR_ENC_FFN_UP, LLM_TENSOR_ENC_OUTPUT_NORM, - LLM_TENSOR_NEXTN_EH_PROJ, + LLM_TENSOR_NEXTN_EH_PROJ, // 90 LLM_TENSOR_NEXTN_EMBED_TOKENS, LLM_TENSOR_NEXTN_ENORM, LLM_TENSOR_NEXTN_HNORM, LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, - LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, + LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, // 95 LLM_TENSOR_INDEXER_K_NORM, LLM_TENSOR_INDEXER_PROJ, LLM_TENSOR_INDEXER_ATTN_K, - LLM_TENSOR_INDEXER_ATTN_Q_B, + LLM_TENSOR_INDEXER_ATTN_Q_B, // 97 LLM_TENSOR_UNKNOWN, }; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index dd88a4cf..a2f3373c 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1852,7 +1852,12 @@ llm_tensor llm_tensor_type(llm_arch arch, const std::string & tensor_name, int i return LLM_TENSOR_UNKNOWN; } for (auto & entry : it->second) { - auto this_name = ::format(entry.second.c_str(), il); + auto base_name = ::format(entry.second.c_str(), il); + auto this_name = base_name + ".weight"; + if (tensor_name.find(this_name) == 0) { + return entry.first; + } + this_name = base_name + ".bias"; if (tensor_name.find(this_name) == 0) { return entry.first; } diff --git a/src/llama.cpp b/src/llama.cpp index c28e04f1..1154ec8f 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -1997,9 +1997,15 @@ static bool is_model_split_supported(const llama_model & model) { return it != k_supported.end(); } +struct expert_tensors { + ggml_tensor * up = nullptr; + ggml_tensor * gate = nullptr; + ggml_tensor * down = nullptr; +}; static std::pair, double> get_layer_sizes(const llama_model_loader & ml, const llama_model & model, - ggml_type cache_type_k, ggml_type cache_type_v, uint32_t max_ctx_size, int mla_attn, int n_seq_max, int n_ubatch, int amb, bool flash_attn) { + ggml_type cache_type_k, ggml_type cache_type_v, uint32_t max_ctx_size, int mla_attn, int n_seq_max, int n_ubatch, int amb, bool flash_attn, + std::vector & experts) { int n_layer = model.hparams.n_layer; std::vector result(n_layer+1, 0); std::vector compute(n_layer+1, 0); @@ -2013,6 +2019,7 @@ static std::pair, double> get_layer_sizes(const llama_model_ if (has_mla) { mla_tensors.resize(n_layer); } + experts.resize(n_layer); std::vector ffn_exps(n_layer, 0), ffn_shexp(n_layer, 0), ffn_dense(n_layer, 0); std::vector attn(n_layer, 0); std::vector has_layer_norm(n_layer, false); @@ -2035,23 +2042,23 @@ static std::pair, double> get_layer_sizes(const llama_model_ } auto pos = name.find("blk."); if (pos != 0) { - printf("Oops: tensor with strange name %s\n", name.c_str()); + LLAMA_LOG_WARN("Oops: tensor with strange name %s\n", name.c_str()); continue; } pos += 4; auto pos1 = name.find('.', pos); if (pos1 == std::string::npos) { - printf("Oops: tensor with strange name %s\n", name.c_str()); + LLAMA_LOG_WARN("Oops: tensor with strange name %s\n", name.c_str()); continue; } auto layer_string = name.substr(pos, pos1-pos); std::istringstream str(layer_string); int il; str >> il; if (str.fail()) { - printf("Oops: failed to read layer index from %s for tensor %s\n", layer_string.c_str(), name.c_str()); + LLAMA_LOG_WARN("Oops: failed to read layer index from %s for tensor %s\n", layer_string.c_str(), name.c_str()); } if (il < 0 || il >= model.hparams.n_layer) { - printf("Oops: strange layer index %d for tensor %s\n", il, name.c_str()); + LLAMA_LOG_WARN("Oops: strange layer index %d for tensor %s\n", il, name.c_str()); continue; } result[il] += size; @@ -2076,10 +2083,24 @@ static std::pair, double> get_layer_sizes(const llama_model_ } if (!is_mla) { auto ttype = llm_tensor_type(model.arch, name, il); - if (ttype == LLM_TENSOR_UNKNOWN) printf("Oops: got unknows for tensor %s\n", name.c_str()); + if (ttype == LLM_TENSOR_UNKNOWN) { + LLAMA_LOG_WARN("Oops: got unknown for tensor %s\n", name.c_str()); + continue; + } if (ttype == LLM_TENSOR_FFN_GATE_UP_EXPS || ttype == LLM_TENSOR_FFN_GATE_EXPS || ttype == LLM_TENSOR_FFN_UP_EXPS || ttype == LLM_TENSOR_FFN_DOWN_EXPS) { auto size = t->ne[1] * n_ubatch * model.hparams.n_expert_used * sizeof(float); ffn_exps[il] += size; + if (name.find(".bias") == std::string::npos) { + if (ttype == LLM_TENSOR_FFN_GATE_UP_EXPS || ttype == LLM_TENSOR_FFN_UP_EXPS) { + experts[il].up = t; + } + else if (ttype == LLM_TENSOR_FFN_GATE_EXPS) { + experts[il].gate = t; + } + else if (ttype == LLM_TENSOR_FFN_DOWN_EXPS) { + experts[il].down = t; + } + } } else if (ttype == LLM_TENSOR_FFN_UP_SHEXP || ttype == LLM_TENSOR_FFN_GATE_SHEXP || ttype == LLM_TENSOR_FFN_DOWN_SHEXP) { ffn_shexp[il] += t->ne[1] * n_ubatch * sizeof(float); @@ -2099,6 +2120,9 @@ static std::pair, double> get_layer_sizes(const llama_model_ ttype == LLM_TENSOR_FFN_NORM || ttype == LLM_TENSOR_FFN_POST_NORM || ttype == LLM_TENSOR_LAYER_OUT_NORM) { has_layer_norm[il] = true; } + //else { + // printf("Unhandled tensor %s of type %d\n", name.c_str(), int(ttype)); + //} } } for (int il = 0; il < n_layer; ++il) { @@ -2178,15 +2202,19 @@ static bool llm_load_tensors( int n_seq_max, int n_ubatch, int amb, + int fit_margin, bool flash_attn, bool use_mlock, bool validate_quants, bool mtp, + bool fit, bool dry_run, llama_progress_callback progress_callback, void * progress_callback_user_data) { model.t_start_us = ggml_time_us(); + constexpr size_t k_default_mem_margin = 1024 * 1024 * 1024; + auto & hparams = model.hparams; if (split_mode == LLAMA_SPLIT_MODE_GRAPH || split_mode == LLAMA_SPLIT_MODE_ATTN) { @@ -2216,12 +2244,24 @@ static bool llm_load_tensors( max_ctx_size = model.hparams.n_ctx_train; } + if (fit) { + if (ml.tensor_buft_overrides) { + throw std::runtime_error("Manual tensor overrides cannot be used with --fit"); + } + if (ml.ncmoe > 0) { + throw std::runtime_error("-ncmoe | --n-cpu-moe cannot be used with --fit"); + } + n_gpu_layers = 999; + } + model.split_mode = split_mode; model.main_gpu = main_gpu; model.max_gpu = max_gpu; model.n_gpu_layers = n_gpu_layers; model.mtp = mtp; + size_t mem_margin = fit_margin > 0 ? size_t(fit_margin)*1024*1024 : k_default_mem_margin; + const int n_layer = hparams.n_layer; const int i_gpu_start = std::max((int) hparams.n_layer - n_gpu_layers, (int) 0); bool use_mmap_buffer = true; @@ -2236,13 +2276,24 @@ static bool llm_load_tensors( model.buft_layer[i] = llama_default_buffer_type_cpu(true); } + std::vector device_mem(model.devices.size()); + for (int i = 0; i < int(device_mem.size()); ++i) { + device_mem[i] = llama_get_device_memory(model, model.devices[i]); + if (device_mem[i] > mem_margin) { + device_mem[i] -= mem_margin; + } else { + LLAMA_LOG_WARN("Free memory %zu MiB on device %d is less the %zu MiB safety margin\n", device_mem[i]/(1024*1024), model.devices[i], mem_margin/(1024*1024)); + device_mem[i] = 0; + } + } + if (int device_count = model.devices.size(); device_count > 1) { bool all_zero = tensor_split == nullptr || std::all_of(tensor_split, tensor_split + device_count, [](float x) { return x == 0.0f; }); std::vector splits(device_count); if (all_zero) { // default split, by free memory for (int i = 0; i < device_count; ++i) { - splits[i] = llama_get_device_memory(model, model.devices[i]); + splits[i] = device_mem[i]; } } else { std::copy(tensor_split, tensor_split + device_count, splits.begin()); @@ -2263,44 +2314,193 @@ static bool llm_load_tensors( } int device_count = model.splits.size(); + if (device_count < 2 && (model.split_mode == LLAMA_SPLIT_MODE_ATTN || model.split_mode == LLAMA_SPLIT_MODE_GRAPH)) { + throw std::runtime_error("It is not possible to use split mode 'graph' with less than 2 devices"); + } + if (fit && device_count > 1) { + model.main_gpu = device_count - 1; + } model.default_layer_device = std::vector(hparams.n_layer+1, device_count-1); int act_gpu_layers = std::min(n_gpu_layers, (int)n_layer + 1); - if (device_count > 1) { - auto [layer_sizes, max_compute] = get_layer_sizes(ml, model, cache_type_k, cache_type_v, max_ctx_size, mla_attn, n_seq_max, n_ubatch, amb, flash_attn); - int n_last = n_layer; - if (n_gpu_layers > n_layer) ++n_last; - double sum = max_compute * device_count; - for (int i = i_gpu_start; i < n_last; ++i) sum += layer_sizes[i]; - int last = i_gpu_start; - float loaded_sum = 0; + std::vector overrides; + if (device_count > 0) { + std::vector experts; + auto [layer_sizes, max_compute] = get_layer_sizes(ml, model, cache_type_k, cache_type_v, max_ctx_size, mla_attn, n_seq_max, n_ubatch, amb, flash_attn, experts); + size_t required_mem = 0; + for (int i = 0; i <= n_layer; ++i) { + required_mem += layer_sizes[i]; + } + size_t available_mem = 0; for (int id = 0; id < device_count; ++id) { - loaded_sum += max_compute; - float split_size = model.splits[id]*sum; - int il = last; - for (; il < n_last; ++il) { - if (loaded_sum + layer_sizes[il] <= split_size) { - model.default_layer_device[il] = id; - loaded_sum += layer_sizes[il]; - LLAMA_LOG_INFO("Setting default device in layer %2d to %d\n", il, id); - } else { - if (loaded_sum + layer_sizes[il] - split_size < split_size - loaded_sum) { - LLAMA_LOG_INFO("Setting default device in layer %2d to %d\n", il, id); - model.default_layer_device[il] = id; - loaded_sum += layer_sizes[il++]; - } - break; + if (device_mem[id] > max_compute) { + available_mem += device_mem[id] - max_compute; + } else { + LLAMA_LOG_WARN("Free memory %zu MiB on device %d is less the required compute buffer size %g MiB\n", device_mem[id]/(1024*1024), id, max_compute/(1024*1024)); + } + } + LLAMA_LOG_INFO("Memory required for model tensors + cache: %.f MiB\n", required_mem/(1024.*1024.)); + LLAMA_LOG_INFO("Memory available on all devices - compute: %.f MiB\n", available_mem/(1024.*1024.)); + if (required_mem > available_mem) { + float sum = 0; + for (int id = 0; id < device_count; ++id) { + device_mem[id] = device_mem[id] > max_compute ? device_mem[id] - max_compute : 0; + sum += device_mem[id]; + model.splits[id] = sum; + } + if (sum > 0) { + for (int id = 0; id < device_count; ++id) { + model.splits[id] /= sum; + } + } + } + if (device_count > 1) { + int n_last = n_layer; + if (n_gpu_layers > n_layer) ++n_last; + double sum = max_compute * device_count; + for (int i = i_gpu_start; i < n_last; ++i) sum += layer_sizes[i]; + int last = i_gpu_start; + float loaded_sum = 0; + for (int id = 0; id < device_count; ++id) { + loaded_sum += max_compute; + float split_size = model.splits[id]*sum; + int il = last; + for (; il < n_last; ++il) { + if (loaded_sum + layer_sizes[il] <= split_size) { + model.default_layer_device[il] = id; + loaded_sum += layer_sizes[il]; + LLAMA_LOG_INFO("Setting default device in layer %2d to %d\n", il, id); + } else { + if (loaded_sum + layer_sizes[il] - split_size < split_size - loaded_sum) { + LLAMA_LOG_INFO("Setting default device in layer %2d to %d\n", il, id); + model.default_layer_device[il] = id; + loaded_sum += layer_sizes[il++]; + } + break; + } + } + last = il; + } + } + if (fit && required_mem > available_mem) { + auto buft = ggml_backend_cpu_buffer_type(); + if (model.split_mode == LLAMA_SPLIT_MODE_GRAPH || model.split_mode == LLAMA_SPLIT_MODE_ATTN) { + auto cur_mem = required_mem; + int n_override = 0; + for (int il = 0; il < n_layer; ++il) { + bool has_experts = false; + if (experts[il].down) { + cur_mem -= ggml_nbytes(experts[il].down); + has_experts = true; + } + if (experts[il].up) { + cur_mem -= ggml_nbytes(experts[il].up); + has_experts = true; + } + if (experts[il].gate) { + cur_mem -= ggml_nbytes(experts[il].gate); + has_experts = true; + } + if (has_experts) { + LLAMA_LOG_INFO("Adding experts CPU overrides for layer %d\n", il); + std::string pattern = "blk\\." + std::to_string(il) + "\\.(ffn_(up|down|gate|gate_up)_exps\\.weight)"; + auto & o = overrides.emplace_back(); + o.pattern = strdup(pattern.c_str()); + o.buft = buft; + ++n_override; + } + if (cur_mem <= available_mem - max_compute) { + LLAMA_LOG_INFO("Estimated memory use for split mode `graph' is %zu MiB after adding %d overrides, which is less than available memory of %zu MiB\n", + cur_mem/(1024*1024), n_override, size_t(available_mem - max_compute)/(1024*1024)); + break; + } + } + if (cur_mem > available_mem) { + if (n_override > 0) { + LLAMA_LOG_ERROR("Required memory %zu MiB is still greater than available memory %zu MiB after overriding all MoE tensors to CPU\n", + cur_mem/(1024*1024), available_mem/(1024*1024)); + } else { + LLAMA_LOG_ERROR("No MoE tensors in model and required memory %zu MiB greater than available memory %zu MiB.\n", + cur_mem/(1024*1024), available_mem/(1024*1024)); + LLAMA_LOG_ERROR("Note: auto-fit currently only works for MoE models\n"); + } + throw std::runtime_error("Unable to auto-fit model"); + } + int id = model.default_layer_device[n_layer]; + if (device_mem[id] <= layer_sizes[n_layer]) { + LLAMA_LOG_ERROR("Not enough memory in device %d to offload the output layer\n", id); + throw std::runtime_error("Unable to auto-fit model"); + } + device_mem[id] -= layer_sizes[id]; + float sum = 0; + for (int id = 0; id < int(model.splits.size()); ++id) { + sum += device_mem[id]; + model.splits[id] = sum; + } + if (sum > 0) { + LLAMA_LOG_INFO("=== Adjusted device splits for split mode graph:\n"); + float last_split = 0; + for (int id = 0; id < int(model.splits.size()); ++id) { + model.splits[id] /= sum; + LLAMA_LOG_INFO("Device %2d: %zu MiB -> split = %g\n", id, device_mem[id]/(1024*1024), model.splits[id] - last_split); + last_split = model.splits[id]; + } + } else { + throw std::runtime_error("Unable to auto-fit model"); + } + } else { + for (int id = 0; id < device_count; ++id) { + size_t mem = 0; + for (int il = 0; il <= n_layer; ++il) { + if (model.default_layer_device[il] == id) { + mem += layer_sizes[il]; + } + } + size_t cur_mem = mem; + int n_override = 0; + if (cur_mem > device_mem[id]) { + for (int il = n_layer-1; il >= 0; --il) { + if (model.default_layer_device[il] != id) continue; + bool has_experts = false; + if (experts[il].down) { + cur_mem -= ggml_nbytes(experts[il].down); + has_experts = true; + } + if (experts[il].up) { + cur_mem -= ggml_nbytes(experts[il].up); + has_experts = true; + } + if (experts[il].gate) { + cur_mem -= ggml_nbytes(experts[il].gate); + has_experts = true; + } + if (has_experts) { + LLAMA_LOG_INFO("Adding experts CPU overrides for layer %d in device %d\n", il, id); + std::string pattern = "blk\\." + std::to_string(il) + "\\.(ffn_(up|down|gate|gate_up)_exps\\.weight)"; + auto & o = overrides.emplace_back(); + o.pattern = strdup(pattern.c_str()); + o.buft = buft; + ++n_override; + } + if (cur_mem <= device_mem[id]) { + LLAMA_LOG_INFO("Memory use in device %d is %zu MiB after adding %d overrides, which is less than available memory of %zu MiB\n", + id, cur_mem/(1024*1024) ,n_override, device_mem[id]/(1024*1024)); + break; + } + } + } + if (cur_mem > device_mem[id]) { + if (n_override > 0) { + LLAMA_LOG_ERROR("Required memory %zu MiB in device %d is still greater than available memory %zu MiB after overriding all MoE tensors to CPU\n", + cur_mem/(1024*1024), id, device_mem[id]/(1024*1024)); + } else { + LLAMA_LOG_ERROR("No MoE tensors in layers assigned to device %d, and required memory %zu MiB greater than available memory %zu MiB.\n", + id, cur_mem/(1024*1024), available_mem/(1024*1024)); + } + throw std::runtime_error("Unable to auto-fit model"); + } } } - last = il; } - //for (int i = i_gpu_start; i < n_layer; ++i) { - // int layer_gpu = std::upper_bound(model.splits.begin(), model.splits.begin() + device_count, float(i - i_gpu_start)/act_gpu_layers) - model.splits.begin(); - // model.default_layer_device[i] = model.devices[layer_gpu]; - //} - //if (n_gpu_layers > n_layer) { - // int layer_gpu = std::upper_bound(model.splits.begin(), model.splits.begin() + device_count, float(act_gpu_layers - 1)/act_gpu_layers) - model.splits.begin(); - // model.default_layer_device[n_layer] = model.devices[layer_gpu]; - //} } // assign the repeating layers to the devices according to the splits if (split_mode == LLAMA_SPLIT_MODE_LAYER) { @@ -2334,7 +2534,7 @@ static bool llm_load_tensors( int layer_gpu = std::upper_bound(model.splits.begin(), model.splits.begin() + device_count, float(i - i_gpu_start)/act_gpu_layers) - model.splits.begin(); model.buft_layer[i] = { split_buft, llama_default_buffer_type_offload(model, model.devices[layer_gpu]) }; - printf("Layer %d: assigning buft_layer to GPU %d\n", i, layer_gpu); + LLAMA_LOG_INFO("Layer %d: assigning buft_layer to GPU %d\n", i, layer_gpu); } else { model.buft_layer[i] = { split_buft, buft_layer }; } @@ -2350,6 +2550,13 @@ static bool llm_load_tensors( } } + if (!overrides.empty()) { + auto & last = overrides.emplace_back(); + last.pattern = nullptr; + ml.tensor_buft_overrides = overrides.data(); + model.tensor_overrides = true; + } + auto cth = create_tensors_helper_interface::instance(ml, model); auto ctx_size = cth->get_ctx_size(); @@ -2524,7 +2731,7 @@ static bool llm_load_tensors( if (iqk_modify_tensor(it.second)) ++n_modified; } } - if (n_modified > 0) printf("============ Modified %d tensors\n", n_modified); + if (n_modified > 0) LLAMA_LOG_INFO("============ Modified %d tensors\n", n_modified); } if (validate_quants) { @@ -2550,7 +2757,7 @@ static bool llm_load_tensors( if (it.second->type != orig_type) ++n_repacked; } } - if (n_repacked > 0) printf("============ Repacked %d tensors\n", n_repacked); + if (n_repacked > 0) LLAMA_LOG_INFO("============ Repacked %d tensors\n", n_repacked); } if (model.arch == LLM_ARCH_BITNET) { @@ -2579,6 +2786,13 @@ static bool llm_load_tensors( } } + if (!overrides.empty()) { + for (auto & o : overrides) { + if (o.pattern) free(const_cast(o.pattern)); + } + ml.tensor_buft_overrides = nullptr; + } + // loading time will be recalculate after the first eval, so // we take page faults deferred by mmap() into consideration model.t_load_us = ggml_time_us() - model.t_start_us; @@ -2642,8 +2856,8 @@ static int llama_model_load(const std::string & fname, llama_model & model, llam if (!llm_load_tensors( ml, model, params.n_gpu_layers, params.mla, params.split_mode, params.main_gpu, params.max_gpu, params.tensor_split, - params.type_k, params.type_v, params.max_ctx_size, params.n_seq_max, params.n_ubatch, params.amb, params.flash_attn, - params.use_mlock, params.validate_quants, params.mtp, params.dry_run, + params.type_k, params.type_v, params.max_ctx_size, params.n_seq_max, params.n_ubatch, params.amb, params.fit_margin, params.flash_attn, + params.use_mlock, params.validate_quants, params.mtp, params.fit, params.dry_run, params.progress_callback, params.progress_callback_user_data )) { return -2; @@ -4612,6 +4826,8 @@ struct llama_model_params llama_model_default_params() { /*.n_seq_max =*/ 1, /*.n_ubatch =*/ 512, /*.amb =*/ 0, + /*.fit_margin =*/ 0, + /*.fit =*/ false, /*.tensor_split =*/ nullptr, /*.rpc_servers =*/ nullptr, /*.progress_callback =*/ nullptr,