Add rope_scaling support for LLama3.1 (#356)

Add support for rope scaling and FusedRoPE in LLama3.1
HabanaAI · Oct 3, 2024 · f848d27 · f848d27
1 parent da03d8b
commit f848d27
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 8 deletions.
diff --git a/requirements-hpu.txt b/requirements-hpu.txt
@@ -8,4 +8,4 @@ pandas
 tabulate
 setuptools>=61
 setuptools-scm>=8
-vllm-hpu-extension @ git+https://github.com/HabanaAI/vllm-hpu-extension.git@940fdb7
+vllm-hpu-extension @ git+https://github.com/HabanaAI/vllm-hpu-extension.git@bb56d3b
diff --git a/vllm/model_executor/layers/rotary_embedding.py b/vllm/model_executor/layers/rotary_embedding.py
@@ -31,7 +31,8 @@
 from vllm.platforms import current_platform
 
 if current_platform.is_hpu():
-    from vllm_hpu_extension.rotary_embed import HpuRotaryEmbedding
+    from vllm_hpu_extension.rotary_embed import (HpuLlama3RotaryEmbedding,
+                                                 HpuRotaryEmbedding)
 
 
 def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
@@ -943,12 +944,23 @@ def get_rope(
             high_freq_factor = rope_scaling["high_freq_factor"]
             original_max_position = rope_scaling[
                 "original_max_position_embeddings"]
-            rotary_emb = Llama3RotaryEmbedding(head_size, rotary_dim,
-                                               max_position, base,
-                                               is_neox_style, dtype,
-                                               scaling_factor, low_freq_factor,
-                                               high_freq_factor,
-                                               original_max_position)
+            if current_platform.is_hpu():
+                rotary_emb = HpuLlama3RotaryEmbedding(
+                    head_size,
+                    rotary_dim,
+                    max_position,
+                    base,
+                    is_neox_style,
+                    scaling_factor,
+                    low_freq_factor,
+                    high_freq_factor,
+                    original_max_position,
+                    RoPEFallback=Llama3RotaryEmbedding)
+            else:
+                rotary_emb = Llama3RotaryEmbedding(
+                    head_size, rotary_dim, max_position, base, is_neox_style,
+                    dtype, scaling_factor, low_freq_factor, high_freq_factor,
+                    original_max_position)
         elif scaling_type == "linear":
             rotary_emb = LinearScalingRotaryEmbedding(head_size, rotary_dim,
                                                       max_position, base,