From 48bcc4dcf9125daa1a65fed896d50dab9ab5aab5 Mon Sep 17 00:00:00 2001 From: xaedes Date: Sun, 7 May 2023 01:27:11 +0200 Subject: [PATCH] fix backward pass for add_at and change arguments to have same order as in view --- ggml.c | 60 ++++++++++++--------- ggml.h | 8 +-- tests/test-grad0.c | 132 ++++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 169 insertions(+), 31 deletions(-) diff --git a/ggml.c b/ggml.c index 39f29f2fc..5a917ae6a 100644 --- a/ggml.c +++ b/ggml.c @@ -5058,10 +5058,10 @@ struct ggml_tensor * ggml_add_at_impl( struct ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * b, - size_t offset, size_t nb1, size_t nb2, size_t nb3, + size_t offset, bool inplace) { GGML_ASSERT(ggml_nelements(b) <= ggml_nelements(a)); GGML_ASSERT(ggml_is_contiguous(a)); @@ -5076,10 +5076,10 @@ struct ggml_tensor * ggml_add_at_impl( struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); struct ggml_tensor * c = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 5); - ((int32_t *) c->data)[0] = offset; - ((int32_t *) c->data)[1] = nb1; - ((int32_t *) c->data)[2] = nb2; - ((int32_t *) c->data)[3] = nb3; + ((int32_t *) c->data)[0] = nb1; + ((int32_t *) c->data)[1] = nb2; + ((int32_t *) c->data)[2] = nb3; + ((int32_t *) c->data)[3] = offset; ((int32_t *) c->data)[4] = inplace ? 1 : 0; result->op = GGML_OP_ADD_AT; @@ -5095,22 +5095,22 @@ struct ggml_tensor * ggml_add_at( struct ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * b, - size_t offset, size_t nb1, size_t nb2, - size_t nb3) { - return ggml_add_at_impl(ctx, a, b, offset, nb1, nb2, nb3, false); + size_t nb3, + size_t offset) { + return ggml_add_at_impl(ctx, a, b, nb1, nb2, nb3, offset, false); } struct ggml_tensor * ggml_add_at_inplace( struct ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * b, - size_t offset, size_t nb1, size_t nb2, - size_t nb3) { - return ggml_add_at_impl(ctx, a, b, offset, nb1, nb2, nb3, true); + size_t nb3, + size_t offset) { + return ggml_add_at_impl(ctx, a, b, nb1, nb2, nb3, offset, true); } // ggml_sub @@ -8135,10 +8135,10 @@ static void ggml_compute_forward_add_at_f32( // view src0 and dst with these strides and data offset inbytes during add_at // nb0 is implicitely element_size because src0 and dst are contiguous - size_t offset = ((int32_t *) opt0->data)[0]; - size_t nb1 = ((int32_t *) opt0->data)[1]; - size_t nb2 = ((int32_t *) opt0->data)[2]; - size_t nb3 = ((int32_t *) opt0->data)[3]; + size_t nb1 = ((int32_t *) opt0->data)[0]; + size_t nb2 = ((int32_t *) opt0->data)[1]; + size_t nb3 = ((int32_t *) opt0->data)[2]; + size_t offset = ((int32_t *) opt0->data)[3]; bool inplace = (bool) ((int32_t *) opt0->data)[4]; if (!inplace && (params->type == GGML_TASK_INIT)) { @@ -13187,19 +13187,27 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor src0->grad = ggml_add_impl(ctx, src0->grad, tensor->grad, inplace); } if (src1->grad) { - size_t offset; - memcpy(&offset, tensor->padding, sizeof(size_t)); + GGML_ASSERT(ggml_nelements(tensor->opt[0]) == 5); + GGML_ASSERT(tensor->opt[0]->type == GGML_TYPE_I32); + const size_t nb1 = (( int32_t * ) tensor->opt[0]->data)[0]; + const size_t nb2 = (( int32_t * ) tensor->opt[0]->data)[1]; + const size_t nb3 = (( int32_t * ) tensor->opt[0]->data)[2]; + const size_t offset = (( int32_t * ) tensor->opt[0]->data)[3]; + + struct ggml_tensor * tensor_grad_view = ggml_view_4d(ctx, + tensor->grad, + src1->grad->ne[0], + src1->grad->ne[1], + src1->grad->ne[2], + src1->grad->ne[3], + nb1, nb2, nb3, offset); + src1->grad = ggml_add_impl(ctx, src1->grad, - ggml_view_3d(ctx, - tensor->grad, - tensor->ne[0], - tensor->ne[1], - tensor->ne[2], - tensor->nb[1], - tensor->nb[2], - offset), + ggml_reshape(ctx, + ggml_cont(ctx, tensor_grad_view), + src1->grad), inplace); } } break; @@ -13572,7 +13580,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor nb3 = (nb3 / n0) * ng; } - src0->grad = ggml_add_at_impl(ctx, src0->grad, tensor->grad, offset, nb1, nb2, nb3, inplace); + src0->grad = ggml_add_at_impl(ctx, src0->grad, tensor->grad, nb1, nb2, nb3, offset, inplace); } } break; case GGML_OP_PERMUTE: diff --git a/ggml.h b/ggml.h index 883915467..b50deb2b9 100644 --- a/ggml.h +++ b/ggml.h @@ -499,19 +499,19 @@ extern "C" { struct ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * b, - size_t offset, size_t nb1, size_t nb2, - size_t nb3); + size_t nb3, + size_t offset); GGML_API struct ggml_tensor * ggml_add_at_inplace( struct ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * b, - size_t offset, size_t nb1, size_t nb2, - size_t nb3); + size_t nb3, + size_t offset); GGML_API struct ggml_tensor * ggml_sub( struct ggml_context * ctx, diff --git a/tests/test-grad0.c b/tests/test-grad0.c index edb3c514a..202a71c8a 100644 --- a/tests/test-grad0.c +++ b/tests/test-grad0.c @@ -45,7 +45,8 @@ float frand() { } int irand(int n) { - return rand()%n; + if (n == 0) return 0; + else return rand()%n; } void get_random_dims(int64_t * dims, int ndims) { @@ -696,6 +697,135 @@ int main(int argc, const char ** argv) { } } + // add_at 1d + { + int64_t ne2[4] = { 1, 1, 1, 1 }; + + const int nargs = 2; + for (int ndims = 1; ndims <= 4; ++ndims) { + + x[0] = get_random_tensor(ctx0, ndims, ne, -1.0f, 1.0f); + ggml_set_param(ctx0, x[0]); + + get_random_dims(ne2, 1); + while ((ne2[0] > ne[0]) || (ne2[0] > ggml_nelements(x[0]))) { + get_random_dims(ne2, 1); + } + + x[1] = get_random_tensor(ctx0, 1, ne2, -1.0f, 1.0f); + ggml_set_param(ctx0, x[1]); + + const int max_offset = MAX(0, ggml_nelements(x[0]) - ggml_nelements(x[1])); + const int offset = irand(max_offset) * ggml_element_size(x[0]); + + struct ggml_tensor * f = ggml_sum(ctx0, ggml_add_at(ctx0, x[0], x[1], x[0]->nb[1], x[0]->nb[2], x[0]->nb[3], offset)); + + check_gradient("add_at 1d", ctx0, x, f, ndims, nargs, 1e-3f, 1e-3f, INFINITY); + } + } + + // add_at 2d + { + int64_t ne2[4] = { 1, 1, 1, 1 }; + int64_t max_offsets[4] = { 0, 0, 0, 0 }; + int64_t offsets[4] = { 0, 0, 0, 0 }; + + const int nargs = 2; + for (int ndims = 2; ndims <= 4; ++ndims) { + + x[0] = get_random_tensor(ctx0, ndims, ne, -1.0f, 1.0f); + ggml_set_param(ctx0, x[0]); + + get_random_dims(ne2, 2); + while ((ne2[0] > ne[0]) || (ne2[1] > ne[1]) || (ne2[0]*ne2[1] > ggml_nelements(x[0]))) { + get_random_dims(ne2, 2); + } + + x[1] = get_random_tensor(ctx0, 2, ne2, -1.0f, 1.0f); + ggml_set_param(ctx0, x[1]); + + max_offsets[0] = MAX(0, x[0]->ne[0] - x[1]->ne[0]); + max_offsets[1] = MAX(0, x[0]->ne[1] - x[1]->ne[1]); + offsets[0] = irand(max_offsets[0]) * x[0]->nb[0]; + offsets[1] = irand(max_offsets[1]) * x[0]->nb[1]; + const int offset = offsets[0] + offsets[1]; + + struct ggml_tensor * f = ggml_sum(ctx0, ggml_add_at(ctx0, x[0], x[1], x[0]->nb[1], x[0]->nb[2], x[0]->nb[3], offset)); + + check_gradient("add_at 2d", ctx0, x, f, ndims, nargs, 1e-3f, 1e-3f, INFINITY); + } + } + + // add_at 3d + { + int64_t ne2[4] = { 1, 1, 1, 1 }; + int64_t max_offsets[4] = { 0, 0, 0, 0 }; + int64_t offsets[4] = { 0, 0, 0, 0 }; + + const int nargs = 2; + for (int ndims = 3; ndims <= 4; ++ndims) { + + x[0] = get_random_tensor(ctx0, ndims, ne, -1.0f, 1.0f); + ggml_set_param(ctx0, x[0]); + + get_random_dims(ne2, 3); + while ((ne2[0] > ne[0]) || (ne2[1] > ne[1]) || (ne2[2] > ne[2]) || (ne2[0]*ne2[1]*ne2[2] > ggml_nelements(x[0]))) { + get_random_dims(ne2, 3); + } + + x[1] = get_random_tensor(ctx0, 3, ne2, -1.0f, 1.0f); + ggml_set_param(ctx0, x[1]); + + max_offsets[0] = MAX(0, x[0]->ne[0] - x[1]->ne[0]); + max_offsets[1] = MAX(0, x[0]->ne[1] - x[1]->ne[1]); + max_offsets[2] = MAX(0, x[0]->ne[2] - x[1]->ne[2]); + offsets[0] = irand(max_offsets[0]) * x[0]->nb[0]; + offsets[1] = irand(max_offsets[1]) * x[0]->nb[1]; + offsets[2] = irand(max_offsets[2]) * x[0]->nb[2]; + const int offset = offsets[0] + offsets[1] + offsets[2]; + + struct ggml_tensor * f = ggml_sum(ctx0, ggml_add_at(ctx0, x[0], x[1], x[0]->nb[1], x[0]->nb[2], x[0]->nb[3], offset)); + + check_gradient("add_at 3d", ctx0, x, f, ndims, nargs, 1e-3f, 1e-3f, INFINITY); + } + } + + // add_at 4d + { + int64_t ne2[4] = { 1, 1, 1, 1 }; + int64_t max_offsets[4] = { 0, 0, 0, 0 }; + int64_t offsets[4] = { 0, 0, 0, 0 }; + + const int nargs = 2; + for (int ndims = 4; ndims <= 4; ++ndims) { + + x[0] = get_random_tensor(ctx0, ndims, ne, -1.0f, 1.0f); + ggml_set_param(ctx0, x[0]); + + get_random_dims(ne2, 4); + while ((ne2[0] > ne[0]) || (ne2[1] > ne[1]) || (ne2[2] > ne[2]) || (ne2[3] > ne[3]) || (ne2[0]*ne2[1]*ne2[2]*ne2[3] > ggml_nelements(x[0]))) { + get_random_dims(ne2, 4); + } + + x[1] = get_random_tensor(ctx0, 4, ne2, -1.0f, 1.0f); + ggml_set_param(ctx0, x[1]); + + max_offsets[0] = MAX(0, x[0]->ne[0] - x[1]->ne[0]); + max_offsets[1] = MAX(0, x[0]->ne[1] - x[1]->ne[1]); + max_offsets[2] = MAX(0, x[0]->ne[2] - x[1]->ne[2]); + max_offsets[3] = MAX(0, x[0]->ne[3] - x[1]->ne[3]); + offsets[0] = irand(max_offsets[0]) * x[0]->nb[0]; + offsets[1] = irand(max_offsets[1]) * x[0]->nb[1]; + offsets[2] = irand(max_offsets[2]) * x[0]->nb[2]; + offsets[3] = irand(max_offsets[3]) * x[0]->nb[3]; + const int offset = offsets[0] + offsets[1] + offsets[2] + offsets[3]; + + struct ggml_tensor * f = ggml_sum(ctx0, ggml_add_at(ctx0, x[0], x[1], x[0]->nb[1], x[0]->nb[2], x[0]->nb[3], offset)); + + check_gradient("add_at 4d", ctx0, x, f, ndims, nargs, 1e-3f, 1e-3f, INFINITY); + } + } + // view_1d { const int nargs = 1;