diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 172128685a1..dc42a55a9fe 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_OUT_PROD_ID_GRP: + { + ggml_compute_forward_out_prod_id_grp(params, tensor); + } break; case GGML_OP_SCALE: { ggml_compute_forward_scale(params, tensor); @@ -2321,6 +2325,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_OUT_PROD_ID_GRP: { n_tasks = n_threads; } break; diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 76c2dda64b3..52c76ccec33 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -467,10 +467,16 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st ((ggml_is_quantized(src0->type) || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) && src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; - case GGML_OP_OUT_PROD_ID: case GGML_OP_OUT_PROD_ID_GRP: - // learning-llamas (S1-25): declared, not yet implemented. The kernels land in S1-26 and - // S1-27, which flip this to a real check. + // learning-llamas (S1-27): d(as). All three operands are F32 on the training path -- + // b is activations, grad is a gradient, and the expert stack it feeds is the F32 LoRA + // A/B (base experts are frozen). A quantized `as` would need a dequantizing variant, + // and there is no caller for one. + return src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_I32 && op->type == GGML_TYPE_F32; + case GGML_OP_OUT_PROD_ID: + // learning-llamas (S1-25): declared, not yet implemented. The kernel lands in S1-26, + // which flips this to a real check. // // This case is NOT redundant, and leaving it out is the trap. The default below returns // TRUE -- so a brand-new op with no dispatch case is reported *supported* by the CPU diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index a881cea706f..a85add85ff4 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -4512,6 +4512,101 @@ void ggml_compute_forward_out_prod( } } +// ggml_compute_forward_out_prod_id_grp (learning-llamas, S1-27) +// +// The WEIGHT half of MUL_MAT_ID's backward: d(as). +// +// forward dst[j,i,t] = sum_k as[k,j, ids[i,t]] * b[k, i % ne_b1, t] +// this op d_as[k,j,e] = sum_{(i,t) : ids[i,t] == e} b[k, i % ne_b1, t] * grad[j,i,t] +// +// src0 = b [n, ne_b1, n_tok] activations +// src1 = grad [m, n_ids, n_tok] incoming gradient +// src2 = ids [n_ids, n_tok] I32 which expert each (slot, token) routed to +// dst [n, m, n_expert] +// +// "grp" is for GROUPED: it reduces over every (slot, token) that routed to a given expert. That +// reduction is the whole difficulty. An expert chosen by many tokens accumulates all of them, and +// a token may route to several experts -- so this is a scatter-add, not a permutation. +// +// It is threaded BY EXPERT, and that is deliberate rather than incidental: each expert's output +// slice is written by exactly one thread, so the accumulation needs no atomics and no reduction +// barrier, and the summation order within a slice is fixed by the (t, i) loop rather than by which +// thread got there first. Determinism is a hard requirement on a gradient path (ADR-0002) -- a +// float sum whose order depends on thread scheduling makes a training run unreproducible, and it +// does so silently. +// +// Experts with no tokens routed to them are not skipped: their slice is still zeroed, because it +// is a gradient and the optimizer will read it whether or not anything routed there this step. +static void ggml_compute_forward_out_prod_id_grp_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; // b + const ggml_tensor * src1 = dst->src[1]; // grad + const ggml_tensor * src2 = dst->src[2]; // ids + + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(src2->type == GGML_TYPE_I32); + + GGML_ASSERT(dst->nb[0] == sizeof(float)); + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t n = src0->ne[0]; // b's row length == dst->ne[0] + const int64_t ne_b1 = src0->ne[1]; // b's broadcast dim + const int64_t n_tok = src0->ne[2]; + const int64_t m = src1->ne[0]; // grad's row length == dst->ne[1] + const int64_t n_ids = src1->ne[1]; + const int64_t n_expert = dst->ne[2]; + + GGML_ASSERT(dst->ne[0] == n); + GGML_ASSERT(dst->ne[1] == m); + GGML_ASSERT(src1->ne[2] == n_tok); + GGML_ASSERT(src2->ne[0] == n_ids); + GGML_ASSERT(src2->ne[1] == n_tok); + GGML_ASSERT(n_ids % ne_b1 == 0); + + for (int64_t e = ith; e < n_expert; e += nth) { + float * d_e = (float *) ((char *) dst->data + e*dst->nb[2]); + + // Zero this expert's whole slice first -- including the experts nothing routes to. + for (int64_t j = 0; j < m; ++j) { + ggml_vec_set_f32(n, (float *) ((char *) d_e + j*dst->nb[1]), 0.0f); + } + + for (int64_t t = 0; t < n_tok; ++t) { + const int32_t * ids_t = (const int32_t *) ((const char *) src2->data + t*src2->nb[1]); + + for (int64_t i = 0; i < n_ids; ++i) { + if (ids_t[i] != (int32_t) e) { + continue; + } + + // The forward broadcasts b's columns across slots when ne_b1 < n_ids, so slot i + // reads column i % ne_b1 -- and the gradient must gather from the same column. + const float * b_col = (const float *) ((const char *) src0->data + + (i % ne_b1)*src0->nb[1] + t*src0->nb[2]); + const float * g_col = (const float *) ((const char *) src1->data + + i*src1->nb[1] + t*src1->nb[2]); + + // d_as[:, j, e] += g_col[j] * b_col[:] + for (int64_t j = 0; j < m; ++j) { + ggml_vec_mad_f32(n, (float *) ((char *) d_e + j*dst->nb[1]), b_col, g_col[j]); + } + } + } + } +} + +void ggml_compute_forward_out_prod_id_grp( + const ggml_compute_params * params, + ggml_tensor * dst) { + ggml_compute_forward_out_prod_id_grp_f32(params, dst); +} + // ggml_compute_forward_scale static void ggml_compute_forward_scale_f32( diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 3a5ebae5c00..1ddc27b4c35 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_out_prod_id_grp(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); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index c4d8e657c00..6655fe2e2af 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4313,9 +4313,12 @@ struct test_mul_mat_id : public test_case { const int64_t m; const int64_t n; const int64_t k; + // learning-llamas: bit 1 = ask for d(as) [OUT_PROD_ID_GRP], bit 2 = ask for d(b) [OUT_PROD_ID]. + // 0 keeps the case eval-only, which is what every pre-existing instantiation wants. + const int grad_param; std::string vars() override { - return VARS_TO_STR8(type_a, type_b, n_mats, n_used, b, m, n, k); + return VARS_TO_STR9(type_a, type_b, n_mats, n_used, b, m, n, k, grad_param); } double max_nmse_err() override { @@ -4337,29 +4340,51 @@ struct test_mul_mat_id : public test_case { test_mul_mat_id(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32, int n_mats = 8, int n_used = 2, bool b = false, - int64_t m = 32, int64_t n = 32, int64_t k = 32) + int64_t m = 32, int64_t n = 32, int64_t k = 32, int grad_param = 0) : type_a(type_a), type_b(type_b), n_mats(n_mats), n_used(n_used), b(b), - m(m), n(n), k(k) { + m(m), n(n), k(k), grad_param(grad_param) { GGML_ASSERT(n_used <= n_mats); } + // MODE_GRAD's default objective is sum(out), and under it THIS OP'S WEIGHT GRADIENT IS + // VACUOUS. With an all-ones incoming gradient, + // + // d_as[i,j,e] = sum_{(s,t): ids=e} b[i,s'] * 1 + // + // is independent of j -- the result is constant along the entire output axis, so a kernel that + // ignored `grad` completely and just summed b-columns per expert would pass every case. A + // weighted sum makes the objective depend on grad's actual (j, slot, token) structure, which is + // the only thing that can catch a transposed or mis-strided read of it. + // + // Same trap, same fix as SOFT_MAX (S1-34): an op can be invisible to sum(out) for structural + // reasons, and then its grad test checks nothing while reporting OK. + ggml_tensor * grad_loss(ggml_context * ctx, ggml_tensor * out) override { + ggml_tensor * w = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, out->ne); + ggml_set_name(w, "grad_loss_weights"); + + return ggml_sum(ctx, ggml_mul(ctx, out, w)); + } + ggml_tensor * build_graph(ggml_context * ctx) override { // C^T = A * B^T: (k, m) * (k, n) => (m, n) ggml_tensor * as = ggml_new_tensor_3d(ctx, type_a, k, m, n_mats); ggml_set_name(as, "as"); - // learning-llamas (S1-25): ask for BOTH gradients. + // learning-llamas (S1-25/S1-27): WHICH gradients are asked for is a test parameter, and it + // has to be, because the two halves of MUL_MAT_ID's backward are separate ops landing in + // separate tickets: // - // `as` is not an exotic param. build_lora_mm_id computes - // mul_mat_id(B, mul_mat_id(A, cur, ids), ids), so the trainable LoRA A/B tensors ARE the - // 3D expert operand -- LoRA-only MoE training needs the weight-grad half, and an - // "activations only" backward would silently train nothing at all. + // grad_param & 1 -> as needs grads -> emits OUT_PROD_ID_GRP (S1-27, implemented) + // grad_param & 2 -> b needs grads -> emits OUT_PROD_ID (S1-26, still pending) // - // These cases build a backward graph containing OUT_PROD_ID / OUT_PROD_ID_GRP, which no - // backend supports yet (the kernels are S1-26 / S1-27). They therefore register and report - // not-supported rather than executing -- which is the point: the wiring is exercised now, - // and the day a kernel lands these turn on with no test change. - if (type_a == GGML_TYPE_F32) { + // An `as`-only case therefore exercises S1-27's kernel ALONE. Asking for both would make + // every case depend on the op that does not exist yet, and S1-27 would have no green test + // to stand on. + // + // `as` is not an exotic param: build_lora_mm_id computes + // mul_mat_id(B, mul_mat_id(A, cur, ids), ids), so the trainable LoRA A/B tensors ARE the + // 3D expert operand. LoRA-only MoE training needs the weight-grad half. + if ((grad_param & 1) && type_a == GGML_TYPE_F32) { ggml_set_param(as); } @@ -4372,7 +4397,7 @@ struct test_mul_mat_id : public test_case { ggml_tensor * b = ggml_new_tensor_3d(ctx, type_b, k, this->b ? 1 : n_used, n); ggml_set_name(b, "b"); - if (type_b == GGML_TYPE_F32) { + if ((grad_param & 2) && type_b == GGML_TYPE_F32) { ggml_set_param(b); } @@ -9154,6 +9179,31 @@ static std::vector> make_test_cases_eval() { 256, 16, 16, {1, 1}, {nr2, 1})); } + // learning-llamas (S1-25/S1-27): MUL_MAT_ID gradient cases. + // + // Deliberately TINY. grad_nmax() is 10000 parameter elements, and every pre-existing + // test_mul_mat_id shape is far above it -- so they would be SILENTLY SKIPPED under MODE_GRAD + // and print OK. (That is trap #2 of ADR-0002, and it is why these shapes are 4x6x8 rather than + // anything realistic: 4*6*3 = 72 elements for `as`.) + // + // grad_param 1 = d(as) only -> exercises OUT_PROD_ID_GRP alone (S1-27, implemented). + // grad_param 2 = d(b) only -> exercises OUT_PROD_ID alone (S1-26, still pending, so these + // register and report not-supported rather than aborting). + // grad_param 3 = both. + // + // `b` toggles the forward's BROADCAST: false gives ne_b1 == n_used, true gives ne_b1 == 1. + // Both are live in the real LoRA MoE graph -- the inner mul_mat_id broadcasts, the outer does + // not -- and they are different index arithmetic in the kernel, so both are covered. + for (int grad_param : {1, 2, 3}) { + for (bool bcast : {false, true}) { + for (int n_used : {1, 2}) { + test_cases.emplace_back(new test_mul_mat_id( + GGML_TYPE_F32, GGML_TYPE_F32, /*n_mats =*/ 3, n_used, bcast, + /*m =*/ 6, /*n =*/ 5, /*k =*/ 4, grad_param)); + } + } + } + // add_id for (ggml_type type_a : {GGML_TYPE_F32}) { for (ggml_type type_b : {GGML_TYPE_F32}) {