diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 545b892c0eed..dc983d94ab3a 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -8,10 +8,10 @@ extern "C" { #define RPC_PROTO_MAJOR_VERSION 5 #define RPC_PROTO_MINOR_VERSION 1 -#define RPC_PROTO_PATCH_VERSION 1 +#define RPC_PROTO_PATCH_VERSION 2 #ifdef __cplusplus -static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); #endif #define GGML_RPC_MAX_SERVERS 16 diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 12d01193c5f1..093d727195f6 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -513,6 +513,7 @@ extern "C" { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID, GGML_OP_OUT_PROD, + GGML_OP_MUL_MAT_ID_BACK, // MUL_MAT_ID's gradient w.r.t. its activations (src1), through frozen experts GGML_OP_SCALE, GGML_OP_SET, @@ -1466,6 +1467,25 @@ extern "C" { struct ggml_tensor * ids, struct ggml_tensor * expert_ptrs); + // The gradient of ggml_mul_mat_id(as, b, ids) with respect to b, through the experts held + // fixed (frozen, possibly quantized): the MoE half of training a LoRA on attention, whose + // gradient reaches earlier layers through every later expert FFN. + // + // as: [cols, rows, n_expert] the forward's experts (F32, F16 or quantized) + // grad: [rows, n_expert_used, n_tok] F32, the gradient of the forward's result + // ids: [n_expert_used, n_tok] I32, the forward's routing + // result: [cols, n_expert_used, n_tok] F32, one row per routed slot + // + // result[:, i, t] = sum_r as[:, r, ids[i, t]] * grad[r, i, t] + // + // A forward whose b broadcast over the slots (b->ne[1] < n_expert_used) sums the slots back + // with ggml_repeat_back, so every slot row has exactly one writer. + GGML_API struct ggml_tensor * ggml_mul_mat_id_back( + struct ggml_context * ctx, + struct ggml_tensor * as, + struct ggml_tensor * grad, + struct ggml_tensor * ids); + // A: m columns, n rows, // B: p columns, n rows, // result is m columns, p rows diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index b14248897227..a125c1eb2d03 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -1845,6 +1845,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_out_prod(params, tensor); } break; + case GGML_OP_MUL_MAT_ID_BACK: + { + ggml_compute_forward_mul_mat_id_back(params, tensor); + } break; case GGML_OP_SCALE: { ggml_compute_forward_scale(params, tensor); @@ -2335,6 +2339,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_MUL_MAT: case GGML_OP_MUL_MAT_ID: case GGML_OP_OUT_PROD: + case GGML_OP_MUL_MAT_ID_BACK: { n_tasks = n_threads; } break; @@ -2889,6 +2894,11 @@ struct ggml_cplan ggml_graph_plan( cur = ggml_type_size(GGML_TYPE_F32) * ggml_out_prod_cpu_wsize(node->src[0]->ne[0]) * n_tasks; } } break; + case GGML_OP_MUL_MAT_ID_BACK: + { + const struct ggml_tensor * ids = node->src[2]; + cur = ggml_mul_mat_id_back_cpu_wsize(node->src[0]->ne[0], node->src[0]->ne[2], ids->ne[0]*ids->ne[1], n_tasks); + } break; case GGML_OP_SET_ROWS: { if (node->src[0]->type == GGML_TYPE_F16 && node->type != GGML_TYPE_F16) { diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 96601da9a074..ab85b0bf2682 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -473,6 +473,9 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st return src0->type == GGML_TYPE_F32 && (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); case GGML_OP_GET_ROWS_BACK: return src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16; + case GGML_OP_MUL_MAT_ID_BACK: + return (src0->type == GGML_TYPE_F32 || ggml_get_type_traits(src0->type)->to_float != NULL) && + src1->type == GGML_TYPE_F32 && op->src[2]->type == GGML_TYPE_I32 && op->type == GGML_TYPE_F32; case GGML_OP_OUT_PROD: return (src0->type == GGML_TYPE_F32 || ((src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && src1->ne[2] % src0->ne[2] == 0 && src1->ne[3] % src0->ne[3] == 0)) && diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index be9d4135cc8f..ee8eb5e35f06 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -4469,6 +4469,108 @@ static void ggml_compute_forward_out_prod_converted_f32( } } +// ggml_compute_forward_mul_mat_id_back +// +// dst[:, i, t] = sum_r as[:, r, ids[i, t]] * grad[r, i, t]: the activations' gradient of +// mul_mat_id through its experts, held fixed. Task 0 groups the routed (slot, token) pairs by +// expert with a counting sort (count, prefix sum, fill: one entry per pair, Fable on fork #82); +// then each task takes whole experts, dequantizes one in tiles of +// GGML_OUT_PROD_KT rows ONCE, and adds each tile row, scaled by its gradient, into every slot +// routed to that expert. A slot has exactly one expert, so each dst row has one writer. +void ggml_compute_forward_mul_mat_id_back( + const ggml_compute_params * params, + ggml_tensor * dst) { + const ggml_tensor * as = dst->src[0]; + const ggml_tensor * grad = dst->src[1]; + const ggml_tensor * ids = dst->src[2]; + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t cols = as->ne[0]; + const int64_t rows = as->ne[1]; + const int64_t n_expert = as->ne[2]; + const int64_t n_used = ids->ne[0]; + const int64_t n_tok = ids->ne[1]; + const int64_t n_slots = n_used*n_tok; + + GGML_ASSERT(dst->type == GGML_TYPE_F32 && dst->nb[0] == sizeof(float)); + GGML_ASSERT(grad->nb[0] == sizeof(float)); + GGML_ASSERT(as->nb[0] == ggml_type_size(as->type)); // expert rows are contiguous + + ggml_to_float_t const to_float = ggml_get_type_traits(as->type)->to_float; + GGML_ASSERT(as->type == GGML_TYPE_F32 || to_float != NULL); + + char * wdata = (char *) params->wdata; + int64_t * offset = (int64_t *) wdata; // [n_expert + 1]: expert e's pairs + int32_t * routed = (int32_t *) (offset + n_expert + 1); // [n_slots] of (i, t), by expert + float * tiles = (float *) GGML_PAD((uintptr_t) (routed + n_slots*2), CACHE_LINE_SIZE); + + const auto expert_of = [&](int64_t i, int64_t t) { + const int32_t e = *(const int32_t *) ((const char *) ids->data + i*ids->nb[0] + t*ids->nb[1]); + GGML_ASSERT(e >= 0 && e < n_expert); + return e; + }; + + if (ith == 0) { + memset(dst->data, 0, ggml_nbytes(dst)); + memset(offset, 0, (n_expert + 1)*sizeof(int64_t)); + for (int64_t t = 0; t < n_tok; ++t) { + for (int64_t i = 0; i < n_used; ++i) { + offset[expert_of(i, t) + 1]++; + } + } + for (int64_t e = 0; e < n_expert; ++e) { + offset[e + 1] += offset[e]; + } + // fill: offset[e] advances to the end of e's range, which is where e + 1 begins + for (int64_t t = 0; t < n_tok; ++t) { + for (int64_t i = 0; i < n_used; ++i) { + int32_t * pair = routed + 2*offset[expert_of(i, t)]++; + pair[0] = (int32_t) i; + pair[1] = (int32_t) t; + } + } + // shift back: offset[e] is again the start of e's range + for (int64_t e = n_expert; e > 0; --e) { + offset[e] = offset[e - 1]; + } + offset[0] = 0; + } + ggml_barrier(params->threadpool); + + const int64_t KT = GGML_OUT_PROD_KT; + float * wt = tiles + ith*KT*cols; // [KT][cols] + + for (int64_t e = ith; e < n_expert; e += nth) { + const int64_t n_routed = offset[e + 1] - offset[e]; + if (n_routed == 0) { + continue; + } + const int32_t * pairs = routed + 2*offset[e]; + for (int64_t k0 = 0; k0 < rows; k0 += KT) { + const int64_t kt = std::min(KT, rows - k0); + for (int64_t r = 0; r < kt; ++r) { + const char * row = (const char *) as->data + (k0 + r)*as->nb[1] + e*as->nb[2]; + if (as->type == GGML_TYPE_F32) { + memcpy(wt + r*cols, row, cols*sizeof(float)); + } else { + to_float(row, wt + r*cols, cols); + } + } + for (int64_t s = 0; s < n_routed; ++s) { + const int64_t i = pairs[2*s + 0]; + const int64_t t = pairs[2*s + 1]; + float * d = (float *) ((char *) dst->data + i*dst->nb[1] + t*dst->nb[2]); + const char * g = (const char *) grad->data + i*grad->nb[1] + t*grad->nb[2]; + for (int64_t r = 0; r < kt; ++r) { + ggml_vec_mad_f32((int) cols, d, wt + r*cols, *(const float *) (g + (k0 + r)*grad->nb[0])); + } + } + } + } +} + void ggml_compute_forward_out_prod( const ggml_compute_params * params, ggml_tensor * dst) { @@ -5204,18 +5306,22 @@ static void ggml_compute_forward_get_rows_back_f32( memset(dst->data, 0, ggml_nbytes(dst)); const int nc = src0->ne[0]; - const int nr = ggml_nelements(src1); GGML_ASSERT( dst->ne[0] == nc); GGML_ASSERT(src0->nb[0] == sizeof(float)); - for (int i = 0; i < nr; ++i) { - const int r = ((int32_t *) src1->data)[i]; + // batch for batch (ggml_get_rows_back): index k of batch (i2, i3) adds its gradient row + // into that batch's row r of dst + for (int64_t i3 = 0; i3 < src1->ne[2]; ++i3) { + for (int64_t i2 = 0; i2 < src1->ne[1]; ++i2) { + for (int64_t k = 0; k < src1->ne[0]; ++k) { + const int32_t r = *(const int32_t *) ((const char *) src1->data + k*src1->nb[0] + i2*src1->nb[1] + i3*src1->nb[2]); + GGML_ASSERT(r >= 0 && r < dst->ne[1]); - ggml_vec_add_f32(nc, - (float *) ((char *) dst->data + r*dst->nb[1]), - (float *) ((char *) dst->data + r*dst->nb[1]), - (float *) ((char *) src0->data + i*src0->nb[1])); + float * d = (float *) ((char *) dst->data + r*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3]); + ggml_vec_add_f32(nc, d, d, (const float *) ((const char *) src0->data + k*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3])); + } + } } } diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 9c1e15dc2765..f3b6c067bd81 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -49,6 +49,7 @@ void ggml_compute_forward_rms_norm_back(const struct ggml_compute_params * param void ggml_compute_forward_group_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_l2_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_out_prod(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_mul_mat_id_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_scale(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_set(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_cpy(const struct ggml_compute_params * params, struct ggml_tensor * dst); @@ -96,6 +97,14 @@ void ggml_compute_forward_flash_attn_back(const struct ggml_compute_params * par static inline int64_t ggml_out_prod_cpu_wsize(int64_t ne0) { return GGML_OUT_PROD_KT*ne0 + GGML_OUT_PROD_BN*GGML_OUT_PROD_KT; } + +// ggml_compute_forward_mul_mat_id_back's work buffer, in bytes: the routing grouped by expert +// with a counting sort (an offset per expert, plus one, then every (slot, token) pair once), +// and each task's dequantized tile of GGML_OUT_PROD_KT expert rows. +static inline size_t ggml_mul_mat_id_back_cpu_wsize(int64_t cols, int64_t n_expert, int64_t n_slots, int n_tasks) { + return (size_t) (n_expert + 1)*sizeof(int64_t) + (size_t) n_slots*2*sizeof(int32_t) + + (size_t) n_tasks*GGML_OUT_PROD_KT*cols*sizeof(float) + 2*CACHE_LINE_SIZE; +} // FLASH_ATTN_BACK's blocks: BR query rows (pass 1) and BC KV positions (pass 2) share one conversion of // each K and V row. Its per-thread scratch, in floats; the work-size planner reserves this per task. #define GGML_FA_BACK_BR 8 diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 8faa9a66da77..592d66d3728e 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1490,7 +1490,11 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te op->src[0]->nb[0] == sizeof(float) && op->src[4]->nb[0] == sizeof(float) && (op->src[3] == NULL || op->src[3]->type == GGML_TYPE_F16); case GGML_OP_GET_ROWS_BACK: + // kernel_get_rows_back walks one flat index list over dst's rows: the 2-D form only; + // a batched one ([1, n_expert, n_tokens], a MoE's router weights) runs where it is + // implemented return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_I32 && + ggml_is_vector(op->src[1]) && op->ne[2] == 1 && op->ne[3] == 1 && ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op); case GGML_OP_CROSS_ENTROPY_LOSS: return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c1d86aaac5c1..d68ce14492a2 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -18275,7 +18275,9 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm } } case GGML_OP_GET_ROWS_BACK: - return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32; + // the shader walks one flat index list: the 2-D form only (a batched one, a MoE's + // router weights, runs where it is implemented), as CUDA's supports_op says + return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->ne[2] == 1 && op->ne[3] == 1; case GGML_OP_SET_ROWS: { if ((op->src[0]->type != GGML_TYPE_F32 && op->src[0]->type != GGML_TYPE_F16) || diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 94bbb9d40fbc..8869f9cedb54 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1022,6 +1022,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "MUL_MAT", "MUL_MAT_ID", "OUT_PROD", + "MUL_MAT_ID_BACK", "SCALE", "SET", @@ -1101,7 +1102,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", }; -static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); +static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1138,6 +1139,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "X*Y", "X[i]*Y", "X*Y", + "X[i]^T*dY", "x*v", "y-\\>view(x)", @@ -1217,7 +1219,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", }; -static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); +static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -3374,6 +3376,31 @@ struct ggml_tensor * ggml_mul_mat_id_gather( return result; } +// ggml_mul_mat_id_back + +struct ggml_tensor * ggml_mul_mat_id_back( + struct ggml_context * ctx, + struct ggml_tensor * as, + struct ggml_tensor * grad, + struct ggml_tensor * ids) { + GGML_ASSERT(!ggml_is_transposed(as)); + GGML_ASSERT(as->ne[3] == 1); + GGML_ASSERT(grad->type == GGML_TYPE_F32 && grad->ne[3] == 1); + GGML_ASSERT(ids->type == GGML_TYPE_I32 && ids->ne[2] == 1 && ids->ne[3] == 1); + GGML_ASSERT(grad->ne[0] == as->ne[1]); // one gradient value per expert row + GGML_ASSERT(grad->ne[1] == ids->ne[0]); // per routed slot + GGML_ASSERT(grad->ne[2] == ids->ne[1]); // per token + + struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, as->ne[0], ids->ne[0], ids->ne[1]); + + result->op = GGML_OP_MUL_MAT_ID_BACK; + result->src[0] = as; + result->src[1] = grad; + result->src[2] = ids; + + return result; +} + // ggml_out_prod static inline bool ggml_can_out_prod(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { @@ -3940,12 +3967,17 @@ struct ggml_tensor * ggml_get_rows_back( struct ggml_tensor * a, struct ggml_tensor * b, struct ggml_tensor * c) { - GGML_ASSERT(ggml_is_matrix(a) && ggml_is_vector(b) && b->type == GGML_TYPE_I32); - GGML_ASSERT(ggml_is_matrix(c) && (a->ne[0] == c->ne[0])); + // The inverse of ggml_get_rows, batch for batch: `a` is the gradient of its result + // [ne0, n_idx, ne2, ne3], `b` the indices [n_idx, ne2, ne3], `c` the gathered source + // [ne0, n_rows, ne2, ne3]; each index's row of `a` adds into its row of the result. A MoE + // gathers its router weights per token this way ([1, n_expert, n_tokens] by [n_used, n_tokens]). + GGML_ASSERT(b->type == GGML_TYPE_I32 && b->ne[3] == 1); + GGML_ASSERT(a->ne[0] == c->ne[0] && a->ne[1] >= b->ne[0]); // index k reads gradient row k + GGML_ASSERT(a->ne[2] == b->ne[1] && a->ne[3] == b->ne[2]); + GGML_ASSERT(c->ne[2] == b->ne[1] && c->ne[3] == b->ne[2]); // TODO: implement non F32 return - //struct ggml_tensor * result = ggml_new_tensor_2d(ctx, a->type, a->ne[0], b->ne[0]); - struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, c->ne[0], c->ne[1]); + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, c->ne[0], c->ne[1], c->ne[2], c->ne[3]); result->op = GGML_OP_GET_ROWS_BACK; result->src[0] = a; @@ -7049,6 +7081,7 @@ bool ggml_node_has_backward(const struct ggml_tensor * tensor) { case GGML_OP_REPEAT_BACK: case GGML_OP_RMS_NORM: case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: case GGML_OP_SCALE: case GGML_OP_SET: case GGML_OP_CPY: @@ -7302,6 +7335,22 @@ static void ggml_compute_backward( grad))); // [m,p,qq,rr] } } break; + case GGML_OP_MUL_MAT_ID: { + // dst[:, i, t] = as[:, :, ids[i, t]]^T b[:, i % n_b, t]. The experts are held fixed + // (a frozen, usually quantized base): only the activations' gradient is formed, so a + // LoRA on attention gets its whole gradient through every later expert FFN. + GGML_ASSERT(!src0_needs_grads && "a gradient for MUL_MAT_ID's experts is not implemented: they are trained frozen"); + // no graph builds the pointer-table gather (ggml_mul_mat_id_gather) for training; + // its backward is designed when one does, never computed from the plain layout + GGML_ASSERT(tensor->src[3] == NULL && "MUL_MAT_ID's backward through the pointer-table gather is not implemented"); + if (src1_needs_grads) { + struct ggml_tensor * back = ggml_mul_mat_id_back(ctx, src0, grad, tensor->src[2]); + if (src1->ne[1] != back->ne[1]) { + back = ggml_repeat_back(ctx, back, src1); // slot i read row i % n_b: sum them back + } + ggml_add_or_set(ctx, cgraph, isrc1, back); + } + } break; case GGML_OP_SCALE: { if (src0_needs_grads) { float s; @@ -7888,6 +7937,13 @@ void ggml_build_backward_expand( ignore_src[1] = true; break; + // the routing (I32 ids) and the gather's expert pointer table (I64) are not + // differentiable; the experts and the activations are + case GGML_OP_MUL_MAT_ID: + ignore_src[2] = true; + ignore_src[3] = true; + break; + default: break; } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index a90000500b1f..097f5f1a6e42 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -2340,6 +2340,67 @@ struct test_get_rows : public test_case { } }; +// GET_ROWS's gradient on batched rows (a MoE gathers its router weights per token: [1, n_expert, +// n_tokens] by [n_used, n_tokens]), checked by finite differences. Every row of every batch is +// gathered (r >= m) and the output is weighted by a fixed random tensor: the suite's metric, +// (a-b)/(a+b) averaged, is NaN on any element whose gradient is 0 on both sides, and a NaN mean +// never exceeds the threshold, so a row nothing gathers would pass the case whatever the kernel did. +struct test_get_rows_grad : public test_case { + const int n; // cols + const int m; // rows + const int r; // rows to get, >= m + const int be1; // batches + + std::string vars() override { + return VARS_TO_STR4(n, m, r, be1); + } + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "GET_ROWS"; + } + + test_get_rows_grad(int n = 1, int m = 6, int r = 9, int be1 = 7) : n(n), m(m), r(r), be1(be1) { + GGML_ASSERT(r >= m); + } + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * in = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n, m, be1); + ggml_set_name(in, "in"); + ggml_set_param(in); + + ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, r, be1); + ggml_set_name(rows, "rows"); + + ggml_tensor * w = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n, r, be1); + ggml_set_name(w, "w"); + + ggml_tensor * out = ggml_mul(ctx, ggml_get_rows(ctx, in, rows), w); + ggml_set_name(out, "out"); + + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + std::random_device rd; + std::default_random_engine rng(rd()); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (t->type == GGML_TYPE_I32) { + for (int b = 0; b < be1; b++) { + std::vector data(r); + for (int k = 0; k < r; k++) { + data[k] = k % m; // every row at least once + } + std::shuffle(data.begin(), data.end(), rng); + ggml_backend_tensor_set(t, data.data(), b*t->nb[1], r*sizeof(int32_t)); + } + } else { + init_tensor_uniform(t); + } + } + } +}; + // GGML_OP_GET_ROWS_BACK struct test_get_rows_back : public test_case { const ggml_type type; @@ -4991,6 +5052,106 @@ struct test_mul_mat_id : public test_case { } }; +// The activations' gradient of mul_mat_id through fixed experts (GGML_OP_MUL_MAT_ID_BACK), +// checked by finite differences. The output is weighted by a fixed random tensor before the +// suite's SUM loss: under a bare SUM every upstream gradient value is 1, so a kernel that reads +// the wrong gradient row passes. F32 experts only: a quantized expert's forward rounds b to its +// vec_dot type, so the finite differences are not exact (test_mul_mat_id_back compares those). +struct test_mul_mat_id_grad : public test_case { + const int n_mats; + const int n_used; + const bool b; // broadcast b matrix + const int64_t m; + const int64_t n; + const int64_t k; + + std::string vars() override { + return VARS_TO_STR6(n_mats, n_used, b, m, n, k); + } + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "MUL_MAT_ID"; + } + + test_mul_mat_id_grad(int n_mats = 4, int n_used = 2, bool b = false, int64_t m = 40, int64_t n = 3, int64_t k = 6) + : n_mats(n_mats), n_used(n_used), b(b), m(m), n(n), k(k) { + GGML_ASSERT(n_used <= n_mats); + } + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * as = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, k, m, n_mats); + ggml_set_name(as, "as"); + + ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_mats, n); + ggml_set_name(ids, "ids"); + if (n_used != n_mats) { + ids = ggml_view_2d(ctx, ids, n_used, n, ids->nb[1], 0); + ggml_set_name(ids, "view_of_ids"); + } + + ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, k, this->b ? 1 : n_used, n); + ggml_set_name(b, "b"); + ggml_set_param(b); + + ggml_tensor * w = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, m, n_used, n); + ggml_set_name(w, "w"); + + ggml_tensor * out = ggml_mul(ctx, ggml_mul_mat_id(ctx, as, b, ids), w); + ggml_set_name(out, "out"); + + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + init_mul_mat_id_tensors(ctx, n_mats); + } +}; + +// GGML_OP_MUL_MAT_ID_BACK: the activations' gradient of mul_mat_id through fixed (frozen, +// quantized) experts, compared across backends +struct test_mul_mat_id_back : public test_case { + const ggml_type type_a; + const int n_mats; + const int n_used; + const int64_t m; // expert rows + const int64_t n; // tokens + const int64_t k; // expert cols + + std::string vars() override { + return VARS_TO_STR6(type_a, n_mats, n_used, m, n, k); + } + + test_mul_mat_id_back(ggml_type type_a = GGML_TYPE_F32, int n_mats = 8, int n_used = 2, int64_t m = 32, int64_t n = 16, int64_t k = 64) + : type_a(type_a), n_mats(n_mats), n_used(n_used), m(m), n(n), k(k) { + GGML_ASSERT(n_used <= n_mats); + } + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * as = ggml_new_tensor_3d(ctx, type_a, k, m, n_mats); + ggml_set_name(as, "as"); + + ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_mats, n); + ggml_set_name(ids, "ids"); + if (n_used != n_mats) { + ids = ggml_view_2d(ctx, ids, n_used, n, ids->nb[1], 0); + ggml_set_name(ids, "view_of_ids"); + } + + ggml_tensor * grad = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, m, n_used, n); + ggml_set_name(grad, "grad"); + + ggml_tensor * out = ggml_mul_mat_id_back(ctx, as, grad, ids); + ggml_set_name(out, "out"); + + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + init_mul_mat_id_tensors(ctx, n_mats); + } +}; + // GGML_OP_MUL_MAT_ID + GGML_OP_ADD or GGML_OP_MUL struct test_mul_mat_id_fusion : public test_case { const ggml_type type_a; @@ -8950,6 +9111,23 @@ static std::vector> make_test_cases_eval() { } test_cases.emplace_back(new test_get_rows_back(GGML_TYPE_F32, 1, 8, 2, 1, false)); + // batched (a MoE's router weights: [1, n_expert, n_tokens] gathered by [n_used, n_tokens]) + test_cases.emplace_back(new test_get_rows_back(GGML_TYPE_F32, 1, 16, 4, 7, false)); + test_cases.emplace_back(new test_get_rows_back(GGML_TYPE_F32, 8, 16, 4, 3, true)); + test_cases.emplace_back(new test_get_rows_grad(1, 6, 9, 7)); + test_cases.emplace_back(new test_get_rows_grad(3, 8, 8, 5)); + + // MoE training: the activations' gradient through fixed experts. A small F32 mul_mat_id + // grad-checks it by finite differences (per-slot and broadcast b); MUL_MAT_ID_BACK itself is + // compared across backends for the expert types a resident base is served in + for (bool bcast : {false, true}) { + test_cases.emplace_back(new test_mul_mat_id_grad(4, 2, bcast, 40, 3, 6)); // m > GGML_OUT_PROD_KT: two row tiles + test_cases.emplace_back(new test_mul_mat_id_grad(8, 3, bcast, 70, 5, 9)); + } + for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K}) { + test_cases.emplace_back(new test_mul_mat_id_back(type_a, 8, 2, 32, 16, 256)); + test_cases.emplace_back(new test_mul_mat_id_back(type_a, 16, 4, 64, 33, 512)); + } test_cases.emplace_back(new test_get_rows_back(GGML_TYPE_F32, 1, 70000, 4, 1, false)); // row count > CUDA grid-y limit (65535) for (ggml_type type : all_types) { for (bool v : {false, true}) {