alibi (#19)

- feat: support alibi (b4903887bd152b045811fb78a5edc369a8db7cb5)


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

embedding.py

@@ -50,7 +50,7 @@ class XLMRobertaEmbeddings(nn.Module):
         embeddings = self.word_embeddings(input_ids)
         if self.max_position_embeddings > 0:
             if position_ids is None:
-                position_ids =create_position_ids_from_input_ids(input_ids, padding_idx=self.word_embeddings.padding_idx).to(input_ids.device)
+                position_ids = create_position_ids_from_input_ids(input_ids, padding_idx=self.word_embeddings.padding_idx).to(input_ids.device)
                 # position_ids = torch.arange(seqlen, dtype=torch.long, device=input_ids.device)
             position_embeddings = self.position_embeddings(position_ids)
             embeddings = embeddings + position_embeddings

modeling_xlm_roberta.py

@@ -109,6 +109,7 @@ def create_mixer_cls(config, cross_attn=False, return_residual=False):
         fused_bias_fc=fused_bias_fc,
         use_flash_attn=use_flash_attn,
         return_residual=return_residual,
+        use_alibi=config.position_embedding_type == 'alibi',
         **rotary_kwargs,
     )
     return mixer_cls
@@ -429,7 +430,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         self.embeddings = XLMRobertaEmbeddings(
             config.hidden_size,
             config.vocab_size,
-            config.max_position_embeddings,
+            config.max_position_embeddings if config.position_embedding_type == 'absolute' else -1,
             config.type_vocab_size,
             padding_idx=config.pad_token_id,
         )