diff --git a/common/arg.cpp b/common/arg.cpp index 3acfbab6e544..3d2fd10294b3 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2904,6 +2904,45 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } } ).set_env("LLAMA_ARG_TENSOR_SPLIT")); + add_opt(common_arg( + {"-as", "--attn-split"}, "N0,N1,N2,...", + "fraction of the attention heads to give each GPU under --split-mode tensor, comma-separated, " + "e.g. 3,1. Only the heads move, the rest of the model still follows --tensor-split. Each share is " + "rounded down to whole heads, so a ratio the head count cannot express is not matched exactly. " + "Useful when the GPUs differ in host bandwidth or in speed (default: follow --tensor-split)", + [](common_params & params, const std::string & value) { + const std::regex regex{ R"([,/]+)" }; + std::sregex_token_iterator it{ value.begin(), value.end(), regex, -1 }; + std::vector split_arg{ it, {} }; + if (split_arg.size() > llama_max_devices()) { + throw std::invalid_argument( + string_format("got %zu input configs, but system only has %zu devices", split_arg.size(), llama_max_devices()) + ); + } + float sum = 0.0f; + for (size_t i = 0; i < llama_max_devices(); ++i) { + float share = 0.0f; + if (i < split_arg.size()) { + size_t n_read = 0; + try { + share = std::stof(split_arg[i], &n_read); + } catch (const std::exception &) { + n_read = 0; + } + if (n_read != split_arg[i].size() || !std::isfinite(share) || share < 0.0f) { + throw std::invalid_argument( + string_format("invalid attention split share '%s'", split_arg[i].c_str()) + ); + } + } + params.attn_split[i] = share; + sum += share; + } + if (sum <= 0.0f) { + throw std::invalid_argument("the attention split shares must add up to more than zero"); + } + } + ).set_env("LLAMA_ARG_ATTN_SPLIT")); add_opt(common_arg( {"-mg", "--main-gpu"}, "INDEX", string_format("the GPU to use for the model (with split-mode = none), or for intermediate results and KV (with split-mode = row) (default: %d)", params.main_gpu), diff --git a/common/common.cpp b/common/common.cpp index 05af52737426..966dcd077218 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1690,6 +1690,8 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.load_mode = params.load_mode; mparams.lazy_mode = params.lazy_mode; mparams.tensor_split = params.tensor_split; + mparams.attn_split = std::any_of(params.attn_split, params.attn_split + llama_max_devices(), + [](float f) { return f != 0.0f; }) ? params.attn_split : nullptr; mparams.check_tensors = params.check_tensors; mparams.use_extra_bufts = !params.no_extra_bufts; mparams.no_host = params.no_host; diff --git a/common/common.h b/common/common.h index 8edd79f3d053..035261b2ee06 100644 --- a/common/common.h +++ b/common/common.h @@ -486,6 +486,7 @@ struct common_params { int32_t n_gpu_layers = -1; // number of layers to store in VRAM, -1 is auto, <= -2 is all int32_t main_gpu = 0; // the GPU that is used for scratch and small tensors float tensor_split[128] = {0}; // how split tensors should be distributed across GPUs + float attn_split[128] = {0}; // how attention heads should be distributed across GPUs bool fit_params = true; // whether to fit unset model/context parameters to free device memory bool fit_params_print = false; // print the estimated required memory to run the model int32_t fit_params_min_ctx = 4096; // minimum context size to set when trying to reduce memory use diff --git a/include/llama.h b/include/llama.h index 386f0fd7ad16..4a1f379d9ee2 100644 --- a/include/llama.h +++ b/include/llama.h @@ -329,6 +329,10 @@ extern "C" { // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() const float * tensor_split; + // proportion of the attention heads to give each GPU under split mode tensor, size: llama_max_devices() + // NULL follows tensor_split. Only the heads move - the rest of the model still follows tensor_split + const float * attn_split; + // Called with a progress value between 0.0 and 1.0. Pass NULL to disable. // If the provided progress_callback returns true, model loading continues. // If it returns false, model loading is immediately aborted. diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 66ae3b617fb6..652c72536f83 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -828,6 +828,18 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str const int64_t ne_axis = unfolded ? tensor->ne[cache_copy.axis]*unit : tensor->ne[split_state.axis]; const int64_t blck_size = ggml_blck_size(tc.tensor_axis_0->type); const float * tensor_split = ud->model->tensor_split(); + + // The share of the attention heads decides how much cache a device gets and how much attention + // work it does, neither of which has to follow the memory split. Only a split that counts heads + // follows it, and those are the ones measured against attn_output.weight - a linear-attention + // layer measures against ssm_out.weight and keeps tensor_split. + if (ud->model->attn_split() != nullptr) { + const std::string attn_out_name = "blk." + std::to_string(tc.il) + ".attn_output.weight"; + const ggml_tensor * attn_out = ud->model->get_tensor(attn_out_name.c_str()); + if (attn_out != nullptr && tc.tensor_axis_0 == attn_out) { + tensor_split = ud->model->attn_split(); + } + } std::vector tensor_split_scan; tensor_split_scan.reserve(ud->n_devices); for (size_t j = 0; j < ud->n_devices; j++) { @@ -1221,6 +1233,7 @@ struct llama_model::impl { bool has_tensor_overrides; std::vector tensor_split_owned; + std::vector attn_split_owned; }; llama_model::llama_model(const llama_model_params & params) : params(params), pimpl(std::make_unique()) { @@ -1230,6 +1243,23 @@ llama_model::llama_model(const llama_model_params & params) : params(params), pi pimpl->tensor_split_owned.assign(params.tensor_split, params.tensor_split + llama_max_devices()); this->params.tensor_split = pimpl->tensor_split_owned.data(); } + if (params.attn_split != nullptr) { + // the shares reach a cumulative division and a conversion to whole heads, so a negative, a + // non-finite or an all-zero set has no meaning here - fall back to the tensor split + float sum = 0.0f; + bool ok = true; + for (size_t i = 0; i < llama_max_devices(); i++) { + ok = ok && std::isfinite(params.attn_split[i]) && params.attn_split[i] >= 0.0f; + sum += params.attn_split[i]; + } + if (!ok || !(sum > 0.0f)) { + LLAMA_LOG_WARN("%s: the attention split is not a set of non-negative shares; ignoring it\n", __func__); + this->params.attn_split = nullptr; + } else { + pimpl->attn_split_owned.assign(params.attn_split, params.attn_split + llama_max_devices()); + this->params.attn_split = pimpl->attn_split_owned.data(); + } + } pimpl->has_tensor_overrides = params.tensor_buft_overrides && params.tensor_buft_overrides[0].pattern; } @@ -1941,6 +1971,10 @@ const float * llama_model::tensor_split() const { return params.tensor_split; } +const float * llama_model::attn_split() const { + return params.attn_split; +} + uint32_t llama_model::n_gpu_layers() const { // note: plus 1 for the "output" layer return params.n_gpu_layers >= 0 ? params.n_gpu_layers : hparams.n_layer_all + 1; @@ -2837,6 +2871,7 @@ llama_model_params llama_model_default_params() { /*.lazy_mode =*/ LLAMA_LAZY_MODE_AUTO, /*.main_gpu =*/ 0, /*.tensor_split =*/ nullptr, + /*.attn_split =*/ nullptr, /*.progress_callback =*/ nullptr, /*.progress_callback_user_data =*/ nullptr, /*.kv_overrides =*/ nullptr, diff --git a/src/llama-model.h b/src/llama-model.h index 2cb74ea23273..76777c2f9e4c 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -729,6 +729,7 @@ struct llama_model { size_t n_tensors() const; size_t n_devices() const; const float * tensor_split() const; + const float * attn_split() const; uint32_t n_gpu_layers() const; llama_split_mode split_mode() const; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index cf193bf168e6..ff7d16191b3a 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -400,10 +400,11 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) } // with offload_kqv=false the cache lives in host memory -// n_seq_max > 1 gives the cache one stream per sequence +// n_seq_max > 1 gives the cache one stream per sequence, attn_split gives the heads their own share struct kv_config { - bool offload_kqv = true; - uint32_t n_seq_max = 1; + bool offload_kqv = true; + uint32_t n_seq_max = 1; + std::vector attn_split; }; static std::pair get_model_and_ctx( @@ -416,6 +417,11 @@ static std::pair get_model_and_ctx( devs_copy.push_back(nullptr); model_params.devices = devs_copy.data(); model_params.split_mode = split_mode; + std::vector attn_split = kvc.attn_split; + if (!attn_split.empty()) { + attn_split.resize(llama_max_devices(), 0.0f); + model_params.attn_split = attn_split.data(); + } llama_context_params ctx_params = llama_context_default_params(); ctx_params.n_ctx = 0; @@ -1459,6 +1465,19 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const in kvc_host_streams.n_seq_max = 2; dev_configs.emplace_back(devices_meta, "Meta -nkvo -np 2", LLAMA_SPLIT_MODE_TENSOR, kvc_host_streams); + // the same, with all attention heads on the first device + if (devices_meta.size() > 1) { + kv_config kvc_attn = kvc_host; + kvc_attn.attn_split.assign(devices_meta.size(), 0.0f); + kvc_attn.attn_split[0] = 1.0f; + dev_configs.emplace_back(devices_meta, "Meta -nkvo -as", LLAMA_SPLIT_MODE_TENSOR, kvc_attn); + + // a custom head split must also hold across streams + kv_config kvc_attn_streams = kvc_attn; + kvc_attn_streams.n_seq_max = 2; + dev_configs.emplace_back(devices_meta, "Meta -nkvo -as -np 2", LLAMA_SPLIT_MODE_TENSOR, kvc_attn_streams); + } + for (const device_config & dc : dev_configs) { max_device_label_length = std::max(max_device_label_length, dc.label.length()); }