From 70c3397b60330d52b026f20135877c22eb2dd446 Mon Sep 17 00:00:00 2001 From: bssrdf Date: Wed, 28 Feb 2024 11:04:03 -0500 Subject: [PATCH] remove MultiheadAttention2; customized MultiheadAttention --- clip.hpp | 26 +++++++----------- ggml_extend.hpp | 72 +++++++------------------------------------------ 2 files changed, 19 insertions(+), 79 deletions(-) diff --git a/clip.hpp b/clip.hpp index 75d6de28..fde9ba29 100644 --- a/clip.hpp +++ b/clip.hpp @@ -530,29 +530,23 @@ public: : d_model(d_model), n_head(n_head), intermediate_size(intermediate_size) { - if(atten1) - blocks["self_attn"] = std::shared_ptr(new MultiheadAttention(d_model, n_head)); - else - blocks["self_attn"] = std::shared_ptr(new MultiheadAttention2(d_model, n_head)); + + blocks["self_attn"] = std::shared_ptr(new MultiheadAttention(d_model, n_head, true, atten1)); + blocks["layer_norm1"] = std::shared_ptr(new LayerNorm(d_model)); blocks["layer_norm2"] = std::shared_ptr(new LayerNorm(d_model)); blocks["mlp"] = std::shared_ptr(new CLIPMLP(d_model, intermediate_size)); } - struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, bool mask = true, bool atten1 = true) { + struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, bool mask = true) { // x: [N, n_token, d_model] + auto self_attn = std::dynamic_pointer_cast(blocks["self_attn"]); auto layer_norm1 = std::dynamic_pointer_cast(blocks["layer_norm1"]); auto layer_norm2 = std::dynamic_pointer_cast(blocks["layer_norm2"]); auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); - if(atten1){ - auto self_attn = std::dynamic_pointer_cast(blocks["self_attn"]); - x = ggml_add(ctx, x, self_attn->forward(ctx, layer_norm1->forward(ctx, x), mask)); - } - else{ - auto self_attn = std::dynamic_pointer_cast(blocks["self_attn"]); - x = ggml_add(ctx, x, self_attn->forward(ctx, layer_norm1->forward(ctx, x), mask)); - } + + x = ggml_add(ctx, x, self_attn->forward(ctx, layer_norm1->forward(ctx, x), mask)); x = ggml_add(ctx, x, mlp->forward(ctx, layer_norm2->forward(ctx, x))); return x; } @@ -575,7 +569,7 @@ public: } } - struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, int clip_skip = -1, bool mask = true, bool atten1 = true) { + struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, int clip_skip = -1, bool mask = true) { // x: [N, n_token, d_model] int layer_idx = n_layer - 1; LOG_DEBUG("clip_skip %d", clip_skip); @@ -590,7 +584,7 @@ public: } std::string name = "layers." + std::to_string(i); auto layer = std::dynamic_pointer_cast(blocks[name]); - x = layer->forward(ctx, x, mask, atten1); // [N, n_token, d_model] + x = layer->forward(ctx, x, mask); // [N, n_token, d_model] // LOG_DEBUG("layer %d", i); } return x; @@ -898,7 +892,7 @@ public: auto x = embeddings->forward(ctx, pixel_values); // [N, num_positions, embed_dim] x = pre_layernorm->forward(ctx, x); - x = encoder->forward(ctx, x, -1, false, false); + x = encoder->forward(ctx, x, -1, false); x = post_layernorm->forward(ctx, x); // [N, n_token, hidden_size] GGML_ASSERT(x->ne[3] == 1); diff --git a/ggml_extend.hpp b/ggml_extend.hpp index b9f1f001..ae03bf62 100644 --- a/ggml_extend.hpp +++ b/ggml_extend.hpp @@ -1255,14 +1255,17 @@ protected: int64_t n_head; bool bias; bool mask; + bool reshape2d; public: MultiheadAttention(int64_t embed_dim, int64_t n_head, - bool bias = true) + bool bias = true, + bool r2d = true) : embed_dim(embed_dim), n_head(n_head), - bias(bias) { + bias(bias), + reshape2d(r2d) { blocks["q_proj"] = std::shared_ptr(new Linear(embed_dim, embed_dim, bias)); blocks["k_proj"] = std::shared_ptr(new Linear(embed_dim, embed_dim, bias)); blocks["v_proj"] = std::shared_ptr(new Linear(embed_dim, embed_dim, bias)); @@ -1300,8 +1303,10 @@ public: kqv = ggml_reshape_4d(ctx, kqv, d_head, n_token, n_head, N); kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3)); // [N, n_token, n_head, d_head] - - x = ggml_reshape_2d(ctx, kqv, d_head * n_head, n_token * N); // [N * n_token, d_head * n_head] + if(reshape2d) + x = ggml_reshape_2d(ctx, kqv, d_head * n_head, n_token * N); // [N * n_token, d_head * n_head] + else + x = ggml_reshape_3d(ctx, kqv, d_head * n_head, n_token, N); // [N, n_token, d_head * n_head] x = out_proj->forward(ctx, x); return x; @@ -1309,63 +1314,4 @@ public: }; -class MultiheadAttention2 : public GGMLBlock { -protected: - int64_t embed_dim; - int64_t n_head; - bool bias; - bool mask; - -public: - MultiheadAttention2(int64_t embed_dim, - int64_t n_head, - bool bias = true) - : embed_dim(embed_dim), - n_head(n_head), - bias(bias) { - // printf("MultiheadAttention2: %d, %d, %d\n", embed_dim, n_head, bias?1:0); - blocks["q_proj"] = std::shared_ptr(new Linear(embed_dim, embed_dim, bias)); - blocks["k_proj"] = std::shared_ptr(new Linear(embed_dim, embed_dim, bias)); - blocks["v_proj"] = std::shared_ptr(new Linear(embed_dim, embed_dim, bias)); - blocks["out_proj"] = std::shared_ptr(new Linear(embed_dim, embed_dim, bias)); - } - - // x: [N, n_token, embed_dim] - struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, bool mask = false) { - auto q_proj = std::dynamic_pointer_cast(blocks["q_proj"]); - auto k_proj = std::dynamic_pointer_cast(blocks["k_proj"]); - auto v_proj = std::dynamic_pointer_cast(blocks["v_proj"]); - auto out_proj = std::dynamic_pointer_cast(blocks["out_proj"]); - - int64_t N = x->ne[2]; - int64_t n_token = x->ne[1]; - int64_t d_head = embed_dim / n_head; - - struct ggml_tensor* q = q_proj->forward(ctx, x); - q = ggml_reshape_4d(ctx, q, d_head, n_head, n_token, N); // [N, n_token, n_head, d_head] - q = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); // [N, n_head, n_token, d_head] - q = ggml_reshape_3d(ctx, q, d_head, n_token, n_head * N); // [N * n_head, n_token, d_head] - - struct ggml_tensor* k = k_proj->forward(ctx, x); - k = ggml_reshape_4d(ctx, k, d_head, n_head, n_token, N); // [N, n_token, n_head, d_head] - k = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [N, n_head, n_token, d_head] - // k = ggml_reshape_3d(ctx, k, d_head, n_token, n_head); // [N * n_head, n_token, d_head] !!BUG!! - k = ggml_reshape_3d(ctx, k, d_head, n_token, n_head * N); // [N * n_head, n_token, d_head] - - struct ggml_tensor* v = v_proj->forward(ctx, x); - v = ggml_reshape_4d(ctx, v, d_head, n_head, n_token, N); // [N, n_token, n_head, d_head] - v = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3)); // [N, n_head, d_head, n_token] - v = ggml_reshape_3d(ctx, v, n_token, d_head, n_head * N); // [N * n_head, d_head, n_token] - - struct ggml_tensor* kqv = ggml_nn_attention(ctx, q, k, v, mask); // [N * n_head, n_token, d_head] - - kqv = ggml_reshape_4d(ctx, kqv, d_head, n_token, n_head, N); - kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3)); // [N, n_token, n_head, d_head] - - x = ggml_reshape_3d(ctx, kqv, d_head * n_head, n_token, N); // [N * n_token, d_head * n_head] - x = out_proj->forward(ctx, x); - return x; - } -}; - #endif // __GGML_EXTEND__HPP__ \ No newline at end of file