vbatts/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

bugfixes for backward pass of flash attention

Browse source
This commit is contained in:
xaedes 2023-05-29 23:45:58 +02:00
parent 22a7279ffb
commit 38560b6d51
No known key found for this signature in database
GPG key ID: 30030EDD817EA2B1
1 changed files with 253 additions and 234 deletions
Download patch file Download diff file Expand all files Collapse all files

93
ggml.c
View file

@ -6221,7 +6221,6 @@ struct ggml_tensor * ggml_flash_attn(
bool is_node = false; bool is_node = false;
if (q->grad || k->grad || v->grad) { if (q->grad || k->grad || v->grad) {
GGML_ASSERT(false); // TODO: implement backward
is_node = true; is_node = true;
} }
@ -12882,10 +12881,15 @@ static void ggml_compute_forward_flash_attn_back_f32(
//const int64_t nev2 = v->ne[2]; //const int64_t nev2 = v->ne[2];
//const int64_t nev3 = v->ne[3]; //const int64_t nev3 = v->ne[3];
const int64_t ned0 = d->ne[0];
const int64_t ned1 = d->ne[1];
//const int64_t ned2 = d->ne[2];
//const int64_t ned3 = d->ne[3];
const int64_t ne0 = dst->ne[0]; const int64_t ne0 = dst->ne[0];
const int64_t ne1 = dst->ne[1]; const int64_t ne1 = dst->ne[1];
//const int64_t ne2 = dst->ne[2]; const int64_t ne2 = dst->ne[2];
//const int64_t ne3 = dst->ne[3]; const int64_t ne3 = dst->ne[3];
const int nbk0 = k->nb[0]; const int nbk0 = k->nb[0];
const int nbk1 = k->nb[1]; const int nbk1 = k->nb[1];
@ -12923,8 +12927,8 @@ static void ggml_compute_forward_flash_attn_back_f32(
const int Mup = ggml_up(M, GGML_SOFT_MAX_UNROLL); const int Mup = ggml_up(M, GGML_SOFT_MAX_UNROLL);
const int mxDM = MAX(D, Mup); const int mxDM = MAX(D, Mup);
GGML_ASSERT(ne0 == D); // GGML_ASSERT(ne0 == D);
GGML_ASSERT(ne1 == N); // GGML_ASSERT(ne1 == N);
GGML_ASSERT(P >= 0); GGML_ASSERT(P >= 0);
GGML_ASSERT(nbq0 == sizeof(float)); GGML_ASSERT(nbq0 == sizeof(float));
@ -12934,10 +12938,12 @@ static void ggml_compute_forward_flash_attn_back_f32(
GGML_ASSERT(neq0 == D); GGML_ASSERT(neq0 == D);
GGML_ASSERT(nek0 == D); GGML_ASSERT(nek0 == D);
GGML_ASSERT(nev1 == D); GGML_ASSERT(nev1 == D);
GGML_ASSERT(ned0 == D);
GGML_ASSERT(neq1 == N); GGML_ASSERT(neq1 == N);
GGML_ASSERT(nek1 == N + P); GGML_ASSERT(nek1 == N + P);
GGML_ASSERT(nev1 == D); GGML_ASSERT(nev1 == D);
GGML_ASSERT(ned1 == N);
// dst cannot be transposed or permuted // dst cannot be transposed or permuted
GGML_ASSERT(nb0 == sizeof(float)); GGML_ASSERT(nb0 == sizeof(float));
@ -12946,6 +12952,9 @@ static void ggml_compute_forward_flash_attn_back_f32(
GGML_ASSERT(nb2 <= nb3); GGML_ASSERT(nb2 <= nb3);
if (params->type == GGML_TASK_INIT) { if (params->type == GGML_TASK_INIT) {
if (ith == 0) {
memset(dst->data, 0, nb0*ne0*ne1*ne2*ne3);
}
return; return;
} }
@ -12956,7 +12965,7 @@ static void ggml_compute_forward_flash_attn_back_f32(
// parallelize by q rows using ggml_vec_dot_f32 // parallelize by q rows using ggml_vec_dot_f32
// total rows in q // total rows in q
const int nr = neq1*neq2*neq3; const int nr = neq2*neq3;
// rows per thread // rows per thread
const int dr = (nr + nth - 1)/nth; const int dr = (nr + nth - 1)/nth;
@ -12971,9 +12980,10 @@ static void ggml_compute_forward_flash_attn_back_f32(
for (int ir = ir0; ir < ir1; ++ir) { for (int ir = ir0; ir < ir1; ++ir) {
// q indices // q indices
const int iq3 = ir/(neq2*neq1); const int iq3 = ir/(neq2);
const int iq2 = (ir - iq3*neq2*neq1)/neq1; const int iq2 = (ir - iq3*neq2)/neq2;
const int iq1 = (ir - iq3*neq2*neq1 - iq2*neq1); for ( int iq1 = 0; iq1 < neq1; ++iq1) {
// not sure about CACHE_LINE_SIZE_F32.. // not sure about CACHE_LINE_SIZE_F32..
// - maybe it must not be multiplied by 2 and excluded from .. in SM 1*(..) offset? // - maybe it must not be multiplied by 2 and excluded from .. in SM 1*(..) offset?
@ -13027,17 +13037,18 @@ static void ggml_compute_forward_flash_attn_back_f32(
ggml_float sump[GGML_SOFT_MAX_UNROLL] = { 0.0 }; ggml_float sump[GGML_SOFT_MAX_UNROLL] = { 0.0 };
for (int i = 0; i < Mup; i += GGML_SOFT_MAX_UNROLL) { for (int i = 0; i < Mup; i += GGML_SOFT_MAX_UNROLL) {
float * SS = SM + i; float * SR = S + i;
float * SW = SM + i;
for (int j = 0; j < GGML_SOFT_MAX_UNROLL; ++j) { for (int j = 0; j < GGML_SOFT_MAX_UNROLL; ++j) {
if (SS[j] == -INFINITY) { if (SR[j] == -INFINITY) {
SS[j] = 0.0f; SW[j] = 0.0f;
} else { } else {
ggml_fp16_t s = GGML_FP32_TO_FP16(SS[j] - max); ggml_fp16_t s = GGML_FP32_TO_FP16(SR[j] - max);
memcpy(&scvt[j], &s, sizeof(uint16_t)); memcpy(&scvt[j], &s, sizeof(uint16_t));
const float val = GGML_FP16_TO_FP32(table_exp_f16[scvt[j]]); const float val = GGML_FP16_TO_FP32(table_exp_f16[scvt[j]]);
sump[j] += (ggml_float)val; sump[j] += (ggml_float)val;
SS[j] = val; SW[j] = val;
} }
} }
} }
@ -13118,8 +13129,8 @@ static void ggml_compute_forward_flash_attn_back_f32(
// S = gradSM = d[:D,iq1,iq2,iq3] @ vcur // S = gradSM = d[:D,iq1,iq2,iq3] @ vcur
// S = d[:D,iq1,iq2,iq3] @ vcur // S = d[:D,iq1,iq2,iq3] @ vcur
// S[:M] += vcur[:,ic] * d[ic,iq1,iq2,iq3] // S[:M] += vcur[:M,ic] * d[ic,iq1,iq2,iq3]
ggml_vec_set_f32(D, S, 0); ggml_vec_set_f32(M, S, 0);
for (int64_t ic = 0; ic < D; ++ic) { for (int64_t ic = 0; ic < D; ++ic) {
// dst indices // dst indices
const int i1 = iq1; const int i1 = iq1;
@ -13129,7 +13140,7 @@ static void ggml_compute_forward_flash_attn_back_f32(
ggml_vec_mad_f32(M, ggml_vec_mad_f32(M,
S, S,
(float *) ((char *) v->data + ( ic*nbv1 + i2*nbv2 + i3*nbv3)), (float *) ((char *) v->data + ( ic*nbv1 + i2*nbv2 + i3*nbv3)),
*(float *) ((char *) d->data + (ic*nbd1 + i1*nbd2 + i2*nbd2 + i3*nbd3))); *(float *) ((char *) d->data + (ic*nbd0 + i1*nbd1 + i2*nbd2 + i3*nbd3)));
} }
// S = SM * (S - dot(SM, S)) // S = SM * (S - dot(SM, S))
@ -13140,10 +13151,15 @@ static void ggml_compute_forward_flash_attn_back_f32(
// S = diag_mask_zero(S, P) * scale // S = diag_mask_zero(S, P) * scale
if (masked) { if (masked) {
for (int64_t i = P + iq1 + 1; i < M; i++) { // for (int64_t i = P + iq1 + 1; i < M; i++) {
// S[i] = 0;
// }
for (int64_t i = P; i < M; i++) {
if (i > P + iq1) {
S[i] = 0; S[i] = 0;
} }
} }
}
ggml_vec_scale_f32(M, S, scale); ggml_vec_scale_f32(M, S, scale);
void * grad_q = (char *) dst->data; void * grad_q = (char *) dst->data;
@ -13170,18 +13186,20 @@ static void ggml_compute_forward_flash_attn_back_f32(
// //
// grad[q][:D,iq1,iq2,iq3] += S @ kcur // grad[q][:D,iq1,iq2,iq3] += S @ kcur
// grad[q][:D,iq1,iq2,iq3] += shape[M,1] @ shape[D,M] // grad[q][:D,iq1,iq2,iq3] += shape[M,1] @ shape[D,M]
// grad[q][ic,iq1,iq2,iq3] += dot(kcur[:,ic],S.T) // grad[q][:D,iq1,iq2,iq3] += S[ic] * kcur[:D,ic]
// grad[q][ic,iq1,iq2,iq3] += dot(k[:D,ic,iq2,iq3],S.T) //
//// grad[q][ic,iq1,iq2,iq3] += dot(kcur[:,ic],S.T)
//// grad[q][ic,iq1,iq2,iq3] += dot(k[:D,ic,iq2,iq3],S.T)
for (int64_t ic = 0; ic < M; ++ic) { for (int64_t ic = 0; ic < M; ++ic) {
// dst indices // dst indices
const int i1 = iq1; const int i1 = iq1;
const int i2 = iq2; const int i2 = iq2;
const int i3 = iq3; const int i3 = iq3;
ggml_vec_dot_f32(D, ggml_vec_mad_f32(D,
(float *) ((char *) grad_q + (ic*nb0 + i1*nbgq1 + i2*nbgq2 + i3*nbgq3)), (float *) ((char *) grad_q + (i1*nbgq1 + i2*nbgq2 + i3*nbgq3)),
(float *) ((char *) k->data + (ic*nbk1 + i2*nbk2 + i3*nbk3)), (float *) ((char *) k->data + (ic*nbk1 + i2*nbk2 + i3*nbk3)),
S); S[ic]);
} }
// grad[k][:D,:M,iq2,iq3] += S.T @ qcur // grad[k][:D,:M,iq2,iq3] += S.T @ qcur
@ -13193,12 +13211,12 @@ static void ggml_compute_forward_flash_attn_back_f32(
const int i2 = iq2; const int i2 = iq2;
const int i3 = iq3; const int i3 = iq3;
ggml_vec_set_f32(D, // ggml_vec_set_f32(D,
(float *) ((char *) grad_k + (ic*nbgk1 + i2*nbgk2 + i3*nbgk3)), // (float *) ((char *) grad_k + (ic*nbgk1 + i2*nbgk2 + i3*nbgk3)),
0); // 0);
ggml_vec_mad_f32(D, ggml_vec_mad_f32(D,
(float *) ((char *) grad_k + (ic*nbgk1 + i2*nbgk2 + i3*nbgk3)), (float *) ((char *) grad_k + (ic*nbgk1 + i2*nbgk2 + i3*nbgk3)),
(float *) ((char *) q->data + (i1*nbk1 + i2*nbk2 + i3*nbk3)), (float *) ((char *) q->data + (i1*nbq1 + i2*nbq2 + i3*nbq3)),
S[ic]); S[ic]);
} }
@ -13211,13 +13229,14 @@ static void ggml_compute_forward_flash_attn_back_f32(
const int i2 = iq2; const int i2 = iq2;
const int i3 = iq3; const int i3 = iq3;
ggml_vec_set_f32(M, // ggml_vec_set_f32(M,
(float *) ((char *) grad_v + ( ic*nbgv1 + i2*nbgv2 + i3*nbgv3)), // (float *) ((char *) grad_v + ( ic*nbgv1 + i2*nbgv2 + i3*nbgv3)),
0); // 0);
ggml_vec_mad_f32(M, ggml_vec_mad_f32(M,
(float *) ((char *) grad_v + ( ic*nbgv1 + i2*nbgv2 + i3*nbgv3)), (float *) ((char *) grad_v + ( ic*nbgv1 + i2*nbgv2 + i3*nbgv3)),
SM, SM,
*(float *) ((char *) d->data + (ic*nbd1 + i1*nbd2 + i2*nbd2 + i3*nbd3))); *(float *) ((char *) d->data + (ic*nbd0 + i1*nbd1 + i2*nbd2 + i3*nbd3)));
}
} }
} }
} }
@ -14475,9 +14494,9 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor
bool masked = t != 0; bool masked = t != 0;
flash_grad = flash_grad =
ggml_flash_attn_back(ctx, ggml_flash_attn_back(ctx,
src0->grad, src0,
src1->grad, src1,
tensor->opt[0]->grad, tensor->opt[0],
tensor->grad, tensor->grad,
masked); masked);
} }
@ -14509,7 +14528,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor
} break; } break;
case 4: case 4:
{ {
grad_q = ggml_view_3d(ctx, grad_q = ggml_view_4d(ctx,
flash_grad, flash_grad,
src0->ne[0], src0->ne[0],
src0->ne[1], src0->ne[1],
@ -14555,7 +14574,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor
} break; } break;
case 4: case 4:
{ {
grad_k = ggml_view_3d(ctx, grad_k = ggml_view_4d(ctx,
flash_grad, flash_grad,
src1->ne[0], src1->ne[0],
src1->ne[1], src1->ne[1],
@ -14604,7 +14623,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor
} break; } break;
case 4: case 4:
{ {
grad_v = ggml_view_3d(ctx, grad_v = ggml_view_4d(ctx,
flash_grad, flash_grad,
opt0->ne[0], opt0->ne[0],
opt0->ne[1], opt0->ne[1],