diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index a82caa5e430..a8eab90a6f5 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -334,6 +334,41 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu(ggml_metal_l return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cross_entropy_loss(ggml_metal_library_t lib, const ggml_tensor * op) { + assert(op->op == GGML_OP_CROSS_ENTROPY_LOSS || op->op == GGML_OP_CROSS_ENTROPY_LOSS_BACK); + + char base[256]; + char name[256]; + + snprintf(base, 256, "kernel_cross_entropy_loss%s_f32", op->op == GGML_OP_CROSS_ENTROPY_LOSS_BACK ? "_back" : ""); + 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; +} + +// RMS_NORM_BACK, SOFT_MAX_BACK, REPEAT_BACK: one f32 kernel each, named after the op. +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_row_back(ggml_metal_library_t lib, const ggml_tensor * op) { + const char * kernel = nullptr; + switch (op->op) { + case GGML_OP_RMS_NORM_BACK: kernel = "kernel_rms_norm_back_f32"; break; + case GGML_OP_SOFT_MAX_BACK: kernel = "kernel_soft_max_back_f32"; break; + case GGML_OP_REPEAT_BACK: kernel = "kernel_repeat_back_f32"; break; + default: GGML_ABORT("fatal error"); + } + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, kernel); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, kernel, kernel, nullptr); + } + + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum(ggml_metal_library_t lib, const ggml_tensor * op) { assert(op->op == GGML_OP_SUM); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 2ce8ce86cc5..007895e451e 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -120,6 +120,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_unary struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_silu_back (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cross_entropy_loss(ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_row_back (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum_rows (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_blk (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_add (ggml_metal_library_t lib, const struct ggml_tensor * op); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 219096f7704..e5556fa6989 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1,4 +1,5 @@ #import "ggml-metal-device.h" +#include "ggml-metal-ops.h" #import "ggml-impl.h" #import "ggml-backend-impl.h" @@ -1467,6 +1468,27 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te default: return false; } + case GGML_OP_RMS_NORM_BACK: + case GGML_OP_SOFT_MAX_BACK: + // src0 the incoming gradient, src1 the forward input/output; rows contiguous. + // SOFT_MAX_BACK carries no max_bias (the CPU reference asserts the same). + return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && + ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op) && + (op->op != GGML_OP_SOFT_MAX_BACK || ggml_get_op_params_f32(op, 1) == 0.0f); + case GGML_OP_REPEAT_BACK: + return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op); + case GGML_OP_CROSS_ENTROPY_LOSS: + return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && + ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]); + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + return op->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && op->src[2]->type == GGML_TYPE_F32 && + ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op->src[2]) && ggml_is_contiguous(op); + case GGML_OP_OUT_PROD: + // The backward through a frozen weight: out_prod(W, grad^T). Composed on this + // backend from a dequantizing copy, two transposing copies and the F32 simdgroup + // 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_SILU_BACK: return (op->src[0]->type == GGML_TYPE_F32) && (op->src[1]->type == GGML_TYPE_F32) && @@ -1782,6 +1804,17 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te default: return false; } + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + switch (op->type) { + case GGML_TYPE_F32: + return true; + default: + return false; + } case GGML_TYPE_I32: return op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_I32; default: diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 62dad0d998d..6359d2bf017 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -880,6 +880,23 @@ typedef struct { int32_t n_head_log2; } ggml_metal_kargs_soft_max; +typedef struct { + int32_t ne00; + int32_t nr; + uint64_t nb01; + uint64_t nb11; + uint64_t nb1; +} ggml_metal_kargs_cross_entropy_loss; + +// a per-row backward with one scalar parameter (eps for RMS_NORM_BACK, scale for SOFT_MAX_BACK) +typedef struct { + int32_t ne00; + uint64_t nb01; + uint64_t nb11; + uint64_t nb1; + float param; +} ggml_metal_kargs_row_back; + typedef struct { int64_t ne00; int64_t ne01; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 2f1f9be8b3d..57963160956 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -300,6 +300,27 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { { n_fuse = ggml_metal_op_unary(ctx, idx); } break; + case GGML_OP_OUT_PROD: + { + n_fuse = ggml_metal_op_out_prod(ctx, idx); + } break; + case GGML_OP_CROSS_ENTROPY_LOSS: + { + n_fuse = ggml_metal_op_cross_entropy_loss(ctx, idx); + } break; + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + { + n_fuse = ggml_metal_op_cross_entropy_loss_back(ctx, idx); + } break; + case GGML_OP_RMS_NORM_BACK: + case GGML_OP_SOFT_MAX_BACK: + { + n_fuse = ggml_metal_op_row_back(ctx, idx); + } break; + case GGML_OP_REPEAT_BACK: + { + n_fuse = ggml_metal_op_repeat_back(ctx, idx); + } break; case GGML_OP_SILU_BACK: { n_fuse = ggml_metal_op_silu_back(ctx, idx); @@ -918,8 +939,7 @@ int ggml_metal_op_glu(ggml_metal_op_t ctx, int idx) { return 1; } -int ggml_metal_op_sum(ggml_metal_op_t ctx, int idx) { - ggml_tensor * op = ctx->node(idx); +static int ggml_metal_op_sum_impl(ggml_metal_op_t ctx, const ggml_tensor * op) { ggml_metal_library_t lib = ctx->lib; ggml_metal_encoder_t enc = ctx->enc; @@ -955,6 +975,10 @@ int ggml_metal_op_sum(ggml_metal_op_t ctx, int idx) { return 1; } +int ggml_metal_op_sum(ggml_metal_op_t ctx, int idx) { + return ggml_metal_op_sum_impl(ctx, ctx->node(idx)); +} + int ggml_metal_op_sum_rows(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); @@ -2106,23 +2130,22 @@ int ggml_metal_op_set(ggml_metal_op_t ctx, int idx) { return 1; } -int ggml_metal_op_cpy(ggml_metal_op_t ctx, int idx) { - ggml_tensor * op = ctx->node(idx); +static int ggml_metal_op_cpy_impl(ggml_metal_op_t ctx, const ggml_tensor * src0, const ggml_tensor * op) { ggml_metal_library_t lib = ctx->lib; ggml_metal_encoder_t enc = ctx->enc; - GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); - GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); + GGML_TENSOR_LOCALS( int32_t, ne0, src0, ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, src0, nb); GGML_TENSOR_LOCALS( int32_t, ne, op, ne); GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); - auto pipeline = ggml_metal_library_get_pipeline_cpy(lib, op->src[0]->type, op->type); + auto pipeline = ggml_metal_library_get_pipeline_cpy(lib, src0->type, op->type); - GGML_ASSERT(ne00 % ggml_blck_size(op->src[0]->type) == 0); + GGML_ASSERT(ne00 % ggml_blck_size(src0->type) == 0); int64_t nk0 = ne00; - if (ggml_is_quantized(op->src[0]->type)) { + if (ggml_is_quantized(src0->type)) { nk0 = ne00/16; } else if (ggml_is_quantized(op->type)) { nk0 = ne00/ggml_blck_size(op->type); @@ -2134,7 +2157,7 @@ int ggml_metal_op_cpy(ggml_metal_op_t ctx, int idx) { int nrptg = 1; // TODO: relax this constraint in the future - if (ggml_blck_size(op->src[0]->type) == 1 && ggml_blck_size(op->type) == 1) { + if (ggml_blck_size(src0->type) == 1 && ggml_blck_size(op->type) == 1) { if (nth > nk0) { nrptg = (nth + nk0 - 1)/nk0; nth = nk0; @@ -2171,7 +2194,7 @@ int ggml_metal_op_cpy(ggml_metal_op_t ctx, int idx) { 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(op->src[0]), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(src0), 1); ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ggml_metal_encoder_dispatch_threadgroups(enc, nw0*(ne01 + nrptg - 1)/nrptg, ne02, ne03, nth, nrptg, 1); @@ -2179,6 +2202,11 @@ int ggml_metal_op_cpy(ggml_metal_op_t ctx, int idx) { return 1; } +int ggml_metal_op_cpy(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + return ggml_metal_op_cpy_impl(ctx, op->src[0], op); +} + int ggml_metal_op_pool_1d(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); @@ -2327,6 +2355,58 @@ int ggml_metal_op_pool_2d(ggml_metal_op_t ctx, int idx) { return 1; } +// The F32/quantized simdgroup matmul encode, over an explicit op (its src[0], src[1] and +// itself may be synthetic views: OUT_PROD composes it over scratch buffers). +static void ggml_metal_encode_mul_mm(ggml_metal_op_t ctx, const ggml_tensor * op) { + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); + GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); + GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + + const int16_t r2 = ne12/ne02; + const int16_t r3 = ne13/ne03; + + auto pipeline = ggml_metal_library_get_pipeline_mul_mm(lib, op); + + ggml_metal_kargs_mul_mm args = { + /*.ne00 =*/ ne00, + /*.ne02 =*/ ne02, + /*.nb01 =*/ nb01, + /*.nb02 =*/ nb02, + /*.nb03 =*/ nb03, + /*.ne12 =*/ ne12, + /*.nb10 =*/ nb10, + /*.nb11 =*/ nb11, + /*.nb12 =*/ nb12, + /*.nb13 =*/ nb13, + /*.ne0 =*/ ne0, + /*.ne1 =*/ ne1, + /*.r2 =*/ r2, + /*.r3 =*/ r3, + }; + + 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(op->src[0]), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + + const size_t smem = pipeline.smem; + + ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); + + const int nr0 = pipeline.nr0; + const int nr1 = pipeline.nr1; + const int nsg = pipeline.nsg; + + ggml_metal_encoder_dispatch_threadgroups(enc, ((ne11 + nr1 - 1) / nr1), ((ne01 + nr0 - 1) / nr0), ne12 * ne13, 32, nsg, 1); +} + int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); @@ -2485,40 +2565,7 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) { // default: break; //} - auto pipeline = ggml_metal_library_get_pipeline_mul_mm(lib, op); - - ggml_metal_kargs_mul_mm args = { - /*.ne00 =*/ ne00, - /*.ne02 =*/ ne02, - /*.nb01 =*/ nb01, - /*.nb02 =*/ nb02, - /*.nb03 =*/ nb03, - /*.ne12 =*/ ne12, - /*.nb10 =*/ nb10, - /*.nb11 =*/ nb11, - /*.nb12 =*/ nb12, - /*.nb13 =*/ nb13, - /*.ne0 =*/ ne0, - /*.ne1 =*/ ne1, - /*.r2 =*/ r2, - /*.r3 =*/ r3, - }; - - 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(op->src[0]), 1); - ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); - ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); - - const size_t smem = pipeline.smem; - - ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); - - const int nr0 = pipeline.nr0; - const int nr1 = pipeline.nr1; - const int nsg = pipeline.nsg; - - ggml_metal_encoder_dispatch_threadgroups(enc, ((ne11 + nr1 - 1) / nr1), ((ne01 + nr0 - 1) / nr0), ne12 * ne13, 32, nsg, 1); + ggml_metal_encode_mul_mm(ctx, op); } else { auto pipeline = ggml_metal_library_get_pipeline_mul_mv(lib, op); @@ -3796,6 +3843,398 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) { return n_fuse; } + +// A plain (unfused) ADD encode over an explicit op whose operands may be synthetic views. +static void ggml_metal_encode_add(ggml_metal_op_t ctx, const ggml_tensor * op) { + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); + GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); + GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + + ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); + ggml_metal_buffer_id bid_src1 = ggml_metal_get_buffer_id(op->src[1]); + ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); + + ggml_metal_kargs_bin args = { + /*.ne00 =*/ ne00, /*.ne01 =*/ ne01, /*.ne02 =*/ ne02, /*.ne03 =*/ ne03, + /*.nb00 =*/ nb00, /*.nb01 =*/ nb01, /*.nb02 =*/ nb02, /*.nb03 =*/ nb03, + /*.ne10 =*/ ne10, /*.ne11 =*/ ne11, /*.ne12 =*/ ne12, /*.ne13 =*/ ne13, + /*.nb10 =*/ nb10, /*.nb11 =*/ nb11, /*.nb12 =*/ nb12, /*.nb13 =*/ nb13, + /*.ne0 =*/ ne0, /*.ne1 =*/ ne1, /*.ne2 =*/ ne2, /*.ne3 =*/ ne3, + /*.nb0 =*/ nb0, /*.nb1 =*/ nb1, /*.nb2 =*/ nb2, /*.nb3 =*/ nb3, + /*.offs =*/ 0, + /*.o1 =*/ { bid_src1.offs }, + }; + bid_src1.offs = 0; + + struct ggml_metal_pipeline_with_params pipeline = ggml_metal_library_get_pipeline_bin(lib, op, 1); + + if (pipeline.c4) { + args.ne00 = ne00/4; + args.ne10 = ne10/4; + args.ne0 = ne0/4; + } + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, bid_src0, 1); + ggml_metal_encoder_set_buffer (enc, bid_src1, 2); + ggml_metal_encoder_set_buffer (enc, bid_dst, 3); + + if (pipeline.cnt) { + ggml_metal_encoder_dispatch_threadgroups(enc, args.ne0, ggml_nrows(op), 1, 1, 1, 1); + } else { + const int nth_max = MIN(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + int nth = 1; + while (2*nth < args.ne0 && nth < nth_max) { + nth *= 2; + } + ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); + } +} + +// OUT_PROD: dst[ne00, ne10] = sum over k of src0[:, k] (x) src1[:, k] โ€” the activation gradient +// through a FROZEN weight in a training graph (out_prod(W, grad^T), the mul_mat backward). +// This backend had no OUT_PROD at all, so a Mac could not train a LoRA on the weights it +// serves (Continuum card 55f314b8, the Metal half of the doctor lane). Composed from kernels +// the backend already has, over scratch that the graph allocator reserves after dst: +// +// for each block of rows k0..k0+rows of W: +// S = cpy(W[k0..], F32) (the dequantizing copy; skipped when W is F32) +// X = cpy(S^T) contiguous [rows, ne00] โ€” contraction dim first +// Y = cpy(src1[:, k0..]^T) contiguous [rows, ne10] +// P = mul_mm(X, Y) [ne00, ne10]; the first block writes dst directly +// dst += P for every block after the first +// +// Row blocks bound the transient to OUT_PROD_BLOCK_BYTES of F32 for any weight (the output +// head of a 27B is 3 GB in F32 whole), the same shape as the CUDA out_prod. First cut: +// correctness by test-backend-ops parity against the CPU backend; a fused dequantize-GEMM +// is the second cut if the ยง0 speed gate asks for it. +static constexpr int64_t OUT_PROD_BLOCK_BYTES = 64ll*1024*1024; + +static int64_t ggml_metal_op_out_prod_block_rows(const ggml_tensor * op) { + const ggml_tensor * src0 = op->src[0]; + const int64_t ne00 = src0->ne[0]; + const int64_t ne01 = src0->ne[1]; + int64_t rows = OUT_PROD_BLOCK_BYTES/(ne00*(int64_t) sizeof(float)); + rows = std::max(64, (rows/64)*64); + return std::min(rows, ne01); +} + +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]; + if (op->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32) { + return false; + } + if (src0->ne[2] != 1 || src0->ne[3] != 1 || src1->ne[2] != 1 || src1->ne[3] != 1) { + return false; + } + // src1 as the graph emits it: contiguous rows, or the TRANSPOSED view of the gradient + // (out_prod(W, ggml_transpose(grad)) is what the mul_mat backward builds; Cormac on + // llama.cpp#5). Either way Y is copied contiguous per block, so both layouts are + // exact; anything else is refused. + if (!ggml_is_contiguous(op)) { + return false; + } + if (src1->nb[0] != sizeof(float) && src1->nb[1] != sizeof(float)) { + return false; + } + if (src0->ne[0] % 16 != 0 || src0->ne[0] % ggml_blck_size(src0->type) != 0) { + return false; // the dequantizing copy works in 4x4 tiles along whole quant blocks of a row + } + switch (src0->type) { + case GGML_TYPE_F32: + return ggml_is_contiguous_rows(src0); + case GGML_TYPE_F16: + case GGML_TYPE_BF16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + return ggml_is_contiguous(src0); // kernel_cpy__f32 exists + default: + return false; + } +} + +size_t ggml_metal_op_out_prod_extra(const ggml_tensor * op) { + const ggml_tensor * src0 = op->src[0]; + const ggml_tensor * src1 = op->src[1]; + const int64_t ne00 = src0->ne[0]; + const int64_t ne10 = src1->ne[0]; + const int64_t rows = ggml_metal_op_out_prod_block_rows(op); + size_t res = 0; + if (src0->type != GGML_TYPE_F32) { + res += GGML_PAD(rows*ne00*sizeof(float), 256); // S + } + res += GGML_PAD(rows*ne00*sizeof(float), 256); // X + res += GGML_PAD(rows*ne10*sizeof(float), 256); // Y + if (rows < src0->ne[1]) { + res += GGML_PAD(ggml_nbytes(op), 256); // P + } + return res; +} + +// A synthetic 2D view over the same Metal buffer as `like`, at `data`. +static ggml_tensor ggml_metal_out_prod_view(const ggml_tensor * like, ggml_type type, + int64_t ne0, int64_t ne1, size_t nb0, size_t nb1, void * data) { + ggml_tensor t = {}; + t.type = type; + t.ne[0] = ne0; t.ne[1] = ne1; t.ne[2] = 1; t.ne[3] = 1; + t.nb[0] = nb0; t.nb[1] = nb1; t.nb[2] = nb1*ne1; t.nb[3] = t.nb[2]; + t.buffer = like->buffer; + t.data = data; + return t; +} + +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]; + const ggml_tensor * src1 = op->src[1]; + + GGML_ASSERT(ggml_metal_op_out_prod_supported(op)); + + const int64_t ne00 = src0->ne[0]; + const int64_t ne01 = src0->ne[1]; + const int64_t ne10 = src1->ne[0]; + GGML_ASSERT(src1->ne[1] == ne01); + GGML_ASSERT(op->ne[0] == ne00 && op->ne[1] == ne10); + + const int64_t rows_blk = ggml_metal_op_out_prod_block_rows(op); + + // scratch after dst, in dst's buffer (ggml_backend_metal_buffer_type_get_alloc_size) + char * scratch = (char *) op->data + ggml_nbytes(op); + char * s_data = nullptr; + if (src0->type != GGML_TYPE_F32) { + s_data = scratch; + scratch += GGML_PAD(rows_blk*ne00*sizeof(float), 256); + } + char * x_data = scratch; scratch += GGML_PAD(rows_blk*ne00*sizeof(float), 256); + char * y_data = scratch; scratch += GGML_PAD(rows_blk*ne10*sizeof(float), 256); + char * p_data = rows_blk < ne01 ? scratch : nullptr; + + for (int64_t k0 = 0; k0 < ne01; k0 += rows_blk) { + const int64_t rows = std::min(rows_blk, ne01 - k0); + + // S: this block of W, dequantized to F32, [ne00, rows] contiguous + ggml_tensor s_view; + if (src0->type == GGML_TYPE_F32) { + s_view = ggml_metal_out_prod_view(src0, GGML_TYPE_F32, ne00, rows, + src0->nb[0], src0->nb[1], (char *) src0->data + k0*src0->nb[1]); + } else { + ggml_tensor w_view = ggml_metal_out_prod_view(src0, src0->type, ne00, rows, + src0->nb[0], src0->nb[1], (char *) src0->data + k0*src0->nb[1]); + s_view = ggml_metal_out_prod_view(op, GGML_TYPE_F32, ne00, rows, + sizeof(float), ne00*sizeof(float), s_data); + ggml_metal_op_cpy_impl(ctx, &w_view, &s_view); + // the encoder dispatches concurrently: every kernel below reads what the one + // above wrote, so each step ends on a memory barrier (as cumsum does) + ggml_metal_op_concurrency_reset(ctx); + } + + // X = S^T: a transposed view of S copied contiguous, [rows, ne00] + ggml_tensor st_view = ggml_metal_out_prod_view(op, GGML_TYPE_F32, rows, ne00, + s_view.nb[1], s_view.nb[0], s_view.data); + ggml_tensor x = ggml_metal_out_prod_view(op, GGML_TYPE_F32, rows, ne00, + sizeof(float), rows*sizeof(float), x_data); + ggml_metal_op_cpy_impl(ctx, &st_view, &x); + ggml_metal_op_concurrency_reset(ctx); + + // Y = src1[:, k0..k0+rows]^T copied contiguous, [rows, ne10] + ggml_tensor bt_view = ggml_metal_out_prod_view(src1, GGML_TYPE_F32, rows, ne10, + src1->nb[1], src1->nb[0], (char *) src1->data + k0*src1->nb[1]); + ggml_tensor y = ggml_metal_out_prod_view(op, GGML_TYPE_F32, rows, ne10, + sizeof(float), rows*sizeof(float), y_data); + ggml_metal_op_cpy_impl(ctx, &bt_view, &y); + ggml_metal_op_concurrency_reset(ctx); + + // P = mul_mat(X, Y): [ne00, ne10], contraction over rows + const bool first = k0 == 0; + ggml_tensor p = ggml_metal_out_prod_view(op, GGML_TYPE_F32, ne00, ne10, + sizeof(float), ne00*sizeof(float), first ? op->data : (void *) p_data); + p.op = GGML_OP_MUL_MAT; + p.src[0] = &x; + p.src[1] = &y; + ggml_metal_encode_mul_mm(ctx, &p); + ggml_metal_op_concurrency_reset(ctx); + + if (!first) { + // dst = dst + P + ggml_tensor d = ggml_metal_out_prod_view(op, GGML_TYPE_F32, ne00, ne10, + sizeof(float), ne00*sizeof(float), op->data); + ggml_tensor acc = d; + acc.op = GGML_OP_ADD; + acc.src[0] = &d; + acc.src[1] = &p; + ggml_metal_encode_add(ctx, &acc); + ggml_metal_op_concurrency_reset(ctx); + } + } + + return 1; +} + + +// CROSS_ENTROPY_LOSS: a per-row kernel into scratch (nrows floats after the scalar dst, see +// ggml_backend_metal_buffer_type_get_alloc_size), then the existing sum kernel into dst[0]. +static int ggml_metal_op_cross_entropy_loss_rows(ggml_metal_op_t ctx, const ggml_tensor * op, ggml_metal_buffer_id bid_dst, uint64_t nb1) { + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + + const bool back = op->op == GGML_OP_CROSS_ENTROPY_LOSS_BACK; + const ggml_tensor * logits = back ? op->src[1] : op->src[0]; + const ggml_tensor * labels = back ? op->src[2] : op->src[1]; + + const int32_t ne00 = logits->ne[0]; + const int32_t nr = ggml_nrows(logits); + + ggml_metal_kargs_cross_entropy_loss args = { + /*.ne00 =*/ ne00, + /*.nr =*/ nr, + /*.nb01 =*/ logits->nb[1], + /*.nb11 =*/ labels->nb[1], + /*.nb1 =*/ nb1, + }; + + auto pipeline = ggml_metal_library_get_pipeline_cross_entropy_loss(lib, op); + + int nth = 32; // SIMD width + while (nth < ne00 && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { + nth *= 2; + } + nth = std::min(nth, 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); + if (back) { + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(logits), 2); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(labels), 3); + ggml_metal_encoder_set_buffer(enc, bid_dst, 4); + } else { + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(logits), 1); + ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(labels), 2); + ggml_metal_encoder_set_buffer(enc, bid_dst, 3); + } + ggml_metal_encoder_set_threadgroup_memory_size(enc, 32*sizeof(float), 0); + ggml_metal_encoder_dispatch_threadgroups(enc, nr, 1, 1, nth, 1, 1); + + return 1; +} + +int ggml_metal_op_cross_entropy_loss(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + const int64_t nr = ggml_nrows(op->src[0]); + + // the per-row losses, in the scratch after dst + ggml_metal_buffer_id bid_rows = ggml_metal_get_buffer_id(op); + bid_rows.offs += ggml_nbytes(op); + ggml_metal_op_cross_entropy_loss_rows(ctx, op, bid_rows, sizeof(float)); + ggml_metal_op_concurrency_reset(ctx); + + // dst[0] = sum(rows) + ggml_tensor rows = {}; + rows.type = GGML_TYPE_F32; + rows.ne[0] = nr; rows.ne[1] = 1; rows.ne[2] = 1; rows.ne[3] = 1; + rows.nb[0] = sizeof(float); rows.nb[1] = nr*sizeof(float); rows.nb[2] = rows.nb[1]; rows.nb[3] = rows.nb[1]; + rows.buffer = op->buffer; + rows.data = (char *) op->data + ggml_nbytes(op); + + ggml_tensor sum = *op; + sum.op = GGML_OP_SUM; + sum.src[0] = &rows; + sum.src[1] = nullptr; + ggml_metal_op_sum_impl(ctx, &sum); + + return 1; +} + +int ggml_metal_op_cross_entropy_loss_back(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + return ggml_metal_op_cross_entropy_loss_rows(ctx, op, ggml_metal_get_buffer_id(op), op->nb[1]); +} + + +// RMS_NORM_BACK and SOFT_MAX_BACK: a per-row kernel over contiguous rows, one scalar param. +int ggml_metal_op_row_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 int32_t ne00 = op->src[0]->ne[0]; + const int32_t nr = ggml_nrows(op->src[0]); + + ggml_metal_kargs_row_back args = { + /*.ne00 =*/ ne00, + /*.nb01 =*/ op->src[0]->nb[1], + /*.nb11 =*/ op->src[1]->nb[1], + /*.nb1 =*/ op->nb[1], + /*.param =*/ ggml_get_op_params_f32(op, 0), // eps, or scale + }; + + auto pipeline = ggml_metal_library_get_pipeline_row_back(lib, op); + + int nth = 32; // SIMD width + while (nth < ne00 && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { + nth *= 2; + } + nth = std::min(nth, 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(op->src[0]), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + ggml_metal_encoder_set_threadgroup_memory_size(enc, 32*sizeof(float), 0); + ggml_metal_encoder_dispatch_threadgroups(enc, nr, 1, 1, nth, 1, 1); + + return 1; +} + +// REPEAT_BACK: one threadgroup per dst row (k1, k2, k3), threads over k0. +int ggml_metal_op_repeat_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; + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + + ggml_metal_kargs_repeat args = { + /*.ne00 =*/ ne00, /*.ne01 =*/ ne01, /*.ne02 =*/ ne02, /*.ne03 =*/ ne03, + /*.nb00 =*/ nb00, /*.nb01 =*/ nb01, /*.nb02 =*/ nb02, /*.nb03 =*/ nb03, + /*.ne0 =*/ ne0, /*.ne1 =*/ ne1, /*.ne2 =*/ ne2, /*.ne3 =*/ ne3, + /*.nb0 =*/ nb0, /*.nb1 =*/ nb1, /*.nb2 =*/ nb2, /*.nb3 =*/ nb3, + }; + + auto pipeline = ggml_metal_library_get_pipeline_row_back(lib, op); + + const int nth = std::min(ne0, 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(op->src[0]), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); + ggml_metal_encoder_dispatch_threadgroups(enc, ne1, ne2, ne3, nth, 1, 1); + + return 1; +} + int ggml_metal_op_silu_back(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index 159a628d04a..a8a89693aa0 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -73,6 +73,13 @@ int ggml_metal_op_mul_mat_id (ggml_metal_op_t ctx, int idx); int ggml_metal_op_add_id (ggml_metal_op_t ctx, int idx); 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); +size_t ggml_metal_op_out_prod_extra(const struct ggml_tensor * op); +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); +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_row_back (ggml_metal_op_t ctx, int idx); +int ggml_metal_op_repeat_back (ggml_metal_op_t ctx, int idx); int ggml_metal_op_silu_back (ggml_metal_op_t ctx, int idx); int ggml_metal_op_l2_norm (ggml_metal_op_t ctx, int idx); int ggml_metal_op_group_norm (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index e2c5b28d6c7..fb829adb260 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -233,6 +233,14 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_ res += ggml_metal_op_flash_attn_ext_extra_tmp(tensor); res += ggml_metal_op_flash_attn_ext_extra_kv_f16(tensor); } break; + case GGML_OP_OUT_PROD: + { + res += ggml_metal_op_out_prod_extra(tensor); + } break; + case GGML_OP_CROSS_ENTROPY_LOSS: + { + res += ggml_nrows(tensor->src[0])*sizeof(float); // the per-row losses before the sum + } break; case GGML_OP_CUMSUM: case GGML_OP_ARGSORT: { diff --git a/ggml/src/ggml-metal/kernels/quantize.metal b/ggml/src/ggml-metal/kernels/quantize.metal index 59d0afe9695..bce881c78a9 100644 --- a/ggml/src/ggml-metal/kernels/quantize.metal +++ b/ggml/src/ggml-metal/kernels/quantize.metal @@ -146,6 +146,13 @@ template [[host_name("kernel_cpy_q4_1_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32< template [[host_name("kernel_cpy_q5_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32; template [[host_name("kernel_cpy_q5_1_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32; template [[host_name("kernel_cpy_q8_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32; +// K-quants: a block is 256 values = 16 tiles of 4x4. The backward of a LoRA through a resident +// K-quant base (OUT_PROD with a quantized src0) dequantizes weight rows through these. +template [[host_name("kernel_cpy_q2_K_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32; +template [[host_name("kernel_cpy_q3_K_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32; +template [[host_name("kernel_cpy_q4_K_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32; +template [[host_name("kernel_cpy_q5_K_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32; +template [[host_name("kernel_cpy_q6_K_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32; template [[host_name("kernel_cpy_tq2_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32; diff --git a/ggml/src/ggml-metal/kernels/softmax.metal b/ggml/src/ggml-metal/kernels/softmax.metal index f32fe293793..a6ad5baaa9b 100644 --- a/ggml/src/ggml-metal/kernels/softmax.metal +++ b/ggml/src/ggml-metal/kernels/softmax.metal @@ -221,3 +221,238 @@ template [[host_name("kernel_soft_max_f16")]] kernel kernel_soft_max_t kerne template [[host_name("kernel_soft_max_f32")]] kernel kernel_soft_max_t kernel_soft_max; template [[host_name("kernel_soft_max_f16_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4; template [[host_name("kernel_soft_max_f32_4")]] kernel kernel_soft_max_4_t kernel_soft_max_4; + + +// CROSS_ENTROPY_LOSS: one threadgroup per row; the row's loss goes to dst[row] (a scratch of +// nrows floats the allocator reserves after the scalar dst), already divided by nrows, and +// the existing sum kernel folds the rows into dst[0]. The back kernel writes the gradient +// row: (softmax(logits) - labels) * grad[0]/nrows. Both mirror the CPU reference exactly +// (ggml-cpu ops.cpp cross_entropy_loss{,_back}_f32) so the Metal backward of a training +// graph matches the CPU oracle (Continuum card 55f314b8, S3b). +template +static inline float ce_reduce_max(float v, threadgroup float * buf, uint sgitg, uint tiisg, uint3 tptg) { + float r = simd_max(v); + if (tptg.x > N_SIMDWIDTH) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sgitg == 0) { + buf[tiisg] = -INFINITY; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) { + buf[sgitg] = r; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + r = simd_max(buf[tiisg]); + } + return r; +} + +static inline float ce_reduce_sum(float v, threadgroup float * buf, uint sgitg, uint tiisg, uint3 tptg) { + float r = simd_sum(v); + if (tptg.x > N_SIMDWIDTH) { + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sgitg == 0) { + buf[tiisg] = 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiisg == 0) { + buf[sgitg] = r; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + r = simd_sum(buf[tiisg]); + } + return r; +} + +kernel void kernel_cross_entropy_loss_f32( + constant ggml_metal_kargs_cross_entropy_loss & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup float * buf [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint sgitg[[simdgroup_index_in_threadgroup]], + uint tiisg[[thread_index_in_simdgroup]], + uint3 tptg[[threads_per_threadgroup]]) { + const int32_t row = tgpig.x; + + device const float * s0 = (device const float *) (src0 + row*args.nb01); + device const float * s1 = (device const float *) (src1 + row*args.nb11); + + float lmax = -INFINITY; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + lmax = MAX(lmax, s0[i]); + } + const float max_val = ce_reduce_max(lmax, buf, sgitg, tiisg, tptg); + + float lsum = 0.0f; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + lsum += exp(s0[i] - max_val); + } + const float lse = log(ce_reduce_sum(lsum, buf, sgitg, tiisg, tptg)); + + float lloss = 0.0f; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + lloss += (s0[i] - max_val - lse)*s1[i]; + } + const float loss = ce_reduce_sum(lloss, buf, sgitg, tiisg, tptg); + + if (tpitg.x == 0) { + ((device float *) dst)[row] = -loss/(float) args.nr; + } +} + +kernel void kernel_cross_entropy_loss_back_f32( + constant ggml_metal_kargs_cross_entropy_loss & args, + device const char * src0, // grad: one float + device const char * src1, // logits + device const char * src2, // labels + device char * dst, + threadgroup float * buf [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint sgitg[[simdgroup_index_in_threadgroup]], + uint tiisg[[thread_index_in_simdgroup]], + uint3 tptg[[threads_per_threadgroup]]) { + const int32_t row = tgpig.x; + + const float d_by_nr = ((device const float *) src0)[0]/(float) args.nr; + + device const float * s0 = (device const float *) (src1 + row*args.nb01); + device const float * s1 = (device const float *) (src2 + row*args.nb11); + device float * d = (device float *) (dst + row*args.nb1); + + float lmax = -INFINITY; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + lmax = MAX(lmax, s0[i]); + } + const float max_val = ce_reduce_max(lmax, buf, sgitg, tiisg, tptg); + + float lsum = 0.0f; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + lsum += exp(s0[i] - max_val); + } + const float inv_sum = 1.0f/ce_reduce_sum(lsum, buf, sgitg, tiisg, tptg); + + // the gradient of -sum(y log softmax(x)) is softmax(x) * sum(y) - y, which equals + // softmax(x) - y only when the labels of a row sum to 1: a masked row (all-zero labels, + // a prompt token under completion-only training) must contribute nothing, not softmax(x) + float ly = 0.0f; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + ly += s1[i]; + } + const float sum_y = ce_reduce_sum(ly, buf, sgitg, tiisg, tptg); + + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + d[i] = (exp(s0[i] - max_val)*inv_sum*sum_y - s1[i])*d_by_nr; + } +} + + +// RMS_NORM_BACK: dx = (dz - x * sum(x*dz)/(sum(x*x) + eps*n)) / sqrt(sum(x*x)/n + eps) +// (the CPU reference, ggml-cpu ops.cpp rms_norm_back_f32). src0 = dz, src1 = x. +kernel void kernel_rms_norm_back_f32( + constant ggml_metal_kargs_row_back & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup float * buf [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint sgitg[[simdgroup_index_in_threadgroup]], + uint tiisg[[thread_index_in_simdgroup]], + uint3 tptg[[threads_per_threadgroup]]) { + const int32_t row = tgpig.x; + + device const float * dz = (device const float *) (src0 + row*args.nb01); + device const float * x = (device const float *) (src1 + row*args.nb11); + device float * dx = (device float *) (dst + row*args.nb1); + + float lxx = 0.0f; + float lxdz = 0.0f; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + lxx += x[i]*x[i]; + lxdz += x[i]*dz[i]; + } + const float sum_xx = ce_reduce_sum(lxx, buf, sgitg, tiisg, tptg); + const float sum_xdz = ce_reduce_sum(lxdz, buf, sgitg, tiisg, tptg); + + const float mean_eps = sum_xx/(float) args.ne00 + args.param; + const float sum_eps = sum_xx + args.param*(float) args.ne00; + const float rrms = 1.0f/sqrt(mean_eps); + const float scale_x = -sum_xdz/sum_eps; + + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + dx[i] = (dz[i] + x[i]*scale_x)*rrms; + } +} + +// SOFT_MAX_BACK: dx = scale * y * (dy - dot(y, dy)); src0 = dy, src1 = y (the forward output). +kernel void kernel_soft_max_back_f32( + constant ggml_metal_kargs_row_back & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup float * buf [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint sgitg[[simdgroup_index_in_threadgroup]], + uint tiisg[[thread_index_in_simdgroup]], + uint3 tptg[[threads_per_threadgroup]]) { + const int32_t row = tgpig.x; + + device const float * dy = (device const float *) (src0 + row*args.nb01); + device const float * y = (device const float *) (src1 + row*args.nb11); + device float * dx = (device float *) (dst + row*args.nb1); + + float ldot = 0.0f; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + ldot += y[i]*dy[i]; + } + const float dot_y_dy = ce_reduce_sum(ldot, buf, sgitg, tiisg, tptg); + + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + dx[i] = (dy[i] - dot_y_dy)*y[i]*args.param; + } +} + +// REPEAT_BACK: dst[k0..k3] = sum over the nr0..nr3 repeats of src0 โ€” the gradient of a +// broadcast. One thread per dst element, the repeats gathered in registers; src0 is read +// through its strides, so a non-contiguous view is exact. +kernel void kernel_repeat_back_f32( + constant ggml_metal_kargs_repeat & args, + device const char * src0, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + ushort3 tpitg[[thread_position_in_threadgroup]], + ushort3 ntg[[threads_per_threadgroup]]) { + const int k3 = tgpig.z; + const int k2 = tgpig.y; + const int k1 = tgpig.x; + + const int nr0 = args.ne00/args.ne0; + const int nr1 = args.ne01/args.ne1; + const int nr2 = args.ne02/args.ne2; + const int nr3 = args.ne03/args.ne3; + + device float * d = (device float *) (dst + k3*args.nb3 + k2*args.nb2 + k1*args.nb1); + + for (int k0 = tpitg.x; k0 < args.ne0; k0 += ntg.x) { + float acc = 0.0f; + for (int i3 = 0; i3 < nr3; i3++) { + for (int i2 = 0; i2 < nr2; i2++) { + for (int i1 = 0; i1 < nr1; i1++) { + device const char * base = src0 + + (i3*args.ne3 + k3)*args.nb03 + + (i2*args.ne2 + k2)*args.nb02 + + (i1*args.ne1 + k1)*args.nb01; + for (int i0 = 0; i0 < nr0; i0++) { + acc += *(device const float *) (base + (i0*args.ne0 + k0)*args.nb00); + } + } + } + } + d[k0] = acc; + } +} diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index f46dc3438d5..39ca569edd1 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9450,8 +9450,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_out_prod(GGML_TYPE_Q4_K, GGML_TYPE_F32, 256, 8, 600000, {1, 1}, {1, 1}, trans_b)); } - test_cases.emplace_back(new test_out_prod(GGML_TYPE_Q8_0, GGML_TYPE_F32, - 256, 8, 600000, {1, 1}, {1, 1}, false)); + // the q8_0 three-block case in both gradient layouts: the training graph hands + // out_prod a TRANSPOSED gradient (out_prod(W, ggml_transpose(grad))) + for (bool trans_b : {false, true}) { + test_cases.emplace_back(new test_out_prod(GGML_TYPE_Q8_0, GGML_TYPE_F32, + 256, 8, 600000, {1, 1}, {1, 1}, trans_b)); + } // add_id for (ggml_type type_a : {GGML_TYPE_F32}) {