fix: pass rope_theta argument when initializing LlamaLikeBlock for mo…

…dels like qwen2, mistral, etc. (#568) Co-authored-by: Shuai Xie <shuaixie@jd.com>
casper-hansen · Aug 4, 2024 · 202b967 · 202b967
1 parent 2e2635b
commit 202b967
Show file tree

Hide file tree

Showing 6 changed files with 6 additions and 0 deletions.
diff --git a/awq/models/aquila.py b/awq/models/aquila.py
@@ -127,6 +127,7 @@ def fuse_transformer(self):
                     norm_2=norm_2,
                     dev=device,
                     max_seq_len=self.model.config.max_seq_len,
+                    rope_theta=self.model.config.rope_theta,
                 )
             )
 

diff --git a/awq/models/llava.py b/awq/models/llava.py
@@ -133,6 +133,7 @@ def fuse_transformer(self):
                     norm_2=norm_2,
                     dev=device,
                     max_seq_len=max_seq_len,
+                    rope_theta=self.model.config.rope_theta,
                 )
             )
 

diff --git a/awq/models/llava_next.py b/awq/models/llava_next.py
@@ -130,6 +130,7 @@ def fuse_transformer(self):
                     norm_2=norm_2,
                     dev=device,
                     max_seq_len=max_seq_len,
+                    rope_theta=self.model.config.rope_theta,
                 )
             )
 

diff --git a/awq/models/mistral.py b/awq/models/mistral.py
@@ -127,6 +127,7 @@ def fuse_transformer(self):
                     norm_2=norm_2,
                     dev=device,
                     max_seq_len=self.model.config.max_seq_len,
+                    rope_theta=self.model.config.rope_theta,
                 )
             )
 

diff --git a/awq/models/qwen2.py b/awq/models/qwen2.py
@@ -125,6 +125,7 @@ def fuse_transformer(self):
                     norm_2=norm_2,
                     dev=device,
                     max_seq_len=self.model.config.max_seq_len,
+                    rope_theta=self.model.config.rope_theta,
                 )
             )
 

diff --git a/awq/models/starcoder2.py b/awq/models/starcoder2.py
@@ -125,6 +125,7 @@ def fuse_transformer(self):
                     norm_2=norm_2,
                     dev=device,
                     max_seq_len=self.model.config.max_seq_len,
+                    rope_theta=self.model.config.rope_theta,
                 )
             )