llama: implement NTK-By-Parts (NTKv2)

Author: cebtenzzre

examples/common.cpp

@@ -180,6 +180,18 @@ bool gpt_params_parse(int argc, char ** argv, gpt_params & params) {
                 break;
             }
             params.rope_freq_scale = std::stof(argv[i]);
+        } else if (arg == "--rope-ntk-factor") {
+            if (++i >= argc) {
+                invalid_param = true;
+                break;
+            }
+            params.rope_ntk_factor = std::stof(argv[i]);
+        } else if (arg == "--rope-extrapolation-factor") {
+            if (++i >= argc) {
+                invalid_param = true;
+                break;
+            }
+            params.rope_extrapolation_factor = std::stof(argv[i]);
         } else if (arg == "--memory-f32") {
             params.memory_f16 = false;
         } else if (arg == "--top-p") {
@@ -513,6 +525,9 @@ void gpt_print_usage(int /*argc*/, char ** argv, const gpt_params & params) {
     fprintf(stderr, "  -c N, --ctx-size N    size of the prompt context (default: %d)\n", params.n_ctx);
     fprintf(stderr, "  --rope-freq-base N    RoPE base frequency (default: %.1f)\n", params.rope_freq_base);
     fprintf(stderr, "  --rope-freq-scale N   RoPE frequency scaling factor (default: %g)\n", params.rope_freq_scale);
+    fprintf(stderr, "  --rope-ntk-factor N   RoPE NTK mix factor (default: %.1f)\n", params.rope_ntk_factor);
+    fprintf(stderr, "  --rope-extrapolation-factor N\n");
+    fprintf(stderr, "                        RoPE extrapolation mix factor (default: %.1f)\n", params.rope_extrapolation_factor);
     fprintf(stderr, "  --ignore-eos          ignore end of stream token and continue generating (implies --logit-bias 2-inf)\n");
     fprintf(stderr, "  --no-penalize-nl      do not penalize newline token\n");
     fprintf(stderr, "  --memory-f32          use f32 instead of f16 for memory key+value (default: disabled)\n");
@@ -596,6 +611,8 @@ struct llama_context_params llama_context_params_from_gpt_params(const gpt_param
     lparams.embedding    = params.embedding;
     lparams.rope_freq_base  = params.rope_freq_base;
     lparams.rope_freq_scale = params.rope_freq_scale;
+    lparams.rope_ntk_factor = params.rope_ntk_factor;
+    lparams.rope_extrapolation_factor = params.rope_extrapolation_factor;
 
     return lparams;
 }

examples/common.h

@@ -35,6 +35,8 @@ struct gpt_params {
     int32_t n_probs                         = 0;   // if greater than 0, output the probabilities of top n_probs tokens.
     float   rope_freq_base                  = 10000.0f; // RoPE base frequency
     float   rope_freq_scale                 = 1.0f;     // RoPE frequency scaling factor
+    float   rope_ntk_factor                 = 1.0f;     // RoPE NTK mix factor
+    float   rope_extrapolation_factor       = 1.0f;     // RoPE extrapolation mix factor
 
     // sampling parameters
     std::unordered_map<llama_token, float> logit_bias; // logit bias for specific tokens

examples/server/server.cpp

@@ -610,6 +610,9 @@ static void server_print_usage(const char *argv0, const gpt_params &params,
     fprintf(stderr, "  -c N, --ctx-size N    size of the prompt context (default: %d)\n", params.n_ctx);
     fprintf(stderr, "  --rope-freq-base N    RoPE base frequency (default: %.1f)\n", params.rope_freq_base);
     fprintf(stderr, "  --rope-freq-scale N   RoPE frequency scaling factor (default: %g)\n", params.rope_freq_scale);
+    fprintf(stderr, "  --rope-ntk-factor N   RoPE NTK mix factor (default: %.1f)\n", params.rope_ntk_factor);
+    fprintf(stderr, "  --rope-extrapolation-factor N\n");
+    fprintf(stderr, "                        RoPE extrapolation mix factor (default: %.1f)\n", params.rope_extrapolation_factor);
     fprintf(stderr, "  -b N, --batch-size N  batch size for prompt processing (default: %d)\n", params.n_batch);
     fprintf(stderr, "  --memory-f32          use f32 instead of f16 for memory key+value (default: disabled)\n");
     fprintf(stderr, "                        not recommended: doubles context memory required and no measurable increase in quality\n");
@@ -740,6 +743,22 @@ static void server_params_parse(int argc, char **argv, server_params &sparams,
             }
             params.rope_freq_scale = std::stof(argv[i]);
         }
+        else if (arg == "--rope-ntk-factor")
+        {
+            if (++i >= argc) {
+                invalid_param = true;
+                break;
+            }
+            params.rope_ntk_factor = std::stof(argv[i]);
+        }
+        else if (arg == "--rope-extrapolation-factor")
+        {
+            if (++i >= argc) {
+                invalid_param = true;
+                break;
+            }
+            params.rope_extrapolation_factor = std::stof(argv[i]);
+        }
         else if (arg == "--memory-f32" || arg == "--memory_f32")
         {
             params.memory_f16 = false;