diff --git a/src/models/lfm-base.cpp b/src/models/lfm-base.cpp index f5ff7d178..0dabf114b 100644 --- a/src/models/lfm-base.cpp +++ b/src/models/lfm-base.cpp @@ -22,11 +22,10 @@ ggml_tensor * llm_build_lfm_base::build_lfm_shortconv(const llama_model & model, const uint32_t d_conv = hparams.n_shortconv_l_cache - 1; const int64_t n_embd = hparams.n_embd; - // {n_embd, n_tokens} => {n_embd, n_seq_tokens, n_seqs} - cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs); - auto * bcx = build_lora_mm(model.layers[il].shortconv.in_proj, cur); cb(bcx, "model.layers.{}.conv.in_proj", il); + // {3*n_embd, n_tokens} => {3*n_embd, n_seq_tokens, n_seqs} + bcx = ggml_reshape_3d(ctx0, bcx, bcx->ne[0], n_seq_tokens, n_seqs); constexpr auto n_chunks = 3; GGML_ASSERT(bcx->ne[0] % n_chunks == 0); @@ -76,10 +75,10 @@ ggml_tensor * llm_build_lfm_base::build_lfm_shortconv(const llama_model & model, cb(conv_out, "model.layers.{}.conv.conv", il); auto * y = ggml_mul(ctx0, c, conv_out); - y = build_lora_mm(model.layers[il].shortconv.out_proj, y); - cb(y, "model.layers.{}.conv.out_proj", il); // {n_embd, n_seq_tokens, n_seqs} => {n_embd, n_tokens} y = ggml_reshape_2d(ctx0, y, y->ne[0], n_seq_tokens * n_seqs); + y = build_lora_mm(model.layers[il].shortconv.out_proj, y); + cb(y, "model.layers.{}.conv.out_proj", il); return y; }