Change RotaryEmbedding to GPTNeoXRotaryEmbedding (togethercomputer#168)

updates GPTNeoX RotaryEmbedding for transformeres==4.31.0
laiqun · Aug 17, 2023 · 6372379 · 6372379
1 parent 759fc49
commit 6372379
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/training/modules/hf_gptneox_modules.py b/training/modules/hf_gptneox_modules.py
@@ -13,7 +13,7 @@
 from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXLayer as _GPTNeoXBlock
 from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXModel as _GPTNeoXModel
 from transformers.models.gpt_neox.configuration_gpt_neox import GPTNeoXConfig as GPTConfig
-from transformers.models.gpt_neox.modeling_gpt_neox import RotaryEmbedding
+from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXRotaryEmbedding
 
 
 try:
@@ -111,7 +111,7 @@ def __init__(self, config):
             ),
         )
         self.register_buffer("masked_bias", torch.tensor(-1e9))
-        self.rotary_emb = RotaryEmbedding(
+        self.rotary_emb = GPTNeoXRotaryEmbedding(
             self.rotary_ndims, config.max_position_embeddings, base=config.rotary_emb_base
         )
         self.register_buffer(