Skip to content

S1-25: MUL_MAT_ID and ADD_ID backward wiring #19

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: learning-llamas-base
Choose a base branch
from
Open

S1-25: MUL_MAT_ID and ADD_ID backward wiring #19

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
39 changes: 39 additions & 0 deletions ggml/include/ggml.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};

Expand Down Expand Up @@ -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
//
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 @@ -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;
}
Expand Down
131 changes: 129 additions & 2 deletions ggml/src/ggml.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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");

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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;
Expand Down
28 changes: 28 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 @@ -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);
Expand Down Expand Up @@ -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) {
Expand All @@ -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");
Expand Down