Set RotaryEmbedding base frequency from config

nemo/collections/nlp/models/language_modeling/megatron/gpt_model.py

@@ -167,6 +167,7 @@ def __init__(
         ub_tp_comm_overlap=False,
         use_flash_attention=False,
         seq_len_interpolation_factor=None,
+        rotary_base=10000,
     ):
         super(GPTModel, self).__init__(config=config, share_token_embeddings=share_embeddings_and_output_weights)
 
@@ -248,6 +249,7 @@ def __init__(
             ub_tp_comm_overlap=ub_tp_comm_overlap,
             use_flash_attention=use_flash_attention,
             seq_len_interpolation_factor=seq_len_interpolation_factor,
+            rotary_base=rotary_base,
         )
 
         if self.share_embeddings_and_output_weights:

nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py

@@ -384,6 +384,7 @@ def model_provider_func(self, pre_process, post_process):
                 use_flash_attention=self.cfg.get('use_flash_attention', False),
                 megatron_legacy=self.cfg.get('megatron_legacy', False),
                 seq_len_interpolation_factor=self.cfg.get('seq_len_interpolation_factor', None),
+                rotary_base=self.cfg.get('rotary_base', 10000),
             )
         return model
 

nemo/collections/nlp/modules/common/megatron/language_model.py

@@ -126,6 +126,7 @@ def get_language_model(
     ub_tp_comm_overlap=False,
     use_flash_attention=False,
     seq_len_interpolation_factor=None,
+    rotary_base=10000,
 ):
     """Build language model and return along with the key to save."""
 
@@ -202,6 +203,7 @@ def get_language_model(
         ub_tp_comm_overlap=ub_tp_comm_overlap,
         use_flash_attention=use_flash_attention,
         seq_len_interpolation_factor=seq_len_interpolation_factor,
+        rotary_base=rotary_base,
     )
     # key used for checkpoints.
     language_model_key = 'language_model'
@@ -502,6 +504,7 @@ def __init__(
         ub_tp_comm_overlap=False,
         use_flash_attention=False,
         seq_len_interpolation_factor=None,
+        rotary_base=10000,
     ):
         super(TransformerLanguageModel, self).__init__(
             config=config, share_token_embeddings=share_embeddings_and_output_weights
@@ -557,6 +560,7 @@ def __init__(
                 rotary_dim,
                 seq_len_interpolation_factor=seq_len_interpolation_factor,
                 pretrained_max_position_embeddings=max_position_embeddings,
+                rotary_base=rotary_base,
             )
 
         elif position_embedding_type == 'alibi':

nemo/collections/nlp/modules/common/megatron/position_embedding/rotary_position_embedding.py

@@ -26,19 +26,25 @@ class RotaryEmbedding(nn.Module):
     """
 
     def __init__(
-        self, dim: int, seq_len_interpolation_factor: int = None, pretrained_max_position_embeddings: int = None
+        self,
+        dim: int,
+        seq_len_interpolation_factor: int = None,
+        rotary_base: int = 10000,
+        pretrained_max_position_embeddings: int = None,
     ):
         """
         Args:
 
             dim (int): rotary embedding dimension
             seq_len_interpolation_factor (int): if not None, discrete positions will be interpolated
-            by this factor via the trick in https://arxiv.org/abs/2306.15595.
+                by this factor via the trick in https://arxiv.org/abs/2306.15595.
+            rotary_base (int): rotary_base for the positional frequency (default: 10000)
             pretrained_max_position_embeddings (int): pre-trained max_position_embeddings before position interpolation.
         """
         super().__init__()
         self.seq_len_interpolation_factor = seq_len_interpolation_factor
-        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.rotary_base = rotary_base
+        inv_freq = 1.0 / (self.rotary_base ** (torch.arange(0, dim, 2).float() / dim))
         self.register_buffer('inv_freq', inv_freq)
         self.pretrained_max_position_embeddings = pretrained_max_position_embeddings