diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 24c6f7e90c1..78b4412e265 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -592,6 +592,11 @@ extern "C" { GGML_OP_CROSS_ENTROPY_LOSS_SPARSE, GGML_OP_CROSS_ENTROPY_LOSS_SPARSE_BACK, + // learning-llamas: the two halves of MUL_MAT_ID's backward (S1-25). Also at the TAIL, + // for the same reason as above. + GGML_OP_OUT_PROD_ID, + GGML_OP_OUT_PROD_ID_GRP, + GGML_OP_COUNT, }; @@ -1454,6 +1459,40 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * b); + // learning-llamas (S1-25): the two halves of MUL_MAT_ID's backward. + // + // For the forward dst[j,i,t] = sum_k as[k,j, ids[i,t]] * b[k, i % ne_b1, t]: + // + // ggml_out_prod_id -> d(b) [n, ne_b1, n_tokens] : gather the expert matrix each + // (i,t) slot used and push the gradient back through it. Slots + // sharing a b column (the ne_b1 == 1 case, which is what + // build_moe_ffn produces) ACCUMULATE into it. + // + // ggml_out_prod_id_grp -> d(as) [n, m, n_expert] : the outer product b (x) grad, + // scattered into the expert slice each (i,t) selected, and + // accumulated -- an expert chosen by many tokens sums them all. + // + // "grp" is for grouped: it reduces over every (i,t) that routed to a given expert. The + // weight-grad half is NOT optional for LoRA-only MoE training, which is the discovery this + // ticket turns on: build_lora_mm_id makes the trainable A/B tensors the 3D expert operand + // of mul_mat_id itself, so they are `as`, not activations. + // ne_b1 is b's middle dim. It cannot be recovered from as/grad/ids -- the forward broadcasts + // b's columns across slots whenever ids->ne[0] is a multiple of it -- so it is passed, not + // assumed to be 1. + GGML_API struct ggml_tensor * ggml_out_prod_id( + struct ggml_context * ctx, + struct ggml_tensor * as, + struct ggml_tensor * grad, + struct ggml_tensor * ids, + int64_t ne_b1); + + GGML_API struct ggml_tensor * ggml_out_prod_id_grp( + struct ggml_context * ctx, + struct ggml_tensor * b, + struct ggml_tensor * grad, + struct ggml_tensor * ids, + int64_t n_expert); + // // operations on tensors without backpropagation // diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 435bd9c9b66..76c2dda64b3 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -467,6 +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. + // + // 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 + // backend, gets scheduled, and then hits ggml_compute_forward's `default: GGML_ABORT`. + // The op would look implemented right up until it killed the process. + return false; default: return true; } diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index f875a3a90d0..e3777a80580 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1097,9 +1097,12 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "CROSS_ENTROPY_LOSS_SPARSE", "CROSS_ENTROPY_LOSS_SPARSE_BACK", + + "OUT_PROD_ID", + "OUT_PROD_ID_GRP", }; -static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99"); +static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1211,9 +1214,12 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "cross_entropy_loss_sparse(x,y,w)", "cross_entropy_loss_sparse_back(x,y,w)", + + "out_prod_id(as,grad,ids)", + "out_prod_id_grp(b,grad,ids)", }; -static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99"); +static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -3366,6 +3372,93 @@ struct ggml_tensor * ggml_out_prod( return result; } +// ggml_out_prod_id / ggml_out_prod_id_grp (learning-llamas, S1-25) +// +// The two halves of MUL_MAT_ID's backward. Given the forward +// +// mul_mat_id(as, b, ids): dst[j,i,t] = sum_k as[k,j, ids[i,t]] * b[k, i % ne_b1, t] +// +// as [n, m, n_expert] the 3D expert stack +// b [n, ne_b1, n_tok] activations; ne_b1 is 1 in every build_moe_ffn graph +// ids [n_ids, n_tok] I32 which experts token t routed to +// dst [m, n_ids, n_tok] +// +// the gradients are +// +// d(b) [k, l, t] = sum_{i : i % ne_b1 == l} sum_j grad[j,i,t] * as[k,j, ids[i,t]] +// d(as)[k, j, e] = sum_{(i,t) : ids[i,t] == e} grad[j,i,t] * b[k, i % ne_b1, t] +// +// Both are gather/scatter-with-accumulation, which is why neither is expressible as an existing +// op: OUT_PROD has no notion of an index tensor, and the expert axis is not a broadcast axis -- +// it is a *selection*. Two experts can be chosen by the same token (different i), and one expert +// by many tokens; both cases accumulate, and getting that wrong is silent. +// +// Kernels land in S1-26 (OUT_PROD_ID) and S1-27 (OUT_PROD_ID_GRP). Until then no backend +// advertises support, so a graph containing them builds but will not schedule. + +struct ggml_tensor * ggml_out_prod_id( + struct ggml_context * ctx, + struct ggml_tensor * as, + struct ggml_tensor * grad, + struct ggml_tensor * ids, + int64_t ne_b1) { + GGML_ASSERT(!ggml_is_transposed(as)); + GGML_ASSERT(ids->type == GGML_TYPE_I32); + GGML_ASSERT(grad->type == GGML_TYPE_F32); + + GGML_ASSERT(as->ne[3] == 1); // as is 3d (one matrix per expert) + GGML_ASSERT(ids->ne[2] == 1 && ids->ne[3] == 1); // ids is 2d + GGML_ASSERT(grad->ne[0] == as->ne[1]); // grad's row dim is the expert output dim + GGML_ASSERT(grad->ne[1] == ids->ne[0]); // one grad column per (slot, token) + GGML_ASSERT(grad->ne[2] == ids->ne[1]); // ...and one ids row per token + + // d(b) has b's shape, and b's middle dim is NOT recoverable from `as`, `grad` or `ids`: the + // forward broadcasts b's columns across slots whenever ids->ne[0] is a multiple of it. Assume + // 1 (which is all build_moe_ffn ever produces) and a broadcasting graph would silently get a + // wrong-shaped gradient, so take it as an argument and assert the forward's rule instead. + GGML_ASSERT(ne_b1 > 0); + GGML_ASSERT(ids->ne[0] % ne_b1 == 0); + + const int64_t ne[4] = { as->ne[0], ne_b1, ids->ne[1], 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_OUT_PROD_ID; + result->src[0] = as; + result->src[1] = grad; + result->src[2] = ids; + + return result; +} + +struct ggml_tensor * ggml_out_prod_id_grp( + struct ggml_context * ctx, + struct ggml_tensor * b, + struct ggml_tensor * grad, + struct ggml_tensor * ids, + int64_t n_expert) { + GGML_ASSERT(ids->type == GGML_TYPE_I32); + GGML_ASSERT(grad->type == GGML_TYPE_F32); + + GGML_ASSERT(b->ne[3] == 1); + GGML_ASSERT(ids->ne[2] == 1 && ids->ne[3] == 1); + GGML_ASSERT(ids->ne[1] == b->ne[2]); // one expert list per token + GGML_ASSERT(ids->ne[0] % b->ne[1] == 0); // the forward's broadcast rule + GGML_ASSERT(grad->ne[1] == ids->ne[0]); + GGML_ASSERT(grad->ne[2] == ids->ne[1]); + GGML_ASSERT(n_expert > 0); + + // d(as) has as's shape: [n, m, n_expert]. + const int64_t ne[4] = { b->ne[0], grad->ne[0], n_expert, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_OUT_PROD_ID_GRP; + result->src[0] = b; + result->src[1] = grad; + result->src[2] = ids; + + return result; +} + // ggml_scale static struct ggml_tensor * ggml_scale_impl( @@ -6738,6 +6831,40 @@ static void ggml_compute_backward( grad))); // [m,p,qq,rr] } } break; + case GGML_OP_MUL_MAT_ID: { + // learning-llamas (S1-25). src0 = as [n,m,n_expert], src1 = b [n,ne_b1,n_tok], + // src2 = ids [n_ids,n_tok] I32. Forward: + // + // dst[j,i,t] = sum_k as[k,j, ids[i,t]] * b[k, i % ne_b1, t] + // + // Both gradients are gather/scatter WITH ACCUMULATION, and neither is expressible in + // existing ops -- the expert axis is a selection, not a broadcast. See ggml_out_prod_id. + // + // src2 is I32 and so is skipped by the grad builder outright; there is nothing to + // ignore_src here. + // + // src0 needing grads is NOT an exotic case. build_lora_mm_id computes + // mul_mat_id(B, mul_mat_id(A, cur, ids), ids), which makes the trainable LoRA A and B + // the 3D expert operand -- so LoRA-only MoE training needs the WEIGHT-grad half too, + // and an "activations only" shortcut would silently train nothing. + if (src1_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc1, + ggml_out_prod_id(ctx, src0, grad, src2, src1->ne[1])); + } + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, + ggml_out_prod_id_grp(ctx, src1, grad, src2, src0->ne[2])); + } + } break; + case GGML_OP_ADD_ID: { + // learning-llamas (S1-25). dst is a dup of src0 with a per-expert bias added, so + // src0's VJP is the identity. src1 is the bias TABLE -- a frozen base weight; training + // it needs a scatter-add and is deferred (ROADMAP E8 / B-09). src2 is I32. + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, grad); + } + GGML_ASSERT(!src1_needs_grads && "per-expert bias grads are not implemented (ROADMAP E8)"); + } break; case GGML_OP_SCALE: { if (src0_needs_grads) { float s; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 6acb7100fc9..c4d8e657c00 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3262,7 +3262,17 @@ struct test_add_id : public test_case { ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * a = ggml_new_tensor_3d(ctx, type_a, n_embd, n_experts_used, n_token); + ggml_set_name(a, "a"); + // learning-llamas (S1-25): ADD_ID's src0 VJP is the identity. Ask for it -- without a + // ggml_set_param the MODE_GRAD case builds no backward at all and reports OK regardless. + // src1 is the bias TABLE (a frozen base weight; per-expert bias grads are ROADMAP E8) and + // ids is I32, so neither is a param here. + if (type_a == GGML_TYPE_F32) { + ggml_set_param(a); + } + ggml_tensor * b = ggml_new_tensor_2d(ctx, type_b, n_embd, n_experts); + ggml_set_name(b, "b"); ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_experts, n_token); if (n_experts_used != n_experts) { ids = ggml_view_2d(ctx, ids, n_experts_used, n_token, ids->nb[1], 0); @@ -4338,6 +4348,21 @@ struct test_mul_mat_id : public test_case { 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. + // + // `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. + // + // 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) { + ggml_set_param(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) { @@ -4347,6 +4372,9 @@ 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) { + ggml_set_param(b); + } ggml_tensor * out = ggml_mul_mat_id(ctx, as, b, ids); ggml_set_name(out, "out");