metal: MUL_MAT_ID_BACK kernel, so a MoE LoRA's backward runs on the GPU (stacked on #82; continuum card 14b64aa9)

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -930,6 +930,21 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_meta
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mat_id_back(ggml_metal_library_t lib, const ggml_tensor * op) {
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_mul_mat_id_back_%s", ggml_type_name(op->src[0]->type));
+    snprintf(name, 256, "%s", base);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+    }
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_out_prod(ggml_metal_library_t lib, const ggml_tensor * op) {
     char base[256];
     char name[256];

ggml/src/ggml-metal/ggml-metal-device.h

@@ -142,6 +142,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext        (ggml_metal_library_t lib, const struct ggml_tensor * op, int nsg, int nxpsg, int r1ptg);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm            (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_out_prod          (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mat_id_back   (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_fp32       (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv            (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0    (ggml_metal_library_t lib, int ne02, int ne20);

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

@@ -1512,6 +1512,10 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
             // matmul (see ggml_metal_op_out_prod). 2D, src1 F32 with contiguous rows, src0 F32
             // or a type with a kernel_cpy_<type>_f32.
             return ggml_metal_op_out_prod_supported(op) && ggml_metal_device_get_props(dev)->has_simdgroup_mm;
+        case GGML_OP_MUL_MAT_ID_BACK:
+            // the activations' gradient through frozen experts (kernels/out_prod.metal)
+            return ggml_metal_op_mul_mat_id_back_supported(op) &&
+                (op->src[0]->type != GGML_TYPE_BF16 || ggml_metal_device_get_props(dev)->has_bfloat);
         case GGML_OP_SILU_BACK:
             return (op->src[0]->type == GGML_TYPE_F32) &&
                 (op->src[1]->type == GGML_TYPE_F32) &&

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -497,6 +497,20 @@ typedef struct {
     int32_t  dps3;
 } ggml_metal_kargs_out_prod;
 
+typedef struct {
+    int32_t  ne00; // expert cols
+    int32_t  ne01; // expert rows
+    uint64_t nb01;
+    uint64_t nb02; // one expert
+    uint64_t nbi0; // ids: slot, token
+    uint64_t nbi1;
+    uint64_t nbg0; // grad: row, slot, token
+    uint64_t nbg1;
+    uint64_t nbg2;
+    uint64_t nb1;  // dst: slot, token
+    uint64_t nb2;
+} ggml_metal_kargs_mul_mat_id_back;
+
 typedef struct {
     int32_t  ne00;
     int32_t  ne01;

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -304,6 +304,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
             {
                 n_fuse = ggml_metal_op_out_prod(ctx, idx);
             } break;
+        case GGML_OP_MUL_MAT_ID_BACK:
+            {
+                n_fuse = ggml_metal_op_mul_mat_id_back(ctx, idx);
+            } break;
         case GGML_OP_CROSS_ENTROPY_LOSS:
             {
                 n_fuse = ggml_metal_op_cross_entropy_loss(ctx, idx);
@@ -3984,6 +3988,8 @@ static void ggml_metal_encode_add(ggml_metal_op_t ctx, const ggml_tensor * op) {
 // dispatch, reading the weight in place in its own block order. Batched and broadcast
 // operands follow the CPU backend: dst and src1 share dims 2-3, and src0 is broadcast over
 // them (dst per src0 = ne2/ne02, ne3/ne03).
+static bool ggml_metal_frozen_weight_readable(const ggml_tensor * w);
+
 bool ggml_metal_op_out_prod_supported(const ggml_tensor * op) {
     const ggml_tensor * src0 = op->src[0];
     const ggml_tensor * src1 = op->src[1];
@@ -3994,11 +4000,18 @@ bool ggml_metal_op_out_prod_supported(const ggml_tensor * op) {
         op->ne[2] % src0->ne[2] != 0 || op->ne[3] % src0->ne[3] != 0) {
         return false;
     }
-    switch (src0->type) {
+    return ggml_metal_frozen_weight_readable(src0);
+}
+
+// A frozen weight the training kernels (kernel_out_prod, kernel_mul_mat_id_back) read in its
+// own layout: element by element for float types, else 16 consecutive elements of a row
+// dequantized from whole blocks.
+static bool ggml_metal_frozen_weight_readable(const ggml_tensor * w) {
+    switch (w->type) {
         case GGML_TYPE_F32:
         case GGML_TYPE_F16:
         case GGML_TYPE_BF16:
-            return src0->nb[0] == ggml_type_size(src0->type); // read element by element, any width
+            return w->nb[0] == ggml_type_size(w->type); // read element by element, any width
         case GGML_TYPE_Q4_0:
         case GGML_TYPE_Q4_1:
         case GGML_TYPE_Q5_0:
@@ -4009,13 +4022,64 @@ bool ggml_metal_op_out_prod_supported(const ggml_tensor * op) {
         case GGML_TYPE_Q4_K:
         case GGML_TYPE_Q5_K:
         case GGML_TYPE_Q6_K:
-            // dequantized 16 consecutive elements at a time, from whole blocks of a row
-            return ggml_is_contiguous_rows(src0) && src0->ne[0] % 16 == 0;
+            return ggml_is_contiguous_rows(w) && w->ne[0] % 16 == 0;
         default:
             return false;
     }
 }
 
+bool ggml_metal_op_mul_mat_id_back_supported(const ggml_tensor * op) {
+    const ggml_tensor * as   = op->src[0];
+    const ggml_tensor * grad = op->src[1];
+    const ggml_tensor * ids  = op->src[2];
+    return op->src[3] == NULL && // the gather's pointer table is not read here
+        op->type == GGML_TYPE_F32 && ggml_is_contiguous(op) &&
+        grad->type == GGML_TYPE_F32 && ids->type == GGML_TYPE_I32 &&
+        ggml_metal_frozen_weight_readable(as);
+}
+
+int ggml_metal_op_mul_mat_id_back(ggml_metal_op_t ctx, int idx) {
+    ggml_tensor * op = ctx->node(idx);
+
+    ggml_metal_library_t lib = ctx->lib;
+    ggml_metal_encoder_t enc = ctx->enc;
+
+    const ggml_tensor * as   = op->src[0];
+    const ggml_tensor * grad = op->src[1];
+    const ggml_tensor * ids  = op->src[2];
+
+    GGML_ASSERT(ggml_metal_op_mul_mat_id_back_supported(op));
+
+    auto pipeline = ggml_metal_library_get_pipeline_mul_mat_id_back(lib, op);
+
+    ggml_metal_kargs_mul_mat_id_back args = {
+        /*.ne00 =*/ (int32_t) as->ne[0],
+        /*.ne01 =*/ (int32_t) as->ne[1],
+        /*.nb01 =*/ as->nb[1],
+        /*.nb02 =*/ as->nb[2],
+        /*.nbi0 =*/ ids->nb[0],
+        /*.nbi1 =*/ ids->nb[1],
+        /*.nbg0 =*/ grad->nb[0],
+        /*.nbg1 =*/ grad->nb[1],
+        /*.nbg2 =*/ grad->nb[2],
+        /*.nb1  =*/ op->nb[1],
+        /*.nb2  =*/ op->nb[2],
+    };
+
+    const int32_t per_thread = ggml_is_quantized(as->type) ? as->ne[0]/16 : as->ne[0];
+    const int nth = std::max(1, std::min<int>(per_thread, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)));
+
+    ggml_metal_encoder_set_pipeline(enc, pipeline);
+    ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(as),   1);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(grad), 2);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(ids),  3);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),   4);
+    ggml_metal_encoder_dispatch_threadgroups(enc, ids->ne[0], ids->ne[1], 1, nth, 1, 1);
+
+    return 1;
+}
+
 int ggml_metal_op_out_prod(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
     const ggml_tensor * src0 = op->src[0];

ggml/src/ggml-metal/ggml-metal-ops.h

@@ -78,6 +78,8 @@ int ggml_metal_op_flash_attn_ext    (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_bin               (ggml_metal_op_t ctx, int idx);
 bool   ggml_metal_op_out_prod_supported(const struct ggml_tensor * op);
 int ggml_metal_op_out_prod          (ggml_metal_op_t ctx, int idx);
+bool   ggml_metal_op_mul_mat_id_back_supported(const struct ggml_tensor * op);
+int ggml_metal_op_mul_mat_id_back   (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_cross_entropy_loss(ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_cross_entropy_loss_back(ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_flash_attn_back(ggml_metal_op_t ctx, int idx);

ggml/src/ggml-metal/kernels/out_prod.metal

@@ -155,3 +155,66 @@ template [[host_name("kernel_out_prod_q3_K")]] kernel out_prod_t kernel_out_prod
 template [[host_name("kernel_out_prod_q4_K")]] kernel out_prod_t kernel_out_prod<float,  block_q4_K, QK_NL, dequantize_q4_K>;
 template [[host_name("kernel_out_prod_q5_K")]] kernel out_prod_t kernel_out_prod<float,  block_q5_K, QK_NL, dequantize_q5_K>;
 template [[host_name("kernel_out_prod_q6_K")]] kernel out_prod_t kernel_out_prod<float,  block_q6_K, QK_NL, dequantize_q6_K>;
+
+// 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 frozen experts (a LoRA on attention trains on a MoE base).
+// One threadgroup per routed slot (i, t); each thread owns columns of the slot's row and walks the
+// expert's rows, so every dst element has one writer and no expert is copied or transposed in
+// device memory. nl == 0 reads F32/F16/BF16 element by element; nl > 0 dequantizes 16
+// consecutive elements of a row in registers (ne00 % 16 == 0), as kernel_out_prod does.
+template<typename T0, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread float4x4 &)>
+kernel void kernel_mul_mat_id_back(
+        constant ggml_metal_kargs_mul_mat_id_back & args,
+        device const char * src0,
+        device const char * grad,
+        device const char * ids,
+        device       char * dst,
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort  tiitg[[thread_index_in_threadgroup]],
+        ushort3 ntg  [[threads_per_threadgroup]]) {
+    const int32_t i = tgpig.x;
+    const int32_t t = tgpig.y;
+    const int32_t e = *((device const int32_t *) (ids + i*args.nbi0 + t*args.nbi1));
+
+    device const char * as_e = src0 + e*args.nb02;
+    device const char * g    = grad + i*args.nbg1 + t*args.nbg2;
+    device       char * d    = dst  + i*args.nb1  + t*args.nb2;
+
+    if (nl == 0) {
+        for (int32_t c = tiitg; c < args.ne00; c += ntg.x) {
+            float acc = 0.0f;
+            for (int32_t r = 0; r < args.ne01; ++r) {
+                acc += (float) ((device const T0 *) (as_e + r*args.nb01))[c] * *((device const float *) (g + r*args.nbg0));
+            }
+            ((device float *) d)[c] = acc;
+        }
+    } else {
+        for (int32_t ind = tiitg; ind < args.ne00/16; ind += ntg.x) {
+            float4x4 acc = float4x4(0.0f);
+            for (int32_t r = 0; r < args.ne01; ++r) {
+                float4x4 w;
+                dequantize_func(((device const block_q *) (as_e + r*args.nb01)) + ind/nl, ind%nl, w);
+                acc += w * *((device const float *) (g + r*args.nbg0));
+            }
+            ((device float4x4 *) d)[ind] = acc;
+        }
+    }
+}
+
+typedef decltype(kernel_mul_mat_id_back<float, float4x4, 0, dequantize_f32>) mul_mat_id_back_t;
+
+template [[host_name("kernel_mul_mat_id_back_f32")]]  kernel mul_mat_id_back_t kernel_mul_mat_id_back<float,  float4x4,   0,     dequantize_f32>;
+template [[host_name("kernel_mul_mat_id_back_f16")]]  kernel mul_mat_id_back_t kernel_mul_mat_id_back<half,   float4x4,   0,     dequantize_f32>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_mul_mat_id_back_bf16")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back<bfloat, float4x4,   0,     dequantize_f32>;
+#endif
+template [[host_name("kernel_mul_mat_id_back_q4_0")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back<float,  block_q4_0, 2,     dequantize_q4_0>;
+template [[host_name("kernel_mul_mat_id_back_q4_1")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back<float,  block_q4_1, 2,     dequantize_q4_1>;
+template [[host_name("kernel_mul_mat_id_back_q5_0")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back<float,  block_q5_0, 2,     dequantize_q5_0>;
+template [[host_name("kernel_mul_mat_id_back_q5_1")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back<float,  block_q5_1, 2,     dequantize_q5_1>;
+template [[host_name("kernel_mul_mat_id_back_q8_0")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back<float,  block_q8_0, 2,     dequantize_q8_0>;
+template [[host_name("kernel_mul_mat_id_back_q2_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back<float,  block_q2_K, QK_NL, dequantize_q2_K>;
+template [[host_name("kernel_mul_mat_id_back_q3_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back<float,  block_q3_K, QK_NL, dequantize_q3_K>;
+template [[host_name("kernel_mul_mat_id_back_q4_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back<float,  block_q4_K, QK_NL, dequantize_q4_K>;
+template [[host_name("kernel_mul_mat_id_back_q5_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back<float,  block_q5_K, QK_NL, dequantize_q5_K>;
+template [[host_name("kernel_mul_mat_id_back_q6_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back<float,  block_q6_K, QK_NL, dequantize_q6_K>;