diff --git a/common/arg.cpp b/common/arg.cpp index 806f70578602..1eead7c308b6 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1769,6 +1769,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.flash_attn_native_quants = value; } ).set_env("LLAMA_ARG_FLASH_ATTN_NATIVE_QUANTS")); + add_opt(common_arg( + {"--split-overlap"}, + {"--no-split-overlap"}, + "with --split-mode tensor: alternate ubatches of at least 64 tokens between two schedulers, so that the " + "AllReduce of one overlaps with the compute of the other. Each scheduler has its own compute buffers: " + "half the --ubatch-size takes the same memory and is usually the fastest choice (default: disabled)", + [](common_params & params, bool value) { + params.split_overlap = value; + } + ).set_env("LLAMA_ARG_SPLIT_OVERLAP")); add_opt(common_arg( {"--cache-idle-slots"}, {"--no-cache-idle-slots"}, diff --git a/common/common.cpp b/common/common.cpp index 879d32a4fa39..8d58f1d40d8c 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1742,6 +1742,7 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.swa_full = params.swa_full; cparams.kv_unified = params.kv_unified; cparams.flash_attn_native_quants = params.flash_attn_native_quants; + cparams.split_overlap = params.split_overlap; cparams.type_k = params.cache_type_k; cparams.type_v = params.cache_type_v; diff --git a/common/common.h b/common/common.h index 0959f5ca2919..0b98923331a0 100644 --- a/common/common.h +++ b/common/common.h @@ -507,6 +507,7 @@ struct common_params { enum llama_flash_attn_type flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO; // whether to use Flash Attention bool flash_attn_native_quants = false; // read a quantized K/V cache in place instead of casting it to F16 + bool split_overlap = false; // split mode tensor: overlap the AllReduce of a ubatch with the compute of the next struct common_params_sampling sampling; struct common_params_speculative speculative; diff --git a/docs/multi-gpu.md b/docs/multi-gpu.md index de7491fcc0e4..c6514373315e 100644 --- a/docs/multi-gpu.md +++ b/docs/multi-gpu.md @@ -37,6 +37,7 @@ Set with `--split-mode` / `-sm`. | `-sm` | `--split-mode` | `none` \| `layer` \| `tensor` | `layer` | See modes above. | | `-ts` | `--tensor-split` | comma-separated proportions, e.g. `3,1` | mode-dependent | How much of the model goes to each GPU. If omitted, `layer`/`row` use automatic splitting proportional to memory, while `tensor` splits tensor segments evenly. With `3,1` on two GPUs, GPU 0 gets 75 %, GPU 1 gets 25 %. The values follow the order in `--device`. | | `-as` | `--attn-split` | comma-separated proportions, e.g. `1,3` | follows `-ts` | With `--split-mode tensor`: the share of the attention heads per GPU. The heads take their weights and their KV cache along, the rest of the model follows `--tensor-split`. With GPUs of different speed, give the faster one more of `--tensor-split` and move KV heads to the other one to keep both in memory. The model load logs the KV heads each GPU got. | +| | `--split-overlap` | - | off | With `--split-mode tensor`: alternate ubatches between two schedulers, so that the AllReduce of one overlaps with the compute of the other. See recipe 8. | | `-mg` | `--main-gpu` | integer device index | `0` | The single GPU used in `--split-mode none`. | | `-ngl` | `--n-gpu-layers` / `--gpu-layers` | integer \| `auto` \| `all` | `auto` | Maximum number of layers to keep in VRAM. Use `999` or `all` to push everything possible to the GPUs. | | `-dev` | `--device` | comma-separated device names, or `none` | auto | Restrict which devices llama.cpp may use. See `--list-devices` for names. | @@ -135,6 +136,33 @@ Environment variables: The copy engines need 16 MiB of VRAM and 8 MiB of pinned memory per GPU and per context that has large reductions. They are allocated with the first large reduction. An MTP draft context processes the prompt batches as well, so it has its own buffers: 32 MiB of VRAM per GPU in total. +### 8. Overlapping the AllReduce with compute (`--split-overlap`) + +```bash +llama-server -m model.gguf -sm tensor -ub 256 --split-overlap +``` + +In `tensor` mode every layer ends with an AllReduce, and the GPUs wait for it. On a slow link this can take a large share of prompt processing: about 40% with an RTX 3060 on PCIe 3.0 x4. + +With `--split-overlap`, consecutive ubatches of at least 64 tokens alternate between two schedulers, also across calls. Each scheduler has its own backends, streams and compute buffers. A subgraph waits on each GPU for the subgraphs of the other ubatch that last read or wrote the same state tensors (KV cache, recurrent state). A subgraph without state, such as an FFN, does not wait. So while one ubatch transfers a layer's result, the other computes. + +It needs all layers, the KV cache and the experts on the GPUs (the CPU parts of a graph would run outside that order), and devices with events. + +The output is the same as without the option. Decode (small ubatches) is not affected. + +Notes: + +- Each scheduler reserves the compute buffers for a full ubatch. With half the `--ubatch-size`, the memory stays the same, and two ubatches are in flight at once. On the GPUs below, 256 per scheduler was faster than 128, 384 or 512. +- Each scheduler captures its own CUDA graphs, which takes some more VRAM at run time (about 50 MiB here). +- The overlap works best with the copy-engine AllReduce of recipe 7: NCCL kernels hold an SM each, and the compute of the other ubatch then runs slower. + +Qwen3.8-27B on an RTX 4070 (PCIe 4.0 x16) + RTX 3060 (PCIe 3.0 x4), `-ts 0.625,0.375 -as 1,3`, 262144 tokens of q5_1 cache in 2 slots, with MTP, prompt processing in t/s: + +| prompt | `-ub 512` | `-ub 256 --split-overlap` | +|---|--:|--:| +| 10k tokens | 733 | 947 | +| 93k tokens | 570 | 734 | + --- ## Troubleshooting diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 7a65d8121b1c..273ee39c4ebf 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -430,6 +430,12 @@ extern "C" { GGML_API ggml_backend_dev_t ggml_backend_meta_device( ggml_backend_dev_t * devs, size_t n_devs, ggml_backend_meta_get_split_state_t get_split_state, void * get_split_state_ud); + // let two backends of the same meta device compute consecutive graphs as lanes: each lane runs on its own streams, + // and a subgraph waits for the subgraphs of the other lane that last touched the same state tensors (KV cache, + // recurrent state), so the AllReduce of one lane overlaps with the compute of the other. The graphs must not + // share intermediate buffers (one scheduler per lane). Returns false if a device has no events. + GGML_API bool ggml_backend_meta_share_lanes(ggml_backend_t meta_backend_a, ggml_backend_t meta_backend_b); + // // Utils // diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 679d5216d255..fbaf913c5be7 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -6,6 +6,7 @@ #include "ggml-cpp.h" #include +#include #include #include #include @@ -1836,6 +1837,20 @@ struct ggml_backend_meta_context { void * comm_ctx = nullptr; ggml_backend_comm_allreduce_tensor_t comm_allreduce = nullptr; + // Lanes: meta backends of one device that compute consecutive graphs on their own streams, so that the AllReduce + // of one graph overlaps with the compute of the other, see ggml_backend_meta_share_lanes(). + struct lane_group { + ggml_backend_meta_context * ctx[2] = { nullptr, nullptr }; + }; + std::shared_ptr lanes; + std::vector> lane_events; // [subgraph][simple backend], recorded once it is computed + + // the state tensors (KV cache, recurrent state: tensors of buffers that are neither weights nor compute buffers) + // that each subgraph of the current graph reads or writes, by the tensor at the root of their views + std::vector> lane_touch; + // the subgraph (index into lane_events) of this lane that last touched a state tensor, over all its graphs + std::unordered_map lane_last_touch; + ggml_backend_meta_context(ggml_backend_dev_t meta_dev, const char * params) { const size_t n_devs = ggml_backend_meta_dev_n_devs(meta_dev); n_reduce_steps = std::ceil(std::log2(n_devs)); @@ -1870,6 +1885,18 @@ struct ggml_backend_meta_context { } ~ggml_backend_meta_context() { + for (auto & events : lane_events) { + for (ggml_backend_event_t event : events) { + ggml_backend_event_free(event); + } + } + if (lanes) { + for (auto & ctx : lanes->ctx) { + if (ctx == this) { + ctx = nullptr; + } + } + } if (comm_ctx != nullptr) { ggml_backend_comm_free_t comm_free = (ggml_backend_comm_free_t) ggml_backend_reg_get_proc_address( ggml_backend_dev_backend_reg(ggml_backend_get_device(backend_configs[0].backend)), "ggml_backend_comm_free"); @@ -2259,6 +2286,41 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, backend_ctx->uid = cgraph->uid; backend_ctx->n_subgraphs = n_subgraphs; + if (backend_ctx->lanes) { + auto state_root = [](const ggml_tensor * t) -> const ggml_tensor * { + while (t->view_src != nullptr) { + t = t->view_src; + } + if (t->buffer == nullptr || !ggml_backend_buffer_is_meta(t->buffer)) { + return nullptr; + } + const enum ggml_backend_buffer_usage usage = ggml_backend_buffer_get_usage(t->buffer); + return usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS || usage == GGML_BACKEND_BUFFER_USAGE_COMPUTE ? nullptr : t; + }; + backend_ctx->lane_touch.assign(n_subgraphs, {}); + for (size_t i_graph = 0; i_graph < n_subgraphs; i_graph++) { + auto & touch = backend_ctx->lane_touch[i_graph]; + const auto & cfg = backend_ctx->backend_configs[0].cgraphs; + const int i_start_g = cfg[i_graph].offset; + const int i_stop_g = i_graph + 1 < n_subgraphs ? cfg[i_graph + 1].offset : cgraph->n_nodes; + for (int i_node = i_start_g; i_node < i_stop_g; i_node++) { + const ggml_tensor * node = cgraph->nodes[i_node]; + if (const ggml_tensor * root = state_root(node)) { + touch.push_back(root); + } + for (int k = 0; k < GGML_MAX_SRC; k++) { + if (node->src[k] != nullptr) { + if (const ggml_tensor * root = state_root(node->src[k])) { + touch.push_back(root); + } + } + } + } + std::sort(touch.begin(), touch.end()); + touch.erase(std::unique(touch.begin(), touch.end()), touch.end()); + } + } + if (max_tmp_size > backend_ctx->max_tmp_size) { for (size_t j = 0; j < n_backends; j++) { auto & bcj = backend_ctx->backend_configs[j]; @@ -2466,13 +2528,56 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, }; + // With lanes, a subgraph waits on each device for the subgraphs of the other lane that last read or wrote the + // state tensors it reads or writes (KV cache, recurrent state), whichever graph and subgraph index that was: + // a layer may read the KV of another layer (shared KV), or read its state a subgraph before it writes it. The + // AllReduce after a subgraph is not waited for, it overlaps with the compute of the other lane. An event that + // the other lane recorded again since is later in its stream, so waiting for it waits for at least as much. + ggml_backend_meta_context * lane_other = nullptr; + if (backend_ctx->lanes) { + lane_other = backend_ctx->lanes->ctx[0] == backend_ctx ? backend_ctx->lanes->ctx[1] : backend_ctx->lanes->ctx[0]; + + while (backend_ctx->lane_events.size() < backend_ctx->n_subgraphs) { + std::vector events(n_backends); + for (size_t j = 0; j < n_backends; j++) { + events[j] = ggml_backend_event_new(ggml_backend_get_device(backend_ctx->backend_configs[j].backend)); + GGML_ASSERT(events[j] != nullptr); + } + backend_ctx->lane_events.push_back(std::move(events)); + } + } + std::vector lane_waits; + for (size_t i = 0; i < backend_ctx->n_subgraphs; i++) { + lane_waits.clear(); + if (lane_other != nullptr) { + for (const ggml_tensor * t : backend_ctx->lane_touch[i]) { + auto it = lane_other->lane_last_touch.find(t); + if (it != lane_other->lane_last_touch.end()) { + lane_waits.push_back(it->second); + } + } + std::sort(lane_waits.begin(), lane_waits.end()); + lane_waits.erase(std::unique(lane_waits.begin(), lane_waits.end()), lane_waits.end()); + } + for (size_t j = 0; j < n_backends; j++) { auto & bcj = backend_ctx->backend_configs[j]; + for (const size_t i_other : lane_waits) { + ggml_backend_event_wait(bcj.backend, lane_other->lane_events[i_other][j]); + } const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, bcj.cgraphs[i].cgraph_main); if (status != GGML_STATUS_SUCCESS) { return status; } + if (backend_ctx->lanes) { + ggml_backend_event_record(backend_ctx->lane_events[i][j], bcj.backend); + } + } + if (backend_ctx->lanes) { + for (const ggml_tensor * t : backend_ctx->lane_touch[i]) { + backend_ctx->lane_last_touch[t] = i; + } } if (n_backends > 1 && i < backend_ctx->n_subgraphs - 1) { @@ -2539,6 +2644,29 @@ size_t ggml_backend_meta_n_backends(ggml_backend_t meta_backend) { return backend_ctx->backend_configs.size(); } +bool ggml_backend_meta_share_lanes(ggml_backend_t meta_backend_a, ggml_backend_t meta_backend_b) { + GGML_ASSERT(ggml_backend_is_meta(meta_backend_a) && ggml_backend_is_meta(meta_backend_b)); + GGML_ASSERT(meta_backend_a->device == meta_backend_b->device); + auto * ctx_a = (ggml_backend_meta_context *) meta_backend_a->context; + auto * ctx_b = (ggml_backend_meta_context *) meta_backend_b->context; + GGML_ASSERT(!ctx_a->lanes && !ctx_b->lanes); + + // the lanes are ordered with events on each device + for (const auto & bc : ctx_a->backend_configs) { + ggml_backend_dev_props props; + ggml_backend_dev_get_props(ggml_backend_get_device(bc.backend), &props); + if (!props.caps.events) { + return false; + } + } + + ctx_a->lanes = std::make_shared(); + ctx_b->lanes = ctx_a->lanes; + ctx_a->lanes->ctx[0] = ctx_a; + ctx_a->lanes->ctx[1] = ctx_b; + return true; +} + ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, size_t index) { GGML_ASSERT(ggml_backend_is_meta(meta_backend)); const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context; diff --git a/include/llama.h b/include/llama.h index 882ec438070d..7f325eb4b1f5 100644 --- a/include/llama.h +++ b/include/llama.h @@ -422,6 +422,9 @@ extern "C" { bool flash_attn_native_quants; // let Flash Attention read a quantized K/V cache in place instead of // casting it to F16 first; removes the transient F16 copy of the // attention window. Where no native loader exists the F16 path is used. + bool split_overlap; // split mode tensor: alternate large ubatches between two schedulers, so that + // the AllReduce of one overlaps with the compute of the other; each has its own + // compute buffers, so a ubatch of half the size takes the same memory // [EXPERIMENTAL] // backend sampler chain configuration (make sure the caller keeps the sampler chains alive) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 57cbb1d6778e..bf417af82589 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -206,6 +206,7 @@ llama_context::llama_context( cparams.embeddings_nextn_masked = false; cparams.offload_kqv = params.offload_kqv; cparams.flash_attn_native_quants = params.flash_attn_native_quants; + cparams.split_overlap = params.split_overlap; cparams.no_perf = params.no_perf; cparams.warmup = false; @@ -440,6 +441,28 @@ llama_context::llama_context( } } + // a second meta backend for the lanes, see llama_context::sched_lane + // a draft context only gets small ubatches, it stays on one lane + // the CPU parts of a graph run synchronously when the scheduler gets to them, outside of the order between + // the lanes: the KV cache, all layers and the experts must be on the devices + if (cparams.split_overlap && cparams.ctx_type == LLAMA_CONTEXT_TYPE_DEFAULT) { + const bool full_offload = model.n_gpu_layers() > model.hparams.n_layer_all; + if (model.devices.size() == 1 && model.devices[0].is_meta && backends.size() == 1 && cparams.n_ubatch >= lanes_min_tokens && + cparams.offload_kqv && full_offload && !model.has_tensor_overrides() && cparams.moe_cache_size == 0) { + backend_lane1.reset(ggml_backend_dev_init(model.devices[0].dev, nullptr)); + if (backend_lane1 && ggml_backend_meta_share_lanes(backends[0].get(), backend_lane1.get())) { + lanes = true; + LLAMA_LOG_INFO("%s: split overlap: ubatches of at least %u tokens alternate between two schedulers\n", __func__, lanes_min_tokens); + } else { + backend_lane1.reset(); + LLAMA_LOG_WARN("%s: split overlap needs devices with events, it is not used\n", __func__); + } + } else { + LLAMA_LOG_WARN("%s: split overlap needs split mode tensor, all layers and the KV cache on the devices and a ubatch of at least %u tokens, it is not used\n", + __func__, lanes_min_tokens); + } + } + // add CPU backend backend_cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); if (backend_cpu == nullptr) { @@ -512,6 +535,11 @@ llama_context::llama_context( backend_buf_exp_size.push_back(0); } + if (lanes) { + lane_other.backend_ptrs = backend_ptrs; + lane_other.backend_ptrs[0] = backend_lane1.get(); + } + LLAMA_LOG_DEBUG("%s: backend_ptrs.size() = %zu\n", __func__, backend_ptrs.size()); // TODO: move these checks to ggml_backend_sched @@ -714,6 +742,23 @@ void llama_context::sched_new(size_t max_nodes, bool parallel) { ggml_backend_sched_set_realloc_callback(sched.get(), sched_graph_realloc, this); } +void llama_context::lane_switch(int lane) { + if (!lanes || lane == lane_active) { + return; + } + std::swap(sched, lane_other.sched); + std::swap(gf_res_prev, lane_other.gf_res_prev); + std::swap(gf_res_prev_active, lane_other.gf_res_prev_active); + std::swap(backend_ptrs, lane_other.backend_ptrs); + lane_active = lane; + + // The meta buffers that both lanes use (weights, KV cache) map the views of a graph to the simple tensors of each + // device in a rotating pair of containers: a rebuild in the meta backend clears the container of the previous + // graph. After the other lane has computed, the views of this lane's previous graph may be gone, so the graph is + // built and allocated again instead of reused. + gf_res_prev_active = nullptr; +} + void llama_context::sched_reserve() { if (!sched_need_reserve) { return; @@ -721,7 +766,17 @@ void llama_context::sched_reserve() { sched_need_reserve = false; - LLAMA_LOG_INFO("%s: reserving ...\n", __func__); + GGML_ASSERT(lane_active == 0); + sched_reserve_lane(); + if (lanes) { + lane_switch(1); + sched_reserve_lane(); + lane_switch(0); + } +} + +void llama_context::sched_reserve_lane() { + LLAMA_LOG_INFO("%s: reserving%s ...\n", __func__, lanes ? (lane_active ? " lane 1" : " lane 0") : ""); synchronize(); @@ -876,6 +931,9 @@ void llama_context::synchronize() { } ggml_backend_sched_synchronize(sched.get()); + if (lane_other.sched) { + ggml_backend_sched_synchronize(lane_other.sched.get()); + } // FIXME: if multiple single tokens are evaluated without a synchronization, // the stats will be added to the prompt evaluation stats @@ -2001,8 +2059,20 @@ int llama_context::decode(const llama_batch_ext & batch_inp) { ggml_status status; + // large ubatches alternate between the lanes, also across calls + if (lanes && ubatch.n_tokens >= lanes_min_tokens) { + lane_switch(lane_next); + lane_next ^= 1; + } else { + lane_switch(0); + } + const auto * res = process_ubatch(ubatch, ctx_type_to_graph_type(cparams.ctx_type), mctx.get(), status); + if (!res) { + lane_switch(0); + } + if (!res) { // the last ubatch failed or was aborted -> remove all positions of that ubatch from the memory module llama_pos pos_min[LLAMA_MAX_SEQ]; @@ -2166,6 +2236,8 @@ int llama_context::decode(const llama_batch_ext & batch_inp) { n_tokens_prev += ubatch.n_tokens; } while (mctx->next()); + lane_switch(0); + // set to total number of outputs in the batch, for use in llama_get_logits_ith n_outputs = n_outputs_all; @@ -2860,10 +2932,11 @@ llm_graph_cb llama_context::graph_get_cb() const { if (ubatch.n_tokens < 32 || full_offload) { if (il != -1 && (strcmp(name, "norm") == 0 || strcmp(name, "l_last") == 0)) { const auto & dev_layer = model.dev_layer(il); - for (const auto & backend : backends) { - if (ggml_backend_get_device(backend.get()) == dev_layer) { - if (ggml_backend_supports_op(backend.get(), cur)) { - ggml_backend_sched_set_tensor_backend(sched.get(), cur, backend.get()); + // the backends of the active scheduler lane + for (ggml_backend_t backend : backend_ptrs) { + if (ggml_backend_get_device(backend) == dev_layer) { + if (ggml_backend_supports_op(backend, cur)) { + ggml_backend_sched_set_tensor_backend(sched.get(), cur, backend); } } } @@ -3777,7 +3850,8 @@ llama_memory_breakdown llama_context::memory_breakdown() const { for (size_t i = 0; i < backends.size(); ++i) { ggml_backend_t backend = backends[i].get(); ggml_backend_buffer_type_t buft = ggml_backend_sched_get_buffer_type(sched.get(), backend); - ret[buft].compute += backend_buf_exp_size[i]; + // the second lane reserves the same buffers + ret[buft].compute += backend_buf_exp_size[i] * (lanes ? 2 : 1); } } else { for (const auto & backend_ptr : backends) { @@ -3785,6 +3859,12 @@ llama_memory_breakdown llama_context::memory_breakdown() const { ggml_backend_buffer_type_t buft = ggml_backend_sched_get_buffer_type(sched.get(), backend); ret[buft].compute += ggml_backend_sched_get_buffer_size(sched.get(), backend); } + if (lane_other.sched) { + for (ggml_backend_t backend : lane_other.backend_ptrs) { + ggml_backend_buffer_type_t buft = ggml_backend_sched_get_buffer_type(lane_other.sched.get(), backend); + ret[buft].compute += ggml_backend_sched_get_buffer_size(lane_other.sched.get(), backend); + } + } } return ret; } @@ -4072,6 +4152,7 @@ llama_context_params llama_context_default_params() { /*.swa_full =*/ true, /*.kv_unified =*/ false, /*.flash_attn_native_quants =*/ false, + /*.split_overlap =*/ false, /*.sampler =*/ nullptr, /*.n_sampler =*/ 0, /*.ctx_other =*/ nullptr, diff --git a/src/llama-context.h b/src/llama-context.h index b21088eaf6f9..f340d6db1381 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -281,6 +281,12 @@ struct llama_context { // create a new backend scheduler for a graph of the specified size and register the callbacks below void sched_new(size_t max_nodes, bool parallel); + // reserve the compute buffers of the active scheduler lane + void sched_reserve_lane(); + + // make a scheduler lane the active one (sched, gf_res_prev, backend_ptrs) + void lane_switch(int lane); + // ggml_backend_sched copy callback, copies only the experts used by MUL_MAT_ID and updates the MoE cache static bool sched_copy_experts(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph, void * user_data); @@ -422,6 +428,23 @@ struct llama_context { llm_graph_result * gf_res_prev_active = nullptr; + // Lanes (split_overlap, split mode tensor): consecutive large ubatches alternate between two schedulers + // whose meta backends compute on their own streams, so that the AllReduce of one ubatch overlaps with the compute + // of the other. Each lane has its own compute buffers. The inactive lane keeps its state here; outside of + // decode() lane 0 is always the active one. + struct sched_lane { + ggml_backend_sched_ptr sched; + std::array gf_res_prev; + llm_graph_result * gf_res_prev_active = nullptr; + std::vector backend_ptrs; + }; + sched_lane lane_other; + int lane_active = 0; + int lane_next = 1; // the lane of the next large ubatch + bool lanes = false; + const uint32_t lanes_min_tokens = 64; // smaller ubatches stay on lane 0 + ggml_backend_ptr backend_lane1; // the meta backend of lane 1 + // host buffer for the model output (logits and embeddings) ggml_backend_buffer_ptr buf_output; diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 6cff5e63bfcd..e1fd415ba19d 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -39,6 +39,7 @@ struct llama_cparams { bool offload_kqv; bool flash_attn; bool flash_attn_native_quants; // read a quantized K/V cache in place instead of casting it to F16 + bool split_overlap; // see llama_context_params::split_overlap bool auto_fa; bool fused_gdn_ar; // use fused gated delta net (autoregressive) bool fused_gdn_ch; // use fused gated delta net (chunked) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index d0250d5aa7d9..f6193014896e 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -516,7 +516,7 @@ static std::pair get_model_and_ctx( const std::vector & devs, const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false, const llama_model_tensor_buft_override * tensor_buft_overrides = nullptr, const size_t moe_cache_size = 0, - const std::vector & attn_split = {}) { + const std::vector & attn_split = {}, const bool split_overlap = false) { GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); llama_model_params model_params = llama_model_default_params(); model_params.progress_callback = silent_model_load_progress; @@ -535,6 +535,8 @@ static std::pair get_model_and_ctx( ctx_params.n_ctx = 0; ctx_params.n_threads = 4; ctx_params.n_threads_batch = 4; + // the ubatches of 64 tokens alternate between the lanes of the split overlap + ctx_params.split_overlap = split_overlap; if (!encode) { ctx_params.n_ubatch = 64; } @@ -959,11 +961,12 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con bool host_experts; // keep the experts in host memory, see host_experts_test size_t moe_cache_size; std::vector attn_split; // empty: follow the tensor split + bool split_overlap; device_config(std::vector devs, std::string name, llama_split_mode split_mode, bool host_experts = false, size_t moe_cache_size = 0, - std::vector attn_split = {}) + std::vector attn_split = {}, bool split_overlap = false) : devs(std::move(devs)), label(std::move(name)), split_mode(split_mode), host_experts(host_experts), moe_cache_size(moe_cache_size), - attn_split(std::move(attn_split)) {} + attn_split(std::move(attn_split)), split_overlap(split_overlap) {} }; const llama_model_tensor_buft_override host_experts_overrides[] = { @@ -1005,6 +1008,9 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con attn_uneven[0] = 1.0f; dev_configs.emplace_back(devices_meta, "Meta -as 1/3", LLAMA_SPLIT_MODE_TENSOR, false, 0, attn_uneven); max_device_label_length = std::max(max_device_label_length, dev_configs.back().label.length()); + + // consecutive ubatches on two schedulers whose AllReduce overlaps + dev_configs.emplace_back(devices_meta, "Meta -so", LLAMA_SPLIT_MODE_TENSOR, false, 0, std::vector{}, true); } } @@ -1111,7 +1117,7 @@ static int test_backends(const std::string & arch_filter, const size_t seed, con } if (dc.split_mode != LLAMA_SPLIT_MODE_TENSOR || llm_arch_supports_sm_tensor(arch)) { test_executed = true; - model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, dc.devs, dc.split_mode, encode, overrides, dc.moe_cache_size, dc.attn_split); + model_and_ctx_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, stdev, dc.devs, dc.split_mode, encode, overrides, dc.moe_cache_size, dc.attn_split, dc.split_overlap); logits_dev = get_logits(model_and_ctx_dev.first.get(), model_and_ctx_dev.second.get(), tokens, encode); const double nmse_val = nmse(logits_cpu, logits_dev); snprintf(nmse_str, sizeof(nmse_str), "(%.2e)", nmse_val); diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 9e6069e174a1..330671f16c79 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -359,6 +359,7 @@ struct cmd_params { std::vector no_kv_offload; std::vector flash_attn; std::vector flash_attn_native_quants; + std::vector split_overlap; std::vector> devices; std::vector> tensor_split; std::vector> attn_split; @@ -407,6 +408,7 @@ static const cmd_params cmd_params_defaults = { /* no_kv_offload */ { false }, /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, /* flash_attn_native_quants */ { false }, + /* split_overlap */ { false }, /* devices */ { {} }, /* tensor_split */ { std::vector(llama_max_devices(), 0.0f) }, /* attn_split */ { std::vector(llama_max_devices(), 0.0f) }, @@ -481,6 +483,7 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str()); printf(" -fa, --flash-attn (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str()); printf(" -fanq, --flash-attn-native-quants <0|1> (default: %s)\n", join(cmd_params_defaults.flash_attn_native_quants, ",").c_str()); + printf(" -so, --split-overlap <0|1> (default: %s)\n", join(cmd_params_defaults.split_overlap, ",").c_str()); printf(" -dev, --device (default: auto)\n"); printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); printf(" -lzm, --lazy-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.lazy_mode, lazy_mode_str), ",").c_str()); @@ -852,6 +855,13 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { } auto p = string_split(argv[i], split_delim); params.flash_attn_native_quants.insert(params.flash_attn_native_quants.end(), p.begin(), p.end()); + } else if (arg == "-so" || arg == "--split-overlap") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = string_split(argv[i], split_delim); + params.split_overlap.insert(params.split_overlap.end(), p.begin(), p.end()); } else if (arg == "-nkvo" || arg == "--no-kv-offload") { if (++i >= argc) { invalid_param = true; @@ -1182,6 +1192,9 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { if (params.flash_attn_native_quants.empty()) { params.flash_attn_native_quants = cmd_params_defaults.flash_attn_native_quants; } + if (params.split_overlap.empty()) { + params.split_overlap = cmd_params_defaults.split_overlap; + } if (params.devices.empty()) { params.devices = cmd_params_defaults.devices; } @@ -1250,6 +1263,7 @@ struct cmd_params_instance { bool no_kv_offload; llama_flash_attn_type flash_attn; bool flash_attn_native_quants; + bool split_overlap; std::vector devices; std::vector tensor_split; std::vector attn_split; @@ -1337,6 +1351,7 @@ struct cmd_params_instance { cparams.offload_kqv = !no_kv_offload; cparams.flash_attn_type = flash_attn; cparams.flash_attn_native_quants = flash_attn_native_quants; + cparams.split_overlap = split_overlap; cparams.embeddings = embeddings; cparams.op_offload = !no_op_offload; cparams.swa_full = false; @@ -1374,6 +1389,7 @@ static std::vector get_cmd_params_instances(const cmd_param for (const auto & nkvo : params.no_kv_offload) for (const auto & fa : params.flash_attn) for (const auto & fanq : params.flash_attn_native_quants) + for (const auto & so : params.split_overlap) for (const auto & nt : params.n_threads) for (const auto & cm : params.cpu_mask) for (const auto & cs : params.cpu_strict) @@ -1405,6 +1421,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, /* .flash_attn_native_quants = */ fanq, + /* .split_overlap = */ so, /* .devices = */ devs, /* .tensor_split = */ ts, /* .attn_split = */ as, @@ -1445,6 +1462,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, /* .flash_attn_native_quants = */ fanq, + /* .split_overlap = */ so, /* .devices = */ devs, /* .tensor_split = */ ts, /* .attn_split = */ as, @@ -1485,6 +1503,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, /* .flash_attn_native_quants = */ fanq, + /* .split_overlap = */ so, /* .devices = */ devs, /* .tensor_split = */ ts, /* .attn_split = */ as, @@ -1530,6 +1549,7 @@ struct test { bool no_kv_offload; llama_flash_attn_type flash_attn; bool flash_attn_native_quants; + bool split_overlap; std::vector devices; std::vector tensor_split; std::vector attn_split; @@ -1573,6 +1593,7 @@ struct test { no_kv_offload = inst.no_kv_offload; flash_attn = inst.flash_attn; flash_attn_native_quants = inst.flash_attn_native_quants; + split_overlap = inst.split_overlap; devices = inst.devices; tensor_split = inst.tensor_split; attn_split = inst.attn_split; @@ -1638,7 +1659,7 @@ struct test { "model_filename", "model_type", "model_size", "model_n_params", "n_batch", "n_ubatch", "n_threads", "cpu_mask", "cpu_strict", "poll", "type_k", "type_v", "n_gpu_layers", "n_cpu_moe", "split_mode", - "main_gpu", "no_kv_offload", "flash_attn", "flash_attn_native_quants", "devices", "tensor_split", + "main_gpu", "no_kv_offload", "flash_attn", "flash_attn_native_quants", "split_overlap", "devices", "tensor_split", "attn_split", "tensor_buft_overrides", "load_mode", "lazy_mode", "embeddings", "no_op_offload", "no_host", "repack", "fit_target", "fit_min_ctx", @@ -1658,7 +1679,7 @@ struct test { field == "fit_target" || field == "fit_min_ctx" || field == "flash_attn") { return INT; } - if (field == "f16_kv" || field == "no_kv_offload" || field == "flash_attn_native_quants" || field == "cpu_strict" || + if (field == "f16_kv" || field == "no_kv_offload" || field == "flash_attn_native_quants" || field == "split_overlap" || field == "cpu_strict" || field == "embeddings" || field == "no_host" || field == "repack") { return BOOL; } @@ -1738,6 +1759,7 @@ struct test { std::to_string(no_kv_offload), std::to_string((int) flash_attn), std::to_string(flash_attn_native_quants), + std::to_string(split_overlap), devices_to_string(devices), tensor_split_str, attn_split_str, @@ -1972,6 +1994,9 @@ struct markdown_printer : public printer { if (field == "flash_attn_native_quants") { return "fanq"; } + if (field == "split_overlap") { + return "so"; + } if (field == "load_mode") { return "lm"; } @@ -2062,6 +2087,9 @@ struct markdown_printer : public printer { if (params.flash_attn_native_quants.size() > 1 || params.flash_attn_native_quants != cmd_params_defaults.flash_attn_native_quants) { fields.emplace_back("flash_attn_native_quants"); } + if (params.split_overlap.size() > 1 || params.split_overlap != cmd_params_defaults.split_overlap) { + fields.emplace_back("split_overlap"); + } if (params.devices.size() > 1 || params.devices != cmd_params_defaults.devices) { fields.emplace_back("devices"); }