Feat: Add rope scaling

axolotl-ai-cloud · Aug 5, 2023 · 517a59a · 517a59a
1 parent c93655c
commit 517a59a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -472,6 +472,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
@@ -223,6 +223,7 @@ def load_model(
                 load_in_4bit=cfg.load_in_4bit and cfg.adapter is not None,
                 torch_dtype=torch_dtype,
                 device_map="auto" if cfg.world_size == 1 else cfg.device_map,
+                rope_scaling=cfg.rope_scaling,
                 **model_kwargs,
             )
         # elif model_type == "GPTNeoXForCausalLM" and cfg.flash_attention: