Feat: Add rope scaling (axolotl-ai-cloud#343)

* Feat: Add rope scaling * fix: move rope config
mkeoliya · Aug 12, 2023 · b500e0c · b500e0c
1 parent 1631ead
commit b500e0c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -474,6 +474,10 @@ landmark_attention:
 # xpos RoPE see https://github.com/kaiokendev/cutoff-len-is-context-len/blob/main/util/xpos_rope_llama_monkey_patch.py
 # llama only
 xpos_rope:
+# RoPE Scaling https://github.com/huggingface/transformers/pull/24653
+rope_scaling:
+  type: # linear | dynamic
+  factor: # float
 
 # resume from a specific checkpoint dir
 resume_from_checkpoint:

diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
@@ -219,7 +219,9 @@ def load_model(
         elif cfg.is_llama_derived_model and not cfg.trust_remote_code:
             from transformers import LlamaForCausalLM
 
-            config = LlamaConfig.from_pretrained(base_model_config)
+            config = LlamaConfig.from_pretrained(
+                base_model_config, rope_scaling=cfg.rope_scaling
+            )
             model = LlamaForCausalLM.from_pretrained(
                 base_model,
                 config=config,