Support YaRN models (#1264)

Signed-off-by: Antoni Baum <antoni.baum@protonmail.com> Co-authored-by: Viktor Ferenczi <viktor@ferenczi.eu> Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
2023-11-03 14:12:48 -07:00
parent 555bdcc5a3
commit 9f669a9a7c
5 changed files with 128 additions and 8 deletions
--- a/vllm/config.py
+++ b/vllm/config.py
@@ -390,6 +390,9 @@ def _get_and_verify_max_len(
    if rope_scaling is not None:
        assert "factor" in rope_scaling
        scaling_factor = rope_scaling["factor"]
+        if rope_scaling["type"] == "yarn":
+            derived_max_model_len = rope_scaling[
+                "original_max_position_embeddings"]
        derived_max_model_len *= scaling_factor

    if max_model_len is None: