Skip to content

S1-29b: the SSM backward wiring S1-29 never landed #23

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dillon-blake wants to merge 1 commit into
base: ticket/S1-28-glu-backward
Choose a base branch
from
Open

S1-29b: the SSM backward wiring S1-29 never landed #23

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 46 additions & 0 deletions ggml/include/ggml.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};

Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down
10 changes: 10 additions & 0 deletions ggml/src/ggml-cpu/ggml-cpu.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
172 changes: 170 additions & 2 deletions ggml/src/ggml.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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");

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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;
Expand Down
68 changes: 68 additions & 0 deletions tests/test-backend-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
};
Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -9282,6 +9330,11 @@ static std::vector<std::unique_ptr<test_case>> 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}));
Expand Down Expand Up @@ -9312,6 +9365,21 @@ static std::vector<std::unique_ptr<test_case>> 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
Expand Down