Command line argument to override rope_theta from config.json

0cc4m · Sep 1, 2023 · a544085 · a544085
1 parent b7f3fac
commit a544085
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/model_init.py b/model_init.py
@@ -15,6 +15,7 @@ def add_args(parser):
     parser.add_argument("-l", "--length", type = int, help = "Maximum sequence length", default = 2048)
     parser.add_argument("-cpe", "--compress_pos_emb", type = float, help = "Compression factor for positional embeddings", default = 1.0)
     parser.add_argument("-a", "--alpha", type = float, help = "alpha for context size extension via embedding extension", default = 1.0)
+    parser.add_argument("-theta", "--theta", type = float, help = "theta (base) for RoPE embeddings")
 
     parser.add_argument("-gpfix", "--gpu_peer_fix", action = "store_true", help = "Prevent direct copies of data between GPUs")
 
@@ -140,6 +141,9 @@ def make_config(args):
     config.silu_no_half2 = args.silu_no_half2
     config.concurrent_streams = args.concurrent_streams
 
+    if args.theta:
+        config.rotary_embedding_base = args.theta
+
     return config