Upload modeling_roberta.py

modeling_roberta.py

@@ -0,0 +1,197 @@
+from typing import Unpack
+import torch
+from transformers import (
+    RobertaModel,
+    Cache,
+    EncoderDecoderCache,
+    DynamicCache,
+    DataCollatorWithFlattening,
+    RobertaForMaskedLM,
+    RobertaForSequenceClassification,
+    RobertaForTokenClassification,
+    RobertaForQuestionAnswering,
+    RobertaForMultipleChoice,
+    RobertaForCausalLM
+)
+from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions
+from transformers.utils import TransformersKwargs
+
+
+def _unpad_input(input_ids: torch.Tensor, attention_mask: torch.Tensor):
+    collator = DataCollatorWithFlattening(return_flash_attn_kwargs=True)
+    features = collator([{"input_ids": i[a.bool()].tolist()} for i, a in zip(input_ids, attention_mask)])
+    return features
+
+
+def _pad_output(inputs: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int,) -> torch.Tensor:
+    if inputs.dim() == 3:
+        inputs = inputs.squeeze()
+    if inputs.dim() == 1:
+        output = torch.zeros(batch * seqlen, dtype=inputs.dtype, device=inputs.device)
+        output[indices] = inputs
+        padded_inputs = output.view(batch, seqlen)
+    else:
+        _, *rest = inputs.shape
+        output = torch.zeros(batch * seqlen, *rest, dtype=inputs.dtype, device=inputs.device)
+        output[indices] = inputs
+        padded_inputs = output.view(batch, seqlen, *rest)
+    return padded_inputs
+
+
+class UnpadRobertaModel(RobertaModel):
+    _no_split_modules = ["RobertaEmbeddings", "RobertaLayer"]
+
+    def __init__(self, config, add_pooling_layer=True):
+        super().__init__(config, add_pooling_layer=add_pooling_layer)
+
+    def forward(
+            self,
+            input_ids: torch.Tensor | None = None,
+            attention_mask: torch.Tensor | None = None,
+            token_type_ids: torch.Tensor | None = None,
+            position_ids: torch.Tensor | None = None,
+            inputs_embeds: torch.Tensor | None = None,
+            encoder_hidden_states: torch.Tensor | None = None,
+            encoder_attention_mask: torch.Tensor | None = None,
+            past_key_values: Cache | None = None,
+            use_cache: bool | None = None,
+            cache_position: torch.Tensor | None = None,
+            **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor] | BaseModelOutputWithPoolingAndCrossAttentions:
+        if self.config.is_decoder:
+            use_cache = use_cache if use_cache is not None else self.config.use_cache
+        else:
+            use_cache = False
+
+        if use_cache and past_key_values is None:
+            past_key_values = (
+                EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+                if encoder_hidden_states is not None or self.config.is_encoder_decoder
+                else DynamicCache(config=self.config)
+            )
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if input_ids is not None:
+            device = input_ids.device
+            seq_length = input_ids.shape[1]
+            batch_size = input_ids.size(0)
+        else:
+            device = inputs_embeds.device
+            seq_length = inputs_embeds.shape[1]
+            batch_size = inputs_embeds.size(0)
+
+        past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+        if cache_position is None:
+            cache_position = torch.arange(past_key_values_length, past_key_values_length + seq_length, device=device)
+
+        indices = None
+        if self.config._attn_implementation.startswith("flash_attention"):
+            if input_ids is None or attention_mask is None:
+                raise ValueError("Unpadding requires both input_ids and attention_mask")
+            with torch.no_grad():
+                indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+                features = _unpad_input(input_ids, attention_mask)
+                input_ids = features["input_ids"].to(device=device)
+                # roberta requires shifting position_ids by 2
+                position_ids = (features["position_ids"] + 2).to(device=device)
+                attention_mask = None
+                kwargs["cu_seq_lens_k"] = features["cu_seq_lens_k"].to(device=device)
+                kwargs["cu_seq_lens_q"] = features["cu_seq_lens_q"].to(device=device)
+                kwargs["max_length_k"] = features["max_length_k"]
+                kwargs["max_length_q"] = features["max_length_q"]
+
+        embedding_output = self.embeddings(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values_length=past_key_values_length,
+        )
+
+        attention_mask, encoder_attention_mask = self._create_attention_masks(
+            attention_mask=attention_mask,
+            encoder_attention_mask=encoder_attention_mask,
+            embedding_output=embedding_output,
+            encoder_hidden_states=encoder_hidden_states,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+        )
+
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask=attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_ids=position_ids,
+            **kwargs,
+        )
+
+        sequence_output = encoder_outputs.last_hidden_state
+        if self.config._attn_implementation.startswith("flash_attention"):
+            sequence_output = _pad_output(
+                inputs=sequence_output, indices=indices, batch=batch_size, seqlen=seq_length
+            )
+
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            past_key_values=encoder_outputs.past_key_values,
+        )
+
+
+class UnpadRobertaForCausalLM(RobertaForCausalLM):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.roberta = UnpadRobertaModel(config, add_pooling_layer=False)
+        self.post_init()
+
+
+class UnpadRobertaForMaskedLM(RobertaForMaskedLM):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.roberta = UnpadRobertaModel(config, add_pooling_layer=False)
+        self.post_init()
+
+
+class UnpadRobertaForSequenceClassification(RobertaForSequenceClassification):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.roberta = UnpadRobertaModel(config, add_pooling_layer=False)
+        self.post_init()
+
+
+class UnpadRobertaForTokenClassification(RobertaForTokenClassification):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.roberta = UnpadRobertaModel(config, add_pooling_layer=False)
+        self.post_init()
+
+
+class UnpadRobertaForMultipleChoice(RobertaForMultipleChoice):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.roberta = UnpadRobertaModel(config)
+        self.post_init()
+
+
+class UnpadRobertaForQuestionAnswering(RobertaForQuestionAnswering):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.roberta = UnpadRobertaModel(config, add_pooling_layer=False)
+        self.post_init()
+
+
+def enable_roberta_unpadding():
+    RobertaModel.forward = UnpadRobertaModel.forward