From 7b700fe1921b57ffe05f80c65b60701ddb0924b9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CDillon?= <“dillontblake@gmail.com”> Date: Tue, 14 Jul 2026 09:06:41 +1000 Subject: [PATCH] S1-27: OUT_PROD_ID_GRP CPU kernel -- the weight half of MUL_MAT_ID's backward d_as[k,j,e] = sum_{(i,t) : ids[i,t] == e} b[k, i % ne_b1, t] * grad[j,i,t] This is the op that makes MoE LoRA possible at all: build_lora_mm_id computes mul_mat_id(B, mul_mat_id(A, cur, ids), ids), so the trainable A/B tensors ARE the 3D expert operand, and without d(as) a LoRA MoE run trains nothing. Threaded BY EXPERT, deliberately rather than incidentally. Each expert's output slice is written by exactly one thread, so the scatter-add needs no atomics and no reduction barrier, and the summation order inside a slice is fixed by the (token, slot) loop rather than by which thread arrives first. Changing the thread count changes WHICH thread owns an expert, never the order of the float sum within it. Verified bit-identical at 1/2/4/8 threads on data chosen to make float addition order-sensitive (magnitudes mixed 1e6 against 1e-6, where a reordered sum would show). Experts nothing routes to still get their slice zeroed -- it is a gradient, and the optimizer reads it whether or not anything routed there this step. THE TEST WOULD HAVE PASSED ON A KERNEL THAT IGNORED THE GRADIENT ENTIRELY. MODE_GRAD's default objective is sum(out), which makes the incoming gradient all-ones, which makes d_as[k,j,e] = sum_{(i,t): ids=e} b[k,i,t] * 1 independent of j -- constant along the whole output axis. So a kernel that never reads `grad` and just sums b-columns per expert satisfies it. Measured, not argued: poisoned kernel (grad ignored) + default sum(out) -> OK, passes poisoned kernel (grad ignored) + weighted grad_loss -> FAIL, MAA 6.56 vs 1e-4 test_mul_mat_id therefore overrides grad_loss with a weighted sum -- the same hook, for the same structural reason, that S1-34 had to add for SOFT_MAX. test_mul_mat_id also takes a grad_param selector now, because the two halves of MUL_MAT_ID's backward are separate ops in separate tickets: `as`-only cases exercise this kernel ALONE, so S1-27 has a green test to stand on while OUT_PROD_ID (S1-26) is still unimplemented and its cases report not-supported. Shapes are deliberately tiny (4x6x8) -- grad_nmax() is 10000 elements and every pre-existing test_mul_mat_id shape is far above it, so they would have been silently skipped. Both forward broadcast modes are covered (ne_b1 == 1 and ne_b1 == n_used); both are live in the real LoRA MoE graph and they are different index arithmetic in the kernel. --- ggml/src/ggml-cpu/ggml-cpu.c | 5 ++ ggml/src/ggml-cpu/ggml-cpu.cpp | 12 +++-- ggml/src/ggml-cpu/ops.cpp | 95 ++++++++++++++++++++++++++++++++++ ggml/src/ggml-cpu/ops.h | 1 + tests/test-backend-ops.cpp | 78 +++++++++++++++++++++++----- 5 files changed, 174 insertions(+), 17 deletions(-) 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}) {