ggml: MUL_MAT_ID has a backward pass through frozen experts, so a LoRA trains on a MoE base (GGML_OP_MUL_MAT_ID_BACK; continuum card 14b64aa9)

ggml/include/ggml-rpc.h

@@ -8,10 +8,10 @@ extern "C" {
 
 #define RPC_PROTO_MAJOR_VERSION    5
 #define RPC_PROTO_MINOR_VERSION    1
-#define RPC_PROTO_PATCH_VERSION    1
+#define RPC_PROTO_PATCH_VERSION    2
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16

ggml/include/ggml.h

@@ -513,6 +513,7 @@ extern "C" {
         GGML_OP_MUL_MAT,
         GGML_OP_MUL_MAT_ID,
         GGML_OP_OUT_PROD,
+        GGML_OP_MUL_MAT_ID_BACK, // MUL_MAT_ID's gradient w.r.t. its activations (src1), through frozen experts
 
         GGML_OP_SCALE,
         GGML_OP_SET,
@@ -1466,6 +1467,25 @@ extern "C" {
             struct ggml_tensor  * ids,
             struct ggml_tensor  * expert_ptrs);
 
+    // The gradient of ggml_mul_mat_id(as, b, ids) with respect to b, through the experts held
+    // fixed (frozen, possibly quantized): the MoE half of training a LoRA on attention, whose
+    // gradient reaches earlier layers through every later expert FFN.
+    //
+    //   as:   [cols, rows, n_expert]       the forward's experts (F32, F16 or quantized)
+    //   grad: [rows, n_expert_used, n_tok] F32, the gradient of the forward's result
+    //   ids:  [n_expert_used, n_tok]       I32, the forward's routing
+    //   result: [cols, n_expert_used, n_tok] F32, one row per routed slot
+    //
+    //   result[:, i, t] = sum_r as[:, r, ids[i, t]] * grad[r, i, t]
+    //
+    // A forward whose b broadcast over the slots (b->ne[1] < n_expert_used) sums the slots back
+    // with ggml_repeat_back, so every slot row has exactly one writer.
+    GGML_API struct ggml_tensor * ggml_mul_mat_id_back(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * as,
+            struct ggml_tensor  * grad,
+            struct ggml_tensor  * ids);
+
     // A: m columns, n rows,
     // B: p columns, n rows,
     // result is m columns, p rows

ggml/src/ggml-cpu/ggml-cpu.c

@@ -1845,6 +1845,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_out_prod(params, tensor);
             } break;
+        case GGML_OP_MUL_MAT_ID_BACK:
+            {
+                ggml_compute_forward_mul_mat_id_back(params, tensor);
+            } break;
         case GGML_OP_SCALE:
             {
                 ggml_compute_forward_scale(params, tensor);
@@ -2335,6 +2339,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_MUL_MAT:
         case GGML_OP_MUL_MAT_ID:
         case GGML_OP_OUT_PROD:
+        case GGML_OP_MUL_MAT_ID_BACK:
             {
                 n_tasks = n_threads;
             } break;
@@ -2889,6 +2894,11 @@ struct ggml_cplan ggml_graph_plan(
                             cur = ggml_type_size(GGML_TYPE_F32) * ggml_out_prod_cpu_wsize(node->src[0]->ne[0]) * n_tasks;
                         }
                     } break;
+                case GGML_OP_MUL_MAT_ID_BACK:
+                    {
+                        const struct ggml_tensor * ids = node->src[2];
+                        cur = ggml_mul_mat_id_back_cpu_wsize(node->src[0]->ne[0], node->src[0]->ne[2], ids->ne[0]*ids->ne[1], n_tasks);
+                    } break;
                 case GGML_OP_SET_ROWS:
                     {
                         if (node->src[0]->type == GGML_TYPE_F16 && node->type != GGML_TYPE_F16) {

ggml/src/ggml-cpu/ggml-cpu.cpp

@@ -473,6 +473,9 @@ 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 || src1->type == GGML_TYPE_F16);
         case GGML_OP_GET_ROWS_BACK:
             return src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16;
+        case GGML_OP_MUL_MAT_ID_BACK:
+            return (src0->type == GGML_TYPE_F32 || ggml_get_type_traits(src0->type)->to_float != NULL) &&
+                src1->type == GGML_TYPE_F32 && op->src[2]->type == GGML_TYPE_I32 && op->type == GGML_TYPE_F32;
         case GGML_OP_OUT_PROD:
             return (src0->type == GGML_TYPE_F32 ||
                     ((src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && src1->ne[2] % src0->ne[2] == 0 && src1->ne[3] % src0->ne[3] == 0)) &&

ggml/src/ggml-cpu/ops.cpp

@@ -4469,6 +4469,108 @@ static void ggml_compute_forward_out_prod_converted_f32(
     }
 }
 
+// ggml_compute_forward_mul_mat_id_back
+//
+// dst[:, i, t] = sum_r as[:, r, ids[i, t]] * grad[r, i, t]: the activations' gradient of
+// mul_mat_id through its experts, held fixed. Task 0 groups the routed (slot, token) pairs by
+// expert with a counting sort (count, prefix sum, fill: one entry per pair, Fable on fork #82);
+// then each task takes whole experts, dequantizes one in tiles of
+// GGML_OUT_PROD_KT rows ONCE, and adds each tile row, scaled by its gradient, into every slot
+// routed to that expert. A slot has exactly one expert, so each dst row has one writer.
+void ggml_compute_forward_mul_mat_id_back(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    const ggml_tensor * as   = dst->src[0];
+    const ggml_tensor * grad = dst->src[1];
+    const ggml_tensor * ids  = dst->src[2];
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    const int64_t cols     = as->ne[0];
+    const int64_t rows     = as->ne[1];
+    const int64_t n_expert = as->ne[2];
+    const int64_t n_used   = ids->ne[0];
+    const int64_t n_tok    = ids->ne[1];
+    const int64_t n_slots  = n_used*n_tok;
+
+    GGML_ASSERT(dst->type == GGML_TYPE_F32 && dst->nb[0] == sizeof(float));
+    GGML_ASSERT(grad->nb[0] == sizeof(float));
+    GGML_ASSERT(as->nb[0] == ggml_type_size(as->type)); // expert rows are contiguous
+
+    ggml_to_float_t const to_float = ggml_get_type_traits(as->type)->to_float;
+    GGML_ASSERT(as->type == GGML_TYPE_F32 || to_float != NULL);
+
+    char *    wdata  = (char *) params->wdata;
+    int64_t * offset = (int64_t *) wdata;                    // [n_expert + 1]: expert e's pairs
+    int32_t * routed = (int32_t *) (offset + n_expert + 1);  // [n_slots] of (i, t), by expert
+    float *   tiles  = (float *) GGML_PAD((uintptr_t) (routed + n_slots*2), CACHE_LINE_SIZE);
+
+    const auto expert_of = [&](int64_t i, int64_t t) {
+        const int32_t e = *(const int32_t *) ((const char *) ids->data + i*ids->nb[0] + t*ids->nb[1]);
+        GGML_ASSERT(e >= 0 && e < n_expert);
+        return e;
+    };
+
+    if (ith == 0) {
+        memset(dst->data, 0, ggml_nbytes(dst));
+        memset(offset, 0, (n_expert + 1)*sizeof(int64_t));
+        for (int64_t t = 0; t < n_tok; ++t) {
+            for (int64_t i = 0; i < n_used; ++i) {
+                offset[expert_of(i, t) + 1]++;
+            }
+        }
+        for (int64_t e = 0; e < n_expert; ++e) {
+            offset[e + 1] += offset[e];
+        }
+        // fill: offset[e] advances to the end of e's range, which is where e + 1 begins
+        for (int64_t t = 0; t < n_tok; ++t) {
+            for (int64_t i = 0; i < n_used; ++i) {
+                int32_t * pair = routed + 2*offset[expert_of(i, t)]++;
+                pair[0] = (int32_t) i;
+                pair[1] = (int32_t) t;
+            }
+        }
+        // shift back: offset[e] is again the start of e's range
+        for (int64_t e = n_expert; e > 0; --e) {
+            offset[e] = offset[e - 1];
+        }
+        offset[0] = 0;
+    }
+    ggml_barrier(params->threadpool);
+
+    const int64_t KT = GGML_OUT_PROD_KT;
+    float * wt = tiles + ith*KT*cols; // [KT][cols]
+
+    for (int64_t e = ith; e < n_expert; e += nth) {
+        const int64_t n_routed = offset[e + 1] - offset[e];
+        if (n_routed == 0) {
+            continue;
+        }
+        const int32_t * pairs = routed + 2*offset[e];
+        for (int64_t k0 = 0; k0 < rows; k0 += KT) {
+            const int64_t kt = std::min(KT, rows - k0);
+            for (int64_t r = 0; r < kt; ++r) {
+                const char * row = (const char *) as->data + (k0 + r)*as->nb[1] + e*as->nb[2];
+                if (as->type == GGML_TYPE_F32) {
+                    memcpy(wt + r*cols, row, cols*sizeof(float));
+                } else {
+                    to_float(row, wt + r*cols, cols);
+                }
+            }
+            for (int64_t s = 0; s < n_routed; ++s) {
+                const int64_t i = pairs[2*s + 0];
+                const int64_t t = pairs[2*s + 1];
+                float * d = (float *) ((char *) dst->data + i*dst->nb[1] + t*dst->nb[2]);
+                const char * g = (const char *) grad->data + i*grad->nb[1] + t*grad->nb[2];
+                for (int64_t r = 0; r < kt; ++r) {
+                    ggml_vec_mad_f32((int) cols, d, wt + r*cols, *(const float *) (g + (k0 + r)*grad->nb[0]));
+                }
+            }
+        }
+    }
+}
+
 void ggml_compute_forward_out_prod(
         const ggml_compute_params * params,
         ggml_tensor * dst) {
@@ -5204,18 +5306,22 @@ static void ggml_compute_forward_get_rows_back_f32(
     memset(dst->data, 0, ggml_nbytes(dst));
 
     const int nc = src0->ne[0];
-    const int nr = ggml_nelements(src1);
 
     GGML_ASSERT( dst->ne[0] == nc);
     GGML_ASSERT(src0->nb[0] == sizeof(float));
 
-    for (int i = 0; i < nr; ++i) {
-        const int r = ((int32_t *) src1->data)[i];
+    // batch for batch (ggml_get_rows_back): index k of batch (i2, i3) adds its gradient row
+    // into that batch's row r of dst
+    for (int64_t i3 = 0; i3 < src1->ne[2]; ++i3) {
+        for (int64_t i2 = 0; i2 < src1->ne[1]; ++i2) {
+            for (int64_t k = 0; k < src1->ne[0]; ++k) {
+                const int32_t r = *(const int32_t *) ((const char *) src1->data + k*src1->nb[0] + i2*src1->nb[1] + i3*src1->nb[2]);
+                GGML_ASSERT(r >= 0 && r < dst->ne[1]);
 
-        ggml_vec_add_f32(nc,
-                (float *) ((char *)  dst->data + r*dst->nb[1]),
-                (float *) ((char *)  dst->data + r*dst->nb[1]),
-                (float *) ((char *) src0->data + i*src0->nb[1]));
+                float * d = (float *) ((char *) dst->data + r*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3]);
+                ggml_vec_add_f32(nc, d, d, (const float *) ((const char *) src0->data + k*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3]));
+            }
+        }
     }
 }
 

ggml/src/ggml-cpu/ops.h

@@ -49,6 +49,7 @@ void ggml_compute_forward_rms_norm_back(const struct ggml_compute_params * param
 void ggml_compute_forward_group_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 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_mul_mat_id_back(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);
@@ -96,6 +97,14 @@ void ggml_compute_forward_flash_attn_back(const struct ggml_compute_params * par
 static inline int64_t ggml_out_prod_cpu_wsize(int64_t ne0) {
     return GGML_OUT_PROD_KT*ne0 + GGML_OUT_PROD_BN*GGML_OUT_PROD_KT;
 }
+
+// ggml_compute_forward_mul_mat_id_back's work buffer, in bytes: the routing grouped by expert
+// with a counting sort (an offset per expert, plus one, then every (slot, token) pair once),
+// and each task's dequantized tile of GGML_OUT_PROD_KT expert rows.
+static inline size_t ggml_mul_mat_id_back_cpu_wsize(int64_t cols, int64_t n_expert, int64_t n_slots, int n_tasks) {
+    return (size_t) (n_expert + 1)*sizeof(int64_t) + (size_t) n_slots*2*sizeof(int32_t) +
+           (size_t) n_tasks*GGML_OUT_PROD_KT*cols*sizeof(float) + 2*CACHE_LINE_SIZE;
+}
 // FLASH_ATTN_BACK's blocks: BR query rows (pass 1) and BC KV positions (pass 2) share one conversion of
 // each K and V row. Its per-thread scratch, in floats; the work-size planner reserves this per task.
 #define GGML_FA_BACK_BR 8

ggml/src/ggml-metal/ggml-metal-device.m

@@ -1490,7 +1490,11 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                    op->src[0]->nb[0] == sizeof(float) && op->src[4]->nb[0] == sizeof(float) &&
                    (op->src[3] == NULL || op->src[3]->type == GGML_TYPE_F16);
         case GGML_OP_GET_ROWS_BACK:
+            // kernel_get_rows_back walks one flat index list over dst's rows: the 2-D form only;
+            // a batched one ([1, n_expert, n_tokens], a MoE's router weights) runs where it is
+            // implemented
             return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_I32 &&
+                   ggml_is_vector(op->src[1]) && op->ne[2] == 1 && op->ne[3] == 1 &&
                    ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op);
         case GGML_OP_CROSS_ENTROPY_LOSS:
             return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 &&

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -18275,7 +18275,9 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
                 }
             }
         case GGML_OP_GET_ROWS_BACK:
-            return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32;
+            // the shader walks one flat index list: the 2-D form only (a batched one, a MoE's
+            // router weights, runs where it is implemented), as CUDA's supports_op says
+            return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->ne[2] == 1 && op->ne[3] == 1;
         case GGML_OP_SET_ROWS:
             {
                 if ((op->src[0]->type != GGML_TYPE_F32 && op->src[0]->type != GGML_TYPE_F16) ||