furyhawk/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-07-21 10:15:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Auto-fit offloaded tensors to available VRAM (MoE models) (#1501)

Browse Source 
* WIP: automatically fit model in available VRAM

* WIP

* This seems pretty solid
This commit is contained in:
Kawrakow
2026-03-25 07:29:29 +01:00
committed by GitHub
parent 5451b149d4
commit 86f4f516e5
6 changed files with 311 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/common.cpp
+21
View File
@@ -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);
common/common.h
+2
View File
@@ -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
include/llama.h
+2
View File
@@ -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;
src/llama-arch.h
+22 -22
View File
@@ -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,
};
src/llama-model.cpp
+6 -1
View File
@@ -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;
}
src/llama.cpp
+258 -42
View File
@@ -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<std::vector<double>, 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<expert_tensors> & experts) {
int n_layer = model.hparams.n_layer;
std::vector<double> result(n_layer+1, 0);
std::vector<double> compute(n_layer+1, 0);
@@ -2013,6 +2019,7 @@ static std::pair<std::vector<double>, double> get_layer_sizes(const llama_model_
if (has_mla) {
mla_tensors.resize(n_layer);
}
experts.resize(n_layer);
std::vector<size_t> ffn_exps(n_layer, 0), ffn_shexp(n_layer, 0), ffn_dense(n_layer, 0);
std::vector<size_t> attn(n_layer, 0);
std::vector<bool> has_layer_norm(n_layer, false);
@@ -2035,23 +2042,23 @@ static std::pair<std::vector<double>, 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<std::vector<double>, 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<std::vector<double>, 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<size_t> 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<float> 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<int32_t>(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<llama_model_tensor_buft_override> overrides;
if (device_count > 0) {
std::vector<expert_tensors> 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);
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 {
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++];
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");
}
break;
}
}
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<char*>(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,