diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index e3071c6cf52d..c8f1ed36fc98 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -600,6 +600,10 @@ extern "C" { // learning-llamas: the VJP for every GLU variant ggml could not differentiate (S1-28). GGML_OP_GLU_BACK, + // learning-llamas: the state-space VJPs (S1-29b wiring, S1-30/S1-31 kernels). + GGML_OP_SSM_CONV_BACK, + GGML_OP_SSM_SCAN_BACK, + GGML_OP_COUNT, }; @@ -2512,6 +2516,24 @@ extern "C" { struct ggml_tensor * sx, struct ggml_tensor * c); + // learning-llamas (S1-29b): d(sx) for ggml_ssm_conv. + // + // The forward is a depthwise causal convolution: + // + // y[i1, t, i3] = sum_{i0 < d_conv} sx[t + i0, i1, i3] * c[i0, i1] + // + // so the gradient w.r.t. the input window is the same convolution run backwards: + // + // d_sx[j, i1, i3] = sum_{t : 0 <= j - t < d_conv, 0 <= t < n_t} dy[i1, t, i3] * c[j - t, i1] + // + // dst has sx's shape. The conv weight `c` is a frozen base weight on this project's LoRA path + // and takes no gradient (ROADMAP E8); the backward case asserts rather than silently skipping. + GGML_API struct ggml_tensor * ggml_ssm_conv_back( + struct ggml_context * ctx, + struct ggml_tensor * dy, // [d_inner, n_t, n_s] + struct ggml_tensor * sx, // for the output shape + struct ggml_tensor * c); + GGML_API struct ggml_tensor * ggml_ssm_scan( struct ggml_context * ctx, struct ggml_tensor * s, @@ -2522,6 +2544,30 @@ extern "C" { struct ggml_tensor * C, struct ggml_tensor * ids); + // learning-llamas (S1-29b): the VJP of the selective scan. + // + // ggml_ssm_scan returns a PACKED 1-D dst -- `y` concatenated with the final states -- and so + // does this: one op cannot return five tensors, so the backward packs + // + // [ d_s | d_x | d_dt | d_B | d_C ] + // + // and the backward case views each region back onto its source. `grad` is the WHOLE packed + // gradient of the forward's dst, state region included: MODE_GRAD's objective sums over it, so + // it is NOT zero there, and a kernel that assumed otherwise would disagree with the finite + // difference and be wrong to. The state-grad region seeds the reverse recurrence at t = n_t. + // + // A (the decay matrix) and ids (I32) take no gradient here. + GGML_API struct ggml_tensor * ggml_ssm_scan_back( + struct ggml_context * ctx, + struct ggml_tensor * grad, // the packed gradient of ssm_scan's dst + struct ggml_tensor * s, + struct ggml_tensor * x, + struct ggml_tensor * dt, + struct ggml_tensor * A, + struct ggml_tensor * B, + struct ggml_tensor * C, + struct ggml_tensor * ids); + // partition into non-overlapping windows with padding if needed // example: // a: 768 64 64 1 diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 2e65b151c5f6..2355b5202006 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -474,6 +474,16 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st // 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_SSM_CONV_BACK: + case GGML_OP_SSM_SCAN_BACK: + // learning-llamas (S1-29b): declared, not yet implemented. The kernels land in S1-30 + // and S1-31, which flip this to a real check. + // + // This case is NOT redundant. The default below returns TRUE, so a 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`. It would look implemented right up + // until it killed the process. Same trap S1-25 hit. + return false; case GGML_OP_GLU_BACK: // learning-llamas (S1-28): F32 throughout. The GLU forwards accept F16, but a gradient // is F32 on this project's training path by policy (ADR-0002), and there is no caller diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index a286c8f4f58a..08c71458b7a3 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1102,9 +1102,12 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "OUT_PROD_ID_GRP", "GLU_BACK", + + "SSM_CONV_BACK", + "SSM_SCAN_BACK", }; -static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); +static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1221,9 +1224,12 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "out_prod_id_grp(b,grad,ids)", "glu_back(grad,a,b)", + + "ssm_conv_back(dy,sx,c)", + "ssm_scan_back(grad,...)", }; -static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); +static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -5807,6 +5813,83 @@ struct ggml_tensor * ggml_ssm_scan( return result; } +// ggml_ssm_conv_back / ggml_ssm_scan_back (learning-llamas, S1-29b) + +struct ggml_tensor * ggml_ssm_conv_back( + struct ggml_context * ctx, + struct ggml_tensor * dy, + struct ggml_tensor * sx, + struct ggml_tensor * c) { + GGML_ASSERT(ggml_is_3d(sx)); + GGML_ASSERT(ggml_is_matrix(c)); + GGML_ASSERT(dy->type == GGML_TYPE_F32); + GGML_ASSERT(sx->type == GGML_TYPE_F32); + GGML_ASSERT(c->type == GGML_TYPE_F32); + + const int64_t d_conv = c->ne[0]; + const int64_t d_inner = c->ne[1]; + const int64_t n_t = sx->ne[0] - d_conv + 1; + const int64_t n_s = sx->ne[2]; + + GGML_ASSERT(sx->ne[1] == d_inner); + GGML_ASSERT(dy->ne[0] == d_inner); + GGML_ASSERT(dy->ne[1] == n_t); + GGML_ASSERT(dy->ne[2] == n_s); + + // d_sx has sx's shape, including the d_conv - 1 columns of leading state: those receive a + // gradient too (they are part of the convolution window), and dropping them would silently + // truncate the gradient at every sequence boundary. + struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, sx->ne[0], sx->ne[1], sx->ne[2]); + + result->op = GGML_OP_SSM_CONV_BACK; + result->src[0] = dy; + result->src[1] = sx; + result->src[2] = c; + + return result; +} + +struct ggml_tensor * ggml_ssm_scan_back( + struct ggml_context * ctx, + struct ggml_tensor * grad, + struct ggml_tensor * s, + struct ggml_tensor * x, + struct ggml_tensor * dt, + struct ggml_tensor * A, + struct ggml_tensor * B, + struct ggml_tensor * C, + struct ggml_tensor * ids) { + GGML_ASSERT(grad->type == GGML_TYPE_F32); + GGML_ASSERT(ids->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_are_same_shape(B, C)); + + // grad is the gradient of ssm_scan's PACKED dst: y, then the final states. + GGML_ASSERT(ggml_nelements(grad) == + ggml_nelements(x) + s->ne[0]*s->ne[1]*s->ne[2]*ids->ne[0]); + + // One op, five gradients. Packed in a fixed order, and the backward case views each region + // back onto its source -- the same trick ggml_ssm_scan itself uses for y + states. + const int64_t n_packed = ggml_nelements(s) + + ggml_nelements(x) + + ggml_nelements(dt) + + ggml_nelements(B) + + ggml_nelements(C); + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_packed); + + result->op = GGML_OP_SSM_SCAN_BACK; + result->src[0] = grad; + result->src[1] = s; + result->src[2] = x; + result->src[3] = dt; + result->src[4] = A; + result->src[5] = B; + result->src[6] = C; + result->src[7] = ids; + + return result; +} + // ggml_win_part struct ggml_tensor * ggml_win_part( @@ -7340,6 +7423,91 @@ static void ggml_compute_backward( ggml_add_or_set(ctx, cgraph, isrc0, gb); } } break; + case GGML_OP_SSM_CONV: { + // learning-llamas (S1-29b). src0 = sx (the padded input window), src1 = c (the conv + // weight). Only sx takes a gradient: c is a frozen base weight on the LoRA path, and + // training it is ROADMAP E8. + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_ssm_conv_back(ctx, grad, src0, src1)); + } + GGML_ASSERT(!src1_needs_grads && "SSM conv-weight gradients are not implemented (ROADMAP E8)"); + } break; + case GGML_OP_SSM_SCAN: { + // learning-llamas (S1-29b). Seven sources: s, x, dt, A, B, C, ids. + // + // Five of them take a gradient, and one op cannot return five tensors -- so + // ggml_ssm_scan_back returns a PACKED 1-D tensor, exactly as ggml_ssm_scan itself packs + // y with the final states, and each region is viewed back onto its source here. + // + // A (the decay matrix) and ids (I32) do not. A is a frozen base weight on the LoRA + // path; asserting is better than silently producing nothing, because "the model trained + // but A never moved" is not a thing anyone would notice. + struct ggml_tensor * ssm_s = tensor->src[0]; + struct ggml_tensor * ssm_x = tensor->src[1]; + struct ggml_tensor * ssm_dt = tensor->src[2]; + struct ggml_tensor * ssm_A = tensor->src[3]; + struct ggml_tensor * ssm_B = tensor->src[4]; + struct ggml_tensor * ssm_C = tensor->src[5]; + struct ggml_tensor * ssm_id = tensor->src[6]; + + const size_t issm_s = ggml_hash_find(hash_set, ssm_s); + const size_t issm_x = ggml_hash_find(hash_set, ssm_x); + const size_t issm_dt = ggml_hash_find(hash_set, ssm_dt); + const size_t issm_A = ggml_hash_find(hash_set, ssm_A); + const size_t issm_B = ggml_hash_find(hash_set, ssm_B); + const size_t issm_C = ggml_hash_find(hash_set, ssm_C); + + const bool need_s = issm_s != GGML_HASHSET_FULL && grads_needed[issm_s]; + const bool need_x = issm_x != GGML_HASHSET_FULL && grads_needed[issm_x]; + const bool need_dt = issm_dt != GGML_HASHSET_FULL && grads_needed[issm_dt]; + const bool need_A = issm_A != GGML_HASHSET_FULL && grads_needed[issm_A]; + const bool need_B = issm_B != GGML_HASHSET_FULL && grads_needed[issm_B]; + const bool need_C = issm_C != GGML_HASHSET_FULL && grads_needed[issm_C]; + + GGML_ASSERT(!need_A && "SSM A-matrix gradients are not implemented (ROADMAP E8)"); + + if (need_s || need_x || need_dt || need_B || need_C) { + struct ggml_tensor * gb = ggml_ssm_scan_back( + ctx, grad, ssm_s, ssm_x, ssm_dt, ssm_A, ssm_B, ssm_C, ssm_id); + + // Packed in this order: d_s | d_x | d_dt | d_B | d_C. + size_t off = 0; + + if (need_s) { + ggml_add_or_set(ctx, cgraph, issm_s, + ggml_reshape_4d(ctx, ggml_cont(ctx, ggml_view_1d(ctx, gb, ggml_nelements(ssm_s), off)), + ssm_s->ne[0], ssm_s->ne[1], ssm_s->ne[2], ssm_s->ne[3])); + } + off += ggml_nelements(ssm_s)*sizeof(float); + + if (need_x) { + ggml_add_or_set(ctx, cgraph, issm_x, + ggml_reshape_4d(ctx, ggml_cont(ctx, ggml_view_1d(ctx, gb, ggml_nelements(ssm_x), off)), + ssm_x->ne[0], ssm_x->ne[1], ssm_x->ne[2], ssm_x->ne[3])); + } + off += ggml_nelements(ssm_x)*sizeof(float); + + if (need_dt) { + ggml_add_or_set(ctx, cgraph, issm_dt, + ggml_reshape_4d(ctx, ggml_cont(ctx, ggml_view_1d(ctx, gb, ggml_nelements(ssm_dt), off)), + ssm_dt->ne[0], ssm_dt->ne[1], ssm_dt->ne[2], ssm_dt->ne[3])); + } + off += ggml_nelements(ssm_dt)*sizeof(float); + + if (need_B) { + ggml_add_or_set(ctx, cgraph, issm_B, + ggml_reshape_4d(ctx, ggml_cont(ctx, ggml_view_1d(ctx, gb, ggml_nelements(ssm_B), off)), + ssm_B->ne[0], ssm_B->ne[1], ssm_B->ne[2], ssm_B->ne[3])); + } + off += ggml_nelements(ssm_B)*sizeof(float); + + if (need_C) { + ggml_add_or_set(ctx, cgraph, issm_C, + ggml_reshape_4d(ctx, ggml_cont(ctx, ggml_view_1d(ctx, gb, ggml_nelements(ssm_C), off)), + ssm_C->ne[0], ssm_C->ne[1], ssm_C->ne[2], ssm_C->ne[3])); + } + } + } break; case GGML_OP_NONE: { // noop } break; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 05e41dae6bd8..2f8fdc8abd3f 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4206,8 +4206,25 @@ struct test_ssm_conv : public test_case { ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne_a.data()); + ggml_set_name(a, "sx"); + + // learning-llamas (S1-29b): ask for the gradient. + // + // This class called ggml_set_param ZERO times, so `grad -o SSM_CONV` requested no + // gradients, compared nothing, and printed `Backend CPU: OK` -- while SSM_CONV had no + // backward case at all and would have aborted the moment one was asked for. + // + // `b` is the conv weight: a frozen base weight on the LoRA path (ROADMAP E8), so it is not + // a param and the backward case asserts as much rather than silently producing nothing. + if (type == GGML_TYPE_F32) { + ggml_set_param(a); + } + ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne_b.data()); + ggml_set_name(b, "c"); + ggml_tensor * out = ggml_ssm_conv(ctx, a, b); + ggml_set_name(out, "out"); return out; } }; @@ -4302,8 +4319,39 @@ struct test_ssm_scan : public test_case { B = ggml_new_tensor_4d(ctx, type, d_state, n_group, n_seq_tokens, n_seqs); C = ggml_new_tensor_4d(ctx, type, d_state, n_group, n_seq_tokens, n_seqs); } + ggml_set_name(s, "s"); + ggml_set_name(x, "x"); + ggml_set_name(dt, "dt"); + ggml_set_name(A, "A"); + ggml_set_name(B, "B"); + ggml_set_name(C, "C"); + + // learning-llamas (S1-29b): ask for the gradients. + // + // This class called ggml_set_param ZERO times, so `grad -o SSM_SCAN` requested no + // gradients, compared nothing, and printed `Backend CPU: OK` -- while SSM_SCAN had no + // backward case at all and would have aborted the moment one was asked for. + // + // Five of the seven sources take a gradient: s, x, dt, B, C. A (the decay matrix) does not + // -- it is a frozen base weight on the LoRA path (ROADMAP E8) -- and ids is I32. + // + // The xbc_overlap variants make x, B and C VIEWS of one tensor, and ggml_set_param asserts + // op == GGML_OP_NONE, so those cases stay eval-only. That is a coverage gap and it is + // stated rather than hidden: overlapping x/B/C is a real llama.cpp layout, and its backward + // aliasing is exactly the sort of thing that goes wrong quietly. S1-31 owns it. + if (type == GGML_TYPE_F32 && !xbc_overlap) { + ggml_set_param(s); + ggml_set_param(x); + ggml_set_param(dt); + ggml_set_param(B); + ggml_set_param(C); + } + ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seqs); + ggml_set_name(ids, "ids"); + ggml_tensor * out = ggml_ssm_scan(ctx, s, x, dt, A, B, C, ids); + ggml_set_name(out, "out"); return out; } @@ -9282,6 +9330,11 @@ static std::vector> make_test_cases_eval() { } for (int64_t d_conv : {3, 4, 9}) { for (int64_t d_inner: {1024, 1536, 2048}) { + // learning-llamas (S1-29b): a TINY shape, so the MODE_GRAD case is not silently + // skipped for size once S1-30's kernel lands. n_t = 4, so the convolution genuinely + // slides rather than degenerating to a single window. + test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {4 - 1 + 4, 8, 2, 1}, {4, 8, 1, 1})); + test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {d_conv, d_inner, 1, 1}, {d_conv, d_inner, 1, 1})); test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {2 * d_conv, d_inner, 1, 1}, {d_conv, d_inner, 1, 1})); test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {d_conv, d_inner, 4, 1}, {d_conv, d_inner, 1, 1})); @@ -9312,6 +9365,21 @@ static std::vector> make_test_cases_eval() { } } + // learning-llamas (S1-29b): TINY gradient shapes. + // + // grad_nmax() is 10000 parameter elements and every shape below is far above it, so under + // MODE_GRAD they are all SILENTLY SKIPPED -- and the case still prints OK. Without these, the + // day S1-31's kernel lands there would be nothing for it to be checked by. + // + // Mamba-1 shape (head_dim == 1, so A is per-state) and Mamba-2 (head_dim > 1, A is scalar per + // head): the two branches are different code in the kernel and both need a gradient case. + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, /*d_state=*/ 4, /*head_dim=*/ 1, + /*n_head=*/ 3, /*n_group=*/ 1, + /*n_seq_tokens=*/ 3, /*n_seqs=*/ 2)); // Mamba-1, grad + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, /*d_state=*/ 4, /*head_dim=*/ 2, + /*n_head=*/ 2, /*n_group=*/ 1, + /*n_seq_tokens=*/ 3, /*n_seqs=*/ 2)); // Mamba-2, grad + test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 16, 1, 1024, 1, 32, 4)); // Mamba-1 test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 32, 4)); // Mamba-2 test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 256, 64, 8, 2, 32, 4)); // Falcon-H1