diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index ded678e68..fe58ea3bb 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -592,7 +592,18 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( GGML_ASSERT(split_states_equal(src_ss[0], src_ss[1])); return {assume_sync ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_PARTIAL, {0}, {1}, 1}; } - GGML_ABORT("fatal error"); + if (src_ss[0].axis == src_ss[1].axis && src_ss[0].axis >= GGML_BACKEND_SPLIT_AXIS_2 && + src_ss[0].axis < GGML_MAX_DIMS) { + GGML_ASSERT(split_states_equal(src_ss[0], src_ss[1])); + return src_ss[0]; + } + // batched matmul with the batches split across devices and a replicated activation + if (src_ss[0].axis >= GGML_BACKEND_SPLIT_AXIS_2 && src_ss[0].axis < GGML_MAX_DIMS && + src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + return src_ss[0]; + } + GGML_ABORT("unsupported mul_mat split states: node=%s src0=%s axis=%d src1=%s axis=%d", + tensor->name, tensor->src[0]->name, (int) src_ss[0].axis, tensor->src[1]->name, (int) src_ss[1].axis); //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1}; }; @@ -760,14 +771,33 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( }; auto handle_flash_attn_ext = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { - GGML_ASSERT( src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2); - GGML_ASSERT( src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2); - GGML_ASSERT( src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2); - GGML_ASSERT(tensor->src[4] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + + if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + GGML_ASSERT(src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + GGML_ASSERT(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; + } + + GGML_ASSERT(src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2); + const bool kv_split = src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2 && + src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2; + const bool kv_mirrored = src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED && + src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED; + GGML_ASSERT(kv_split || kv_mirrored); GGML_ASSERT(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_0); return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; }; + auto handle_lightning_indexer = [&]( + const std::vector & src_ss) -> ggml_backend_meta_split_state { + for (size_t i = 0; i < 4; i++) { + GGML_ASSERT(src_ss[i].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } + return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1}; + }; + auto handle_ssm_conv = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { if (src_ss[0].axis == src_ss[1].axis) { if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { @@ -938,7 +968,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( split_state = handle_rope(src_ss); } break; case GGML_OP_ROPE_BACK: { - split_state = handle_generic(src_ss, /*scalar_only =*/ true); + split_state = handle_rope(src_ss); } break; case GGML_OP_CLAMP: { split_state = handle_generic(src_ss, /*scalar_only =*/ false); @@ -1002,6 +1032,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( case GGML_OP_GATED_DELTA_NET: { split_state = handle_gated_delta_net(src_ss); } break; + case GGML_OP_LIGHTNING_INDEXER: { + split_state = handle_lightning_indexer(src_ss); + } break; case GGML_OP_DSV4_HC_COMB: case GGML_OP_DSV4_HC_PRE: case GGML_OP_DSV4_HC_POST: { @@ -1086,13 +1119,14 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( if (buf_ctx->debug > 0) { std::string srcs_info; for (size_t i = 0; i < GGML_MAX_SRC; i++) { - if (tensor->src[i] == nullptr) { + if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { continue; } if (!srcs_info.empty()) { srcs_info += ", "; } - const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor->src[0], true); + const ggml_backend_meta_split_state split_state = + ggml_backend_meta_get_split_state(tensor->src[i], true); GGML_ASSERT(split_state.n_segments == 1); const char * axis_name = ggml_backend_meta_split_axis_name(split_state.axis); std::string ne_info; @@ -1271,6 +1305,108 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor(ggml_backend_buffer return ggml_backend_meta_buffer_init_tensor_impl(buf_ctx->get_simple_tensor_container(tensor), tensor); } +static void ggml_backend_meta_buffer_memset_tensor( + ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { + const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); + const ggml_backend_meta_split_state split_state = + ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); + GGML_ASSERT(ggml_is_contiguous(tensor) || split_state.axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + + if (split_state.n_segments != 1 || split_state.nr[0] != 1) { + GGML_ASSERT(split_state.axis >= 0 && split_state.axis < GGML_MAX_DIMS); + GGML_ASSERT(split_state.nr[0] != 0); + GGML_ASSERT(tensor->ne[3] == 1); + + std::vector simple_offsets(n_bufs, 0); + if (split_state.axis == GGML_BACKEND_SPLIT_AXIS_0) { + GGML_ASSERT(tensor->ne[2] == 1); + + const size_t row_stride = tensor->nb[1]; + GGML_ASSERT(offset % row_stride == 0); + GGML_ASSERT(size % row_stride == 0); + const int64_t row_start = offset / row_stride; + const int64_t row_count = size / row_stride; + GGML_ASSERT(row_start + row_count <= tensor->ne[1]); + + const int64_t blck_size = ggml_blck_size(tensor->type); + for (size_t s = 0; s < split_state.n_segments; s++) { + for (size_t r = 0; r < split_state.nr[s]; r++) { + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + GGML_ASSERT(split_state.ne[s*n_bufs + j] % blck_size == 0); + const size_t nbytes = split_state.ne[s*n_bufs + j]/blck_size * tensor->nb[0]; + for (int64_t row = 0; row < row_count; row++) { + ggml_backend_tensor_memset(simple_tensor, value, + simple_offsets[j] + (row_start + row)*simple_tensor->nb[1], nbytes); + } + simple_offsets[j] += nbytes; + } + } + } + return; + } + + GGML_ASSERT(split_state.axis == GGML_BACKEND_SPLIT_AXIS_1); + + const size_t row_stride = tensor->nb[2]; + GGML_ASSERT(offset % row_stride == 0); + GGML_ASSERT(size % row_stride == 0); + const int64_t row_start = offset / row_stride; + const int64_t row_count = size / row_stride; + GGML_ASSERT(row_start + row_count <= tensor->ne[2]); + + for (size_t s = 0; s < split_state.n_segments; s++) { + for (size_t r = 0; r < split_state.nr[s]; r++) { + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t nbytes = split_state.ne[s*n_bufs + j] * tensor->nb[1]; + for (int64_t row = 0; row < row_count; row++) { + ggml_backend_tensor_memset(simple_tensor, value, + simple_offsets[j] + (row_start + row)*simple_tensor->nb[2], nbytes); + } + simple_offsets[j] += nbytes; + } + } + } + return; + } + + switch (split_state.axis) { + case GGML_BACKEND_SPLIT_AXIS_0: + case GGML_BACKEND_SPLIT_AXIS_1: + case GGML_BACKEND_SPLIT_AXIS_2: { + const size_t chunk_size_full = tensor->nb[split_state.axis + 1]; + GGML_ASSERT(offset % chunk_size_full == 0); + GGML_ASSERT(size % chunk_size_full == 0); + const int64_t i_start = offset / chunk_size_full; + const int64_t i_stop = (offset + size) / chunk_size_full; + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + const size_t chunk_size = simple_tensor->nb[split_state.axis + 1]; + if (chunk_size == 0) { + continue; + } + for (int64_t i = i_start; i < i_stop; i++) { + ggml_backend_tensor_memset(simple_tensor, value, i*chunk_size, chunk_size); + } + } + } break; + case GGML_BACKEND_SPLIT_AXIS_PARTIAL: { + GGML_ASSERT(value == 0); + [[fallthrough]]; + } + case GGML_BACKEND_SPLIT_AXIS_MIRRORED: { + for (size_t j = 0; j < n_bufs; j++) { + ggml_tensor * simple_tensor = ggml_backend_meta_buffer_simple_tensor(tensor, j); + ggml_backend_tensor_memset(simple_tensor, value, offset, size); + } + } break; + default: { + GGML_ABORT("fatal error"); + } + } +} + static void ggml_backend_meta_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { const size_t n_bufs = ggml_backend_meta_buffer_n_bufs(buffer); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(tensor, /*assume_sync =*/ false); @@ -1518,7 +1654,7 @@ static const ggml_backend_buffer_i ggml_backend_meta_buffer_iface = { /* .free_buffer = */ ggml_backend_meta_buffer_free_buffer, /* .get_base = */ ggml_backend_meta_buffer_get_base, /* .init_tensor = */ ggml_backend_meta_buffer_init_tensor, - /* .memset_tensor = */ nullptr, // TODO implement + /* .memset_tensor = */ ggml_backend_meta_buffer_memset_tensor, /* .set_tensor = */ ggml_backend_meta_buffer_set_tensor, /* .get_tensor = */ ggml_backend_meta_buffer_get_tensor, /* .set_tensor_2d = */ nullptr, @@ -1871,7 +2007,7 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, { // For MoE models it may make sense to delay the AllReduce in order to reduce I/O: - auto get_i_delayed = [&](const int i) -> int { + auto get_i_delayed_branch = [&](const int i) -> int { int id = i; // i_delayed int idr = i; // i_delayed return, last safe return value @@ -1971,6 +2107,62 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend, return idr; }; + // AllReduce(a) + AllReduce(b) == AllReduce(a + b) for independent partial branches. + auto get_i_delayed = [&](const int i) -> int { + const int i_delayed = get_i_delayed_branch(i); + ggml_tensor * node = cgraph->nodes[i_delayed]; + + if (ggml_node_get_use_count(cgraph, i_delayed) != 1) { + return i_delayed; + } + + for (int id = i_delayed + 1; id < cgraph->n_nodes; id++) { + ggml_tensor * next = cgraph->nodes[id]; + if (next->view_src == node) { + return i_delayed; + } + for (int s = 0; s < GGML_MAX_SRC; s++) { + if (next->src[s] == node) { + return i_delayed; + } + } + + if (next->view_src != nullptr && next->view_src->op == GGML_OP_NONE && ggml_backend_buffer_is_host(next->view_src->buffer)) { + continue; + } + if (ggml_backend_meta_get_split_state(next, false).axis != GGML_BACKEND_SPLIT_AXIS_PARTIAL) { + continue; + } + + const int i_other = id; + const int i_other_delayed = get_i_delayed_branch(i_other); + ggml_tensor * other = cgraph->nodes[i_other_delayed]; + if (ggml_node_get_use_count(cgraph, i_other_delayed) != 1 || i_other_delayed + 1 >= cgraph->n_nodes) { + return i_delayed; + } + + ggml_tensor * sum = cgraph->nodes[i_other_delayed + 1]; + if (sum->op != GGML_OP_ADD || + !ggml_are_same_shape(node, other) || node->type != other->type || sum->type != node->type || + !((sum->src[0] == node && sum->src[1] == other) || + (sum->src[0] == other && sum->src[1] == node)) || + ggml_backend_meta_get_split_state(sum, false).axis != GGML_BACKEND_SPLIT_AXIS_MIRRORED) { + return i_delayed; + } + + for (size_t j = 0; j < n_backends; j++) { + auto & bcj = backend_ctx->backend_configs[j]; + const bool compute = bcj.nodes[i]->flags & GGML_TENSOR_FLAG_COMPUTE; + const bool compute_other = bcj.nodes[i_other]->flags & GGML_TENSOR_FLAG_COMPUTE; + if (compute != compute_other) { + return i_delayed; + } + } + return i_other_delayed + 1; + } + return i_delayed; + }; + int i_start = 0; for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i];