change rotary base (#31)

- feat: rotary base as a property (c1200891411b6198ca6448cfebf5123d15bf2c31)
- Merge branch 'main' into pr/31 (c2ead96805f8278295d48fda36eba1d96ed3bffb)


Co-authored-by: Jack Min Ong <[email protected]>

configuration_xlm_roberta.py

@@ -20,6 +20,7 @@ class XLMRobertaFlashConfig(PretrainedConfig):
             bos_token_id=0,
             eos_token_id=2,
             position_embedding_type="absolute",
+            rotary_emb_base=10000.0,
             use_cache=True,
             classifier_dropout=None,
             lora_adaptations=None,
@@ -52,6 +53,7 @@ class XLMRobertaFlashConfig(PretrainedConfig):
         self.initializer_range = initializer_range
         self.layer_norm_eps = layer_norm_eps
         self.position_embedding_type = position_embedding_type
+        self.rotary_emb_base = rotary_emb_base
         self.use_cache = use_cache
         self.classifier_dropout = classifier_dropout
         self.load_trained_adapters = load_trained_adapters

modeling_lora.py

@@ -262,6 +262,14 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         self.main_params_trainable = config.lora_main_params_trainable
 
 
+    @property
+    def rotary_emb_base(self):
+        return self.roberta.rotary_emb_base
+
+    @rotary_emb_base.setter
+    def rotary_emb_base(self, base):
+        self.roberta.rotary_emb_base = base
+
     @property
     def main_params_trainable(self):
         return self._main_params_trainable

modeling_xlm_roberta.py

@@ -93,7 +93,7 @@ def create_mixer_cls(config, cross_attn=False, return_residual=False):
         rotary_kwargs["rotary_emb_dim"] = getattr(
             config, "rotary_emb_dim", config.hidden_size / config.num_attention_heads
         )
-        rotary_kwargs["rotary_emb_base"] = getattr(config, "rotary_emb_base", 10000.0)
+        rotary_kwargs["rotary_emb_base"] = config.rotary_emb_base
         rotary_kwargs["rotary_emb_scale_base"] = getattr(
             config, "rotary_emb_scale_base", None
         )
@@ -450,6 +450,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
 
         self.apply(partial(_init_weights, initializer_range=config.initializer_range))
         self.tokenizer = AutoTokenizer.from_pretrained(self.name_or_path, trust_remote_code=True)
+        self._rotary_emb_base = config.rotary_emb_base
 
     @torch.inference_mode()
     def encode(
@@ -599,7 +600,6 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         self.train(is_training)
         return all_embeddings
 
-
     def truncate_embeddings(self, embeddings, truncate_dim):
         if not self.config.matryoshka_dimensions:
             logger.warning(
@@ -622,12 +622,23 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             input_mask_expanded.sum(1), min=1e-9
         )
 
-
     def cls_pooling(
         self, token_embeddings: torch.Tensor, attention_mask: torch.Tensor
     ):
         return token_embeddings[:,0]
 
+    @property
+    def rotary_emb_base(self):
+        return self._rotary_emb_base
+
+    @rotary_emb_base.setter
+    def rotary_emb_base(self, base):
+        if not isinstance(base, (int, float)):
+            raise TypeError("Base must be an integer or float")
+        logger.info(f'Changing RoPE base value to {base}')
+        for layer in self.encoder.layers:
+            layer.mixer.rotary_emb.base = base
+        self._rotary_emb_base = base
 
     def forward(
         self,

rotary.py

@@ -443,7 +443,7 @@ class RotaryEmbedding(torch.nn.Module):
         """
         super().__init__()
         self.dim = dim
-        self.base = float(base)
+        self._base = float(base)
         self.pos_idx_in_fp32 = pos_idx_in_fp32
         # Generate and save the inverse frequency buffer (non trainable)
         inv_freq = self._compute_inv_freq(device)
@@ -463,6 +463,17 @@ class RotaryEmbedding(torch.nn.Module):
         self._cos_k_cached = None
         self._sin_k_cached = None
 
+    @property
+    def base(self):
+        return self._base
+
+    @base.setter
+    def base(self, new_base):
+        if new_base > 0:
+            self._base = float(new_base)
+        else:
+            raise ValueError("Rotary base value must be positive")
+
     def _compute_inv_freq(self, device=None):
         return 1.0 / (
             self.base