diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 88eab7b4dc84..def72aa53dc6 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -340,6 +340,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ROPE_SCALING_ATTN_FACTOR, "%s.rope.scaling.attn_factor" }, { LLM_KV_ROPE_SCALING_ORIG_CTX_LEN, "%s.rope.scaling.original_context_length" }, { LLM_KV_ROPE_SCALING_FINETUNED, "%s.rope.scaling.finetuned" }, + { LLM_KV_ROPE_SCALING_YARN_TRUNCATE, "%s.rope.scaling.yarn_truncate" }, { LLM_KV_ROPE_SCALING_YARN_LOG_MUL, "%s.rope.scaling.yarn_log_multiplier" }, { LLM_KV_ROPE_SCALING_YARN_EXT_FACTOR, "%s.rope.scaling.yarn_ext_factor" }, { LLM_KV_ROPE_SCALING_YARN_ATTN_FACTOR, "%s.rope.scaling.yarn_attn_factor" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 2d8ad6b07702..d7c6a7a0aaf8 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -345,6 +345,7 @@ enum llm_kv { LLM_KV_ROPE_SCALING_ATTN_FACTOR, LLM_KV_ROPE_SCALING_ORIG_CTX_LEN, LLM_KV_ROPE_SCALING_FINETUNED, + LLM_KV_ROPE_SCALING_YARN_TRUNCATE, LLM_KV_ROPE_SCALING_YARN_LOG_MUL, LLM_KV_ROPE_SCALING_YARN_EXT_FACTOR, LLM_KV_ROPE_SCALING_YARN_ATTN_FACTOR, diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 064de68f36c2..d9b5d98e8ee4 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -63,6 +63,7 @@ struct llama_hparams { bool vocab_only; bool no_alloc; bool rope_finetuned; + bool yarn_truncate = true; bool use_par_res; bool swin_norm; bool norm_before_residual = false; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 463510caa768..a3b62b6bc558 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1974,16 +1974,23 @@ ggml_tensor * llama_kv_cache::build_rope_shift( tmp = ggml_rope_ext(ctx, tmp, shift, factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, yarn_ext_factor, yarn_attn_factor, yarn_beta_fast, yarn_beta_slow); + ggml_rope_set_truncate(tmp, hparams.yarn_truncate); // rotate fwd tmp = llama_mul_mat_hadamard(ctx, tmp, rot); tmp = ggml_cpy(ctx, tmp, cur); - } else { + } else if (hparams.yarn_truncate) { // we rotate only the first n_rot dimensions tmp = ggml_rope_ext_inplace(ctx, cur, shift, factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, yarn_ext_factor, yarn_attn_factor, yarn_beta_fast, yarn_beta_slow); + } else { + tmp = ggml_rope_ext(ctx, cur, + shift, factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + yarn_ext_factor, yarn_attn_factor, yarn_beta_fast, yarn_beta_slow); + ggml_rope_set_truncate(tmp, false); + tmp = ggml_cpy(ctx, tmp, cur); } return tmp; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 1c7cbd9005d9..c9ae3319f9c4 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -374,6 +374,7 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_ROPE_SCALING_ATTN_FACTOR, hparams.rope_attn_factor); add_kv(LLM_KV_ROPE_SCALING_ORIG_CTX_LEN, hparams.n_ctx_orig_yarn); add_kv(LLM_KV_ROPE_SCALING_FINETUNED, hparams.rope_finetuned); + add_kv(LLM_KV_ROPE_SCALING_YARN_TRUNCATE, hparams.yarn_truncate); add_kv(LLM_KV_ROPE_SCALING_YARN_LOG_MUL, hparams.rope_yarn_log_mul); add_kv(LLM_KV_ROPE_SCALING_YARN_EXT_FACTOR, hparams.yarn_ext_factor); add_kv(LLM_KV_ROPE_SCALING_YARN_ATTN_FACTOR, hparams.yarn_attn_factor); diff --git a/src/models/plamo3.cpp b/src/models/plamo3.cpp index f8235f8a8323..c818742b1e3e 100644 --- a/src/models/plamo3.cpp +++ b/src/models/plamo3.cpp @@ -2,6 +2,7 @@ void llama_model_plamo3::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ROPE_SCALING_YARN_TRUNCATE, hparams.yarn_truncate, false); const bool found_swa = ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false); if (found_swa && hparams.n_swa > 0) { hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; @@ -131,9 +132,11 @@ llama_model_plamo3::graph::graph(const llama_model & model, const llm_grap Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, ext_factor, attn_factor, beta_fast, beta_slow); + ggml_rope_set_truncate(Qcur, hparams.yarn_truncate); Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, ext_factor, attn_factor, beta_fast, beta_slow); + ggml_rope_set_truncate(Kcur, hparams.yarn_truncate); const float attn_scale = 1.0f / sqrtf(float(head_dim_q));