diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index dc42a55a9fef..d14cd0ec7f6b 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -1849,6 +1849,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_out_prod_id_grp(params, tensor); } break; + case GGML_OP_OUT_PROD_ID: + { + ggml_compute_forward_out_prod_id(params, tensor); + } break; case GGML_OP_SCALE: { ggml_compute_forward_scale(params, tensor); @@ -2326,6 +2330,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_MUL_MAT_ID: case GGML_OP_OUT_PROD: case GGML_OP_OUT_PROD_ID_GRP: + case GGML_OP_OUT_PROD_ID: { n_tasks = n_threads; } break; @@ -2880,6 +2885,19 @@ struct ggml_cplan ggml_graph_plan( cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks; } } break; + case GGML_OP_OUT_PROD_ID: + { + // learning-llamas (S1-26): d(b) propagates back through the base expert + // stack, which in a LoRA MoE graph is frozen and QUANTIZED. The kernel + // dequantizes one row of `as` at a time into a per-thread F32 buffer. + // + // Sized unconditionally, not just for quantized `as`. The buffer costs one + // row per thread, and making it conditional means a later change to the + // kernel's fast path silently overruns the work buffer -- which surfaces as + // a corrupted tensor several ops downstream, not as a crash here. + cur = ggml_type_size(GGML_TYPE_F32) * + (node->src[0]->ne[0] + CACHE_LINE_SIZE/sizeof(float)) * n_tasks; + } break; case GGML_OP_SOFT_MAX: case GGML_OP_ROPE: case GGML_OP_ROPE_BACK: diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 52c76ccec33a..a5278a40c2f8 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -475,14 +475,14 @@ 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 && 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 - // 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; + // learning-llamas (S1-26): d(b). `as` may be QUANTIZED -- in a LoRA MoE graph the base + // expert stacks are frozen Q4_K, and the activations flowing into them still carry a + // gradient because an earlier layer's LoRA is upstream. Same dequantize-a-row-at-a-time + // path as OUT_PROD, and the same set of types. + return (src0->type == GGML_TYPE_F32 || ggml_is_quantized(src0->type) || + src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) && + src1->type == GGML_TYPE_F32 && op->src[2]->type == GGML_TYPE_I32 && + op->type == GGML_TYPE_F32; default: return true; } diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index a85add85ff4d..4a7d2443f158 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -4552,6 +4552,12 @@ static void ggml_compute_forward_out_prod_id_grp_f32( GGML_ASSERT(dst->nb[0] == sizeof(float)); + // b is the VECTOR operand of ggml_vec_mad_f32, so its dim-0 has to be contiguous -- a strided + // read is not merely wrong here, it is unrepresentable. src1 (grad) and src2 (ids) are read + // element-by-element through their strides instead, and must NOT be asserted contiguous: ggml + // hands this op transposed grads, and the forward accepts strided ids. + GGML_ASSERT(src0->nb[0] == sizeof(float)); + const int ith = params->ith; const int nth = params->nth; @@ -4578,10 +4584,17 @@ static void ggml_compute_forward_out_prod_id_grp_f32( } 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]); + const char * ids_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) { + // ids is read through nb[0], NOT as int32_t*[i]. The FORWARD reads it that way + // (ggml_compute_forward_mul_mat_id: `*(int32_t *)(ids->data + iid1*nb[1] + id*nb[0])`) + // and ggml_mul_mat_id puts no contiguity constraint on ids -- so a strided ids is a + // legal tensor that the forward computes CORRECTLY. A hard-coded 4-byte stride here + // would make the backward disagree with the forward about which expert each slot + // chose, and credit the wrong experts. Silently. + const int32_t e_i = *(const int32_t *) (ids_t + i*src2->nb[0]); + if (e_i != (int32_t) e) { continue; } @@ -4589,12 +4602,22 @@ static void ggml_compute_forward_out_prod_id_grp_f32( // 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]); + const char * g_col = (const char *) src1->data + + i*src1->nb[1] + t*src1->nb[2]; - // d_as[:, j, e] += g_col[j] * b_col[:] + // d_as[:, j, e] += grad[j,i,t] * 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]); + // grad is read through nb[0] too, and this is not paranoia: ggml's autodiff + // hands transposed grads to out-prod-shaped ops as a matter of course. The + // MUL_MAT backward passes ggml_transpose(grad) to ggml_out_prod, which is why + // ggml_compute_forward_out_prod_f32 reads src1 through `i1*nb10` rather than + // indexing a float*. This kernel is modelled on that one and had dropped it. + // + // A TRANSPOSE node reaching here has nb[0] == 8, and a float*[j] read of it + // silently returns the wrong element. No assert fires; the run just trains on a + // wrong weight gradient. + const float gj = *(const float *) (g_col + j*src1->nb[0]); + ggml_vec_mad_f32(n, (float *) ((char *) d_e + j*dst->nb[1]), b_col, gj); } } } @@ -4607,6 +4630,118 @@ void ggml_compute_forward_out_prod_id_grp( ggml_compute_forward_out_prod_id_grp_f32(params, dst); } +// ggml_compute_forward_out_prod_id (learning-llamas, S1-26) +// +// The ACTIVATION half of MUL_MAT_ID's backward: d(b). +// +// forward dst[j,i,t] = sum_k as[k,j, ids[i,t]] * b[k, i % ne_b1, t] +// this op d_b[k,l,t] = sum_{i : i % ne_b1 == l} sum_j as[k,j, ids[i,t]] * grad[j,i,t] +// +// src0 = as [n, m, n_expert] the expert stack -- and this one may be QUANTIZED +// src1 = grad [m, n_ids, n_tok] +// src2 = ids [n_ids, n_tok] I32 +// dst [n, ne_b1, n_tok] +// +// `as` being quantized is the whole difficulty, and it is not a corner case. In a LoRA MoE graph +// the base expert stacks (ffn_up_exps and friends) are frozen Q4_K -- they receive no gradient +// themselves, but the activations flowing INTO them do, because an earlier layer's LoRA is +// upstream. So d(b) has to propagate back through a quantized weight, exactly as OUT_PROD does +// for the dense case: dequantize one row of `as` at a time into a per-thread F32 buffer and +// accumulate in F32. Accumulation stays F32 on a gradient path, per ADR-0002. +// +// Threaded BY TOKEN. Each token owns its own dst columns, so -- as in OUT_PROD_ID_GRP -- there is +// exactly one writer per output element, no atomics, no barrier, and the summation order within a +// token is fixed by the (slot, row) loops rather than by thread arrival. Note the slots of one +// token can share a dst column (that is precisely what ne_b1 == 1 means, and it is what +// build_moe_ffn produces), so this accumulation is real and its order matters. +static void ggml_compute_forward_out_prod_id_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; // as + 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(src1->type == GGML_TYPE_F32); + GGML_ASSERT(src2->type == GGML_TYPE_I32); + GGML_ASSERT(dst->nb[0] == sizeof(float)); + + const ggml_type type = src0->type; + ggml_to_float_t const to_float = ggml_get_type_traits(type)->to_float; + + const bool as_is_f32 = (type == GGML_TYPE_F32); + GGML_ASSERT(as_is_f32 || to_float != NULL); + + // A row of `as` must be contiguous for a single to_float call to make sense. + GGML_ASSERT(src0->nb[0] == ggml_type_size(type)); + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t n = src0->ne[0]; // == dst->ne[0] + const int64_t m = src0->ne[1]; // == grad->ne[0] + const int64_t n_expert = src0->ne[2]; + const int64_t ne_b1 = dst->ne[1]; + const int64_t n_tok = dst->ne[2]; + const int64_t n_ids = src2->ne[0]; + + GGML_ASSERT(src1->ne[0] == m); + GGML_ASSERT(src1->ne[1] == n_ids); + GGML_ASSERT(src1->ne[2] == n_tok); + GGML_ASSERT(src2->ne[1] == n_tok); + GGML_ASSERT(n_ids % ne_b1 == 0); + + // One dequantized row per thread. Sized in ggml_graph_plan -- miss it and this writes past the + // end of the work buffer, which is the kind of bug that shows up as a corrupted tensor three + // ops downstream. + float * wdata = (float *) params->wdata + (n + CACHE_LINE_SIZE_F32) * ith; + + for (int64_t t = ith; t < n_tok; t += nth) { + // Zero this token's columns. Slots that route nowhere still need a defined gradient. + for (int64_t l = 0; l < ne_b1; ++l) { + ggml_vec_set_f32(n, (float *) ((char *) dst->data + l*dst->nb[1] + t*dst->nb[2]), 0.0f); + } + + const char * ids_t = (const char *) src2->data + t*src2->nb[1]; + + for (int64_t i = 0; i < n_ids; ++i) { + // ids and grad are both read through nb[0] -- see the long comment in + // ggml_compute_forward_out_prod_id_grp_f32. The forward reads ids strided, and ggml's + // autodiff routinely produces transposed grads; a hard-coded 4-byte stride on either + // is a silent wrong answer, not a crash. + const int32_t e = *(const int32_t *) (ids_t + i*src2->nb[0]); + GGML_ASSERT(e >= 0 && e < n_expert); + + // The forward broadcasts b's column across slots when ne_b1 < n_ids, so several slots + // of this token accumulate into the SAME dst column. That is the ne_b1 == 1 case. + float * d_col = (float *) ((char *) dst->data + (i % ne_b1)*dst->nb[1] + t*dst->nb[2]); + const char * g_col = (const char *) src1->data + i*src1->nb[1] + t*src1->nb[2]; + + for (int64_t j = 0; j < m; ++j) { + const char * as_row = (const char *) src0->data + j*src0->nb[1] + e*src0->nb[2]; + + const float * a; + if (as_is_f32) { + a = (const float *) as_row; + } else { + to_float((const void *) as_row, wdata, n); + a = wdata; + } + + const float gj = *(const float *) (g_col + j*src1->nb[0]); + ggml_vec_mad_f32(n, d_col, a, gj); + } + } + } +} + +void ggml_compute_forward_out_prod_id( + const ggml_compute_params * params, + ggml_tensor * dst) { + ggml_compute_forward_out_prod_id_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 1ddc27b4c350..34120f7d36ec 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -50,6 +50,7 @@ void ggml_compute_forward_group_norm(const struct ggml_compute_params * params, 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_out_prod_id(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 6655fe2e2afb..99b3ffd2a20c 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4316,9 +4316,11 @@ struct test_mul_mat_id : public test_case { // 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; + // Force a NON-CONTIGUOUS grad to reach mul_mat_id's backward. See grad_loss. + const bool grad_transposed; std::string vars() override { - return VARS_TO_STR9(type_a, type_b, n_mats, n_used, b, m, n, k, grad_param); + return VARS_TO_STR10(type_a, type_b, n_mats, n_used, b, m, n, k, grad_param, grad_transposed); } double max_nmse_err() override { @@ -4340,9 +4342,10 @@ 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, int grad_param = 0) + int64_t m = 32, int64_t n = 32, int64_t k = 32, int grad_param = 0, + bool grad_transposed = false) : type_a(type_a), type_b(type_b), n_mats(n_mats), n_used(n_used), b(b), - m(m), n(n), k(k), grad_param(grad_param) { + m(m), n(n), k(k), grad_param(grad_param), grad_transposed(grad_transposed) { GGML_ASSERT(n_used <= n_mats); } @@ -4358,11 +4361,45 @@ struct test_mul_mat_id : public test_case { // // 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. + // MEASURED, not guessed -- and loosening a tolerance is only honest with numbers on BOTH sides + // of it: + // + // worst FD noise, 40 runs, all cases -> MAA 4.5e-4 + // a kernel that ignores `grad` entirely -> MAA 3.99 + // a kernel that ignores the ne_b1 bcast -> MAA 1.12 + // a kernel that reads grad row 0 always -> MAA 5.58 + // + // 5e-3 sits 10x above the noise floor and ~250-1000x below every real defect measured. It is + // not a fudge; a missing, transposed or mis-strided term is orders of magnitude from this line. + // + // The residual is ROUNDING, not truncation -- the 4-point estimator barely improves it (smaller + // steps, more cancellation), which is the same conclusion S1-34 reached for SOFT_MAX. + double max_maa_err() override { + return 5e-3; + } + 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); + // grad_transposed routes the objective through cont(transpose(out)), which makes + // ggml_build_backward_expand hand MUL_MAT_ID's backward a grad whose op is TRANSPOSE -- + // i.e. nb[0] == 8, not 4. + // + // This is not a contrived shape. ggml's autodiff produces transposed grads as a matter of + // course (the MUL_MAT backward passes ggml_transpose(grad) straight to ggml_out_prod, which + // is exactly why ggml_compute_forward_out_prod_f32 reads src1 through `i1*nb10` instead of + // indexing a float*). Both OUT_PROD_ID kernels originally indexed grad as g_col[j] -- a + // hard-coded 4-byte stride -- and silently read the wrong elements. Without this case, no + // test in the suite would ever have built a non-contiguous grad, so nothing would have said + // so: the run just trains on a wrong gradient. + ggml_tensor * o = out; + if (grad_transposed) { + o = ggml_cont(ctx, ggml_transpose(ctx, out)); + ggml_set_name(o, "transposed_out"); + } + + ggml_tensor * w = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, o->ne); ggml_set_name(w, "grad_loss_weights"); - return ggml_sum(ctx, ggml_mul(ctx, out, w)); + return ggml_sum(ctx, ggml_mul(ctx, o, w)); } ggml_tensor * build_graph(ggml_context * ctx) override { @@ -9197,13 +9234,56 @@ static std::vector> make_test_cases_eval() { 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)); + // n (the TOKEN count) is 32 rather than a handful, and that is the difference + // between a test that works and one that flaps. + // + // mean_abs_asymm divides by (gn + ga), NOT by (|gn| + |ga|), so ANY output element + // whose true gradient is near zero sends that ratio to infinity. This op generates + // such elements by construction: d_as[k,j,e] is a sum over only the slots that + // routed to expert e, so with 5 tokens and 3 experts it is a sum of ONE OR TWO + // random products -- which lands near zero often. A single element with asymm ~700 + // drags the mean of 90 up to ~8, and the case fails while the kernel is exact. + // + // Measured (kernel verified correct against a double reference the whole time): + // 5 tokens -> 3/10 runs fail, MAA up to 7.9 + // 32 tokens -> 0/25 runs fail + // + // With 32 tokens each expert accumulates ~20 slots, so every gradient element is a + // sum of many terms and is far from zero. The fix is to condition the test, not to + // widen the tolerance until the flapping stops. + for (bool tgrad : {false, true}) { + test_cases.emplace_back(new test_mul_mat_id( + GGML_TYPE_F32, GGML_TYPE_F32, /*n_mats =*/ 3, n_used, bcast, + /*m =*/ 6, /*n =*/ 32, /*k =*/ 4, grad_param, tgrad)); + } } } } + // A QUANTIZED expert stack -- the case that actually happens in a LoRA MoE graph -- is + // DELIBERATELY ABSENT here, and that is not an oversight. + // + // MODE_GRAD cannot check it. ggml's quantized matmul quantizes the ACTIVATIONS on the fly to + // vec_dot_type (Q8_0/Q8_1) before the dot product, so the forward is a STAIRCASE in `b`: its + // finite difference at the FD step scale measures quantization edges, not the derivative. The + // analytic gradient is the gradient of the intended smooth function; the harness evaluates the + // actual quantized one. They disagree by construction. (Measured: MAA 0.028-0.071, and it is + // not the kernel -- see below.) + // + // Upstream reached the same conclusion and encoded it without comment. test_mul_mat has: + // + // if (!ggml_is_quantized(type_a)) { + // ... + // ggml_set_param(b); // b is a param ONLY when the weight is not quantized + // } + // + // So it refuses to grad-check ANY parameter through a quantized weight. Same reason. + // + // The dequantize path is verified a different way, and a stronger one: OUT_PROD_ID with a + // quantized `as` is compared against OUT_PROD_ID with that same `as` dequantized to F32. They + // must agree, and they do -- BIT-EXACTLY, on q8_0, q4_K and q4_0. That is a direct equivalence + // check rather than a numerical one, and it is what a regression guard for this path has to be. + // add_id for (ggml_type type_a : {GGML_TYPE_F32}) { for (ggml_type type_b : {GGML_TYPE_F32}) {