diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 38f14ab4691a..856f62a610a3 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/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]; diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 42e47cf3ecdf..6b015e66101c 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/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); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 592d66d3728e..e1ecdca64c98 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/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__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) && diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 74a4bcbd15a6..d0da2ffbf759 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/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; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index a704f4e40e89..f63f86d2d4c5 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/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(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]; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index 2aaf2f196da5..319cb70ba661 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/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); diff --git a/ggml/src/ggml-metal/kernels/out_prod.metal b/ggml/src/ggml-metal/kernels/out_prod.metal index 5156dd0f0539..b63a87f4eeea 100644 --- a/ggml/src/ggml-metal/kernels/out_prod.metal +++ b/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; template [[host_name("kernel_out_prod_q5_K")]] kernel out_prod_t kernel_out_prod; template [[host_name("kernel_out_prod_q6_K")]] kernel out_prod_t kernel_out_prod; + +// 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 +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) 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; +template [[host_name("kernel_mul_mat_id_back_f16")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +#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; +#endif +template [[host_name("kernel_mul_mat_id_back_q4_0")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q4_1")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q5_0")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q5_1")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q8_0")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q2_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q3_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q4_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q5_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back; +template [[host_name("kernel_mul_mat_id_back_q6_K")]] kernel mul_mat_id_back_t kernel_mul_mat_id_back;