From 0d0afcb919e4b27e30d2e9f595e3f31ccab94f2e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CDillon?= <“dillontblake@gmail.com”> Date: Tue, 14 Jul 2026 09:43:00 +1000 Subject: [PATCH 1/2] S1-26: OUT_PROD_ID CPU kernel -- the activation half of MUL_MAT_ID's backward d_b[k,l,t] = sum_{i : i % ne_b1 == l} sum_j as[k,j, ids[i,t]] * grad[j,i,t] `as` may be QUANTIZED, and that is the whole difficulty rather than a corner case. In a LoRA MoE graph the base expert stacks are frozen Q4_K. They take 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. Same dequantize-a-row-at-a-time path as OUT_PROD, accumulating in F32 (ADR-0002). Threaded by token: one writer per output column, no atomics, no barrier. MODE_GRAD CANNOT CHECK THE QUANTIZED PATH, and it took a while to work out why. ggml's quantized matmul quantizes the ACTIVATIONS on the fly to vec_dot_type before the dot product. So the forward is a STAIRCASE in b: its finite difference at the FD step scale measures quantization edges, not a 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 the kernel is exact). Upstream reached the same conclusion and encoded it without a word of comment: // test_mul_mat if (!ggml_is_quantized(type_a)) { ggml_set_param(b); // b is a param ONLY when the weight is not quantized } So the quantized cases are deliberately absent from MODE_GRAD, and the dequantize path is verified a stronger way instead: OUT_PROD_ID with a quantized `as` against OUT_PROD_ID with that same `as` dequantized to F32. BIT-EXACT on q8_0, q4_K and q4_0. TWO TEST FIXES, both found by measuring rather than assuming. 1. The token count is 32, not 5. mean_abs_asymm divides by (gn + ga), NOT by (|gn| + |ga|), so any output element whose true gradient is near zero sends the ratio to infinity. This op manufactures such elements: d_as[k,j,e] sums only over the slots that routed to expert e, so with 5 tokens and 3 experts it is a sum of ONE OR TWO random products. One element with asymm ~700 drags the mean of 90 to ~8. 5 tokens -> 3/10 runs fail, MAA up to 7.9 (kernel exact throughout) 32 tokens -> 0/30 runs fail Condition the test; do not widen the tolerance until the flapping stops. 2. max_maa_err = 5e-3, measured on both sides: worst FD noise, 40 runs 4.5e-4 ignore grad entirely 9.21 ignore the ne_b1 broadcast 3.39 / 1.15 read grad row 0 always 4.41 route every slot everywhere 5.30 10x above the noise floor, 200-1800x below every real defect. Five mutations introduced, five caught. Both kernels verified against a naive DOUBLE reference across 15 configurations covering both broadcast modes: relative error ~1e-7, i.e. float32 epsilon. That reference is what proved the FD failures were the harness and not the arithmetic -- though note it shares the kernels' assumption about the broadcast rule, so it cannot check THAT. The rule was confirmed by reading the forward: row_mapping.i1 is the SLOT index (the in-code comment calling it "selected expert index" is wrong), so the forward's b column is slot % ne_b1. --- ggml/src/ggml-cpu/ggml-cpu.c | 18 ++++++ ggml/src/ggml-cpu/ggml-cpu.cpp | 16 ++--- ggml/src/ggml-cpu/ops.cpp | 108 +++++++++++++++++++++++++++++++++ ggml/src/ggml-cpu/ops.h | 1 + tests/test-backend-ops.cpp | 60 +++++++++++++++++- 5 files changed, 194 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index dc42a55a9fe..d14cd0ec7f6 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 52c76ccec33..a5278a40c2f 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 a85add85ff4..73a83a58395 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -4607,6 +4607,114 @@ 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 int32_t * ids_t = (const int32_t *) ((const char *) src2->data + t*src2->nb[1]); + + for (int64_t i = 0; i < n_ids; ++i) { + const int32_t e = ids_t[i]; + 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 float * g_col = (const float *) ((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; + } + + ggml_vec_mad_f32(n, d_col, a, g_col[j]); + } + } + } +} + +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 1ddc27b4c35..34120f7d36e 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 6655fe2e2af..127f0c66b9d 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4358,6 +4358,23 @@ 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); ggml_set_name(w, "grad_loss_weights"); @@ -9197,13 +9214,54 @@ 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}) { + // 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. 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)); + /*m =*/ 6, /*n =*/ 32, /*k =*/ 4, grad_param)); } } } + // 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}) { From b8e007507a71b17d2504953b081addc08164b6fc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CDillon?= <“dillontblake@gmail.com”> Date: Tue, 14 Jul 2026 10:06:01 +1000 Subject: [PATCH 2/2] fix(S1-26/S1-27): both kernels read grad and ids with a hard-coded 4-byte stride Found by adversarial review, with a working reproduction. Two real bugs, both silent, both in the kernels I had just declared verified. 1. GRAD (src1) was indexed as g_col[j] -- a hard-coded 4-byte dim-0 stride. ggml's autodiff hands out-prod-shaped ops 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 -- the kernel mine is modelled on -- reads src1 through `i1*nb10` rather than indexing a float*. I dropped that. A TRANSPOSE node reaching here has nb[0] == 8. Reproduced: 24 of 36 output elements wrong, max abs err 9.2. No assert fires. The run just trains on a wrong gradient. 2. IDS (src2) was indexed as ids_t[i], likewise hard-coded. The FORWARD reads ids through nb[0] (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 and the backward misread -- crediting the wrong experts with each other's gradients. Reproduced: forward bit-identical across two ids layouts, backward disagreeing by 1.03. Both are now read through their nb[0]. b keeps a contiguity assert, because it is the VECTOR operand of ggml_vec_mad_f32 -- a strided read there is not merely wrong, it is unrepresentable. WHY NOTHING CAUGHT THIS. test-backend-ops only ever builds a CONTIGUOUS grad, so no test in the suite could see it -- and neither could my double-precision reference, which built its own contiguous tensors. Every check I had said the kernels were exact, and they were, on the only inputs anyone was feeding them. So the guard is a new test case, not a comment. 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. Reintroducing the exact bug: contiguous-grad cases (the entire old suite) -> PASS. invisible. transposed-grad cases (the new guard) -> FAIL, MAA 21.0 / 2.57 / 2.55 Re-verified after the fix: double-precision reference across 15 configs (rel ~1e-7), bit-identical across 1/2/4/8 threads, quantized-vs-dequantized bit-exact on q8_0/q4_K/q4_0, strided ids identical to contiguous ids, MUL_MAT_ID grad 0/20 runs failing, the grad allowlist and the MUL_MAT_ID forward unregressed. --- ggml/src/ggml-cpu/ops.cpp | 49 +++++++++++++++++++++++++++++--------- tests/test-backend-ops.cpp | 38 ++++++++++++++++++++++------- 2 files changed, 68 insertions(+), 19 deletions(-) diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 73a83a58395..4a7d2443f15 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); } } } @@ -4680,17 +4703,20 @@ static void ggml_compute_forward_out_prod_id_f32( ggml_vec_set_f32(n, (float *) ((char *) dst->data + l*dst->nb[1] + t*dst->nb[2]), 0.0f); } - 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) { - const int32_t e = ids_t[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 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]; for (int64_t j = 0; j < m; ++j) { const char * as_row = (const char *) src0->data + j*src0->nb[1] + e*src0->nb[2]; @@ -4703,7 +4729,8 @@ static void ggml_compute_forward_out_prod_id_f32( a = wdata; } - ggml_vec_mad_f32(n, d_col, a, g_col[j]); + const float gj = *(const float *) (g_col + j*src1->nb[0]); + ggml_vec_mad_f32(n, d_col, a, gj); } } } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 127f0c66b9d..99b3ffd2a20 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); } @@ -4376,10 +4379,27 @@ struct test_mul_mat_id : public test_case { } 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 { @@ -9231,9 +9251,11 @@ static std::vector> make_test_cases_eval() { // 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. - 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)); + 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)); + } } } }