Support SDPA, fix embeddings, output attention probs.

modeling.py

@@ -374,7 +374,7 @@ class NewEmbeddings(nn.Module):
         if position_ids is None:
             if seq_length > self.position_ids.size(0):
                 self.register_buffer(
-                    "position_ids", torch.arange(seq_length), persistent=False
+                    "position_ids", torch.arange(seq_length, device=embeddings.device), persistent=False
                 )
             if unpad_inputs:
                 # [1, cumsum_seq_len]
@@ -397,16 +397,19 @@ class NewEmbeddings(nn.Module):
         if self.type_vocab_size > 0:
             if token_type_ids is None:
                 token_type_ids = position_ids.mul(0)
-            elif unpad_inputs:
-                token_type_ids = token_type_ids[attention_mask_bool].unsqueeze(0)
+            else:
+                if self.type_vocab_size < 2:
+                    token_type_ids.mul_(0)
+                if unpad_inputs:
+                    token_type_ids = token_type_ids[attention_mask_bool].unsqueeze(0)
 
             token_type_embeddings = self.token_type_embeddings(token_type_ids)
-            embeddings += token_type_embeddings
+            embeddings = embeddings + token_type_embeddings
 
         # BERT position
         if self.position_embedding_type == "absolute":
             position_embeddings = self.position_embeddings(position_ids)
-            embeddings += position_embeddings
+            embeddings = embeddings + position_embeddings
 
         embeddings = self.LayerNorm(embeddings)
         embeddings = self.dropout(embeddings)
@@ -449,19 +452,17 @@ class NewAttention(nn.Module):
         self.memory_efficient_attention = None if xops is None else xops.memory_efficient_attention
         if self.use_memory_efficient_attention:
             assert self.memory_efficient_attention is not None, 'please install xformers'
-        if self.config.unpad_inputs:
-            assert self.config.use_memory_efficient_attention, 'unpad only with xformers'
 
     def forward(
         self,
         hidden_states: torch.Tensor,
         attention_bias: torch.FloatTensor,
         rope_embeds: Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] = None,
+        padding_inputs: Optional[Tuple] = None,  # indices, batch, seqlen
         attention_scale: Optional[torch.FloatTensor] = None,
         head_mask: Optional[torch.FloatTensor] = None,
         output_attentions: Optional[bool] = False,
         qkv_inputs: Optional[Tuple] = None,  # For RetroMAE
-        padding_inputs: Optional[Tuple] = None,  # indices, batch, seqlen
     ) -> Tuple[torch.Tensor, ...]:
         shape_hd = (self.num_attention_heads, self.attention_head_size)
         # qkv
@@ -504,7 +505,11 @@ class NewAttention(nn.Module):
                 p=self.dropout.p
             )
         else:
-            context_layer = self._attention(query_states, key_states, value_states, attention_bias, head_mask)
+            if output_attentions and isinstance(self, NewSdpaAttention):
+                raise RuntimeError("SDPA do not output attentions")
+            context_layer, attention_probs = self._attention(
+                query_states, key_states, value_states, attention_bias, head_mask
+            )
 
         if padding_inputs is not None:
             context_layer = unpad_input(context_layer, indices=padding_inputs[0])
@@ -542,7 +547,8 @@ class NewAttention(nn.Module):
 
         # This is actually dropping out entire tokens to attend to, which might
         # seem a bit unusual, but is taken from the original Transformer paper.
-        attention_probs = self.dropout(attention_probs)
+        if self.dropout.p > 0:
+            attention_probs = self.dropout(attention_probs)
 
         # Mask heads if we want to
         if head_mask is not None:
@@ -551,7 +557,7 @@ class NewAttention(nn.Module):
         context_layer = torch.matmul(attention_probs, value_states)
 
         context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
-        return context_layer
+        return context_layer, attention_probs
 
 
 class NewSdpaAttention(NewAttention):
@@ -562,11 +568,11 @@ class NewSdpaAttention(NewAttention):
     """
     def __init__(self, config: NewConfig, **kwargs):
         super().__init__(config, **kwargs)
-        torch.backends.cuda.enable_mem_efficient_sdp(False)
-        logger.warning(
-            "Disable memory efficient attention kernel for `NewSdpaAttention`, you can set "
-            "`use_memory_efficient_attention=True` if it expected to use."
-        )
+        # torch.backends.cuda.enable_mem_efficient_sdp(False)
+        # logger.warning(
+        #     "Disable memory efficient attention kernel for `NewSdpaAttention`, you can set "
+        #     "`use_memory_efficient_attention=True` if it expected to use."
+        # )
 
     def _attention(self, query_states, key_states, value_states, attention_bias, head_mask):
         attn_output = torch.nn.functional.scaled_dot_product_attention(
@@ -577,12 +583,12 @@ class NewSdpaAttention(NewAttention):
             dropout_p=self.dropout.p if self.training else 0.0,
         )
         attn_output = attn_output.permute(0, 2, 1, 3).contiguous()
-        return attn_output
+        return attn_output, None
 
 
 NEW_ATTENTION_CLASSES = {
     "eager": NewAttention,
-    # "flash_attention_2": ,  # TODO: xformers will dispatch to flash_attn
+    # "flash_attention_2": ,  # TODO
     "sdpa": NewSdpaAttention,
 }
 
@@ -625,8 +631,12 @@ class NewLayer(nn.Module):
         super().__init__()
         if attn_implementation is None:
             attn_implementation = config._attn_implementation
-        if attn_implementation != 'eager':
-            use_memory_efficient_attention = False
+        if use_memory_efficient_attention is None:
+            use_memory_efficient_attention = config.use_memory_efficient_attention
+        if use_memory_efficient_attention:
+            if attn_implementation != 'eager':
+                logger.warning_once(f"Override {attn_implementation=} to 'eager' as {use_memory_efficient_attention=}")
+                attn_implementation = 'eager'  # Since it will be SDPA by default for torch>=2.1.1
         self.attention = NEW_ATTENTION_CLASSES[attn_implementation](
             config, pack_qkv=pack_qkv, use_memory_efficient_attention=use_memory_efficient_attention
         )
@@ -646,12 +656,12 @@ class NewLayer(nn.Module):
         hidden_states: torch.Tensor,
         attention_bias: torch.FloatTensor,
         rope_embeds: Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] = None,
+        padding_inputs: Optional[Tuple] = None,  # indices, batch, seqlen
         attention_scale: Optional[torch.FloatTensor] = None,
         subset_indices: Optional[torch.LongTensor] = None,
         head_mask: Optional[torch.FloatTensor] = None,
         output_attentions: Optional[bool] = False,
         qkv_inputs: Optional[Tuple] = None,  # For RetroMAE
-        padding_inputs: Optional[Tuple] = None,
     ) -> Tuple[torch.Tensor, ...]:
         # Multi head self attention
         residual = hidden_states if qkv_inputs is None else qkv_inputs[0]
@@ -659,11 +669,11 @@ class NewLayer(nn.Module):
             hidden_states,
             attention_bias,
             rope_embeds,
+            padding_inputs,
             attention_scale,
             head_mask,
             output_attentions=output_attentions,
             qkv_inputs=qkv_inputs,
-            padding_inputs=padding_inputs,
         )
         hidden_states = attention_outputs[0]
         if self.hidden_dropout is not None:
@@ -701,6 +711,7 @@ class NewEncoder(nn.Module):
         hidden_states: torch.Tensor,
         attention_bias: Optional[torch.FloatTensor] = None,
         rope_embeds: Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] = None,
+        padding_inputs: Optional[Tuple] = None,  # indices, batch, seqlen
         attention_scale: Optional[torch.FloatTensor] = None,
         subset_indices: Optional[torch.LongTensor] = None,
         head_mask: Optional[torch.FloatTensor] = None,
@@ -728,6 +739,7 @@ class NewEncoder(nn.Module):
                     hidden_states,
                     attention_bias,
                     rope_embeds,
+                    padding_inputs,
                     attention_scale,
                     layer_subset_indices,
                     layer_head_mask,
@@ -737,6 +749,7 @@ class NewEncoder(nn.Module):
                     hidden_states,
                     attention_bias,
                     rope_embeds,
+                    padding_inputs,
                     attention_scale,
                     layer_subset_indices,
                     layer_head_mask,
@@ -792,6 +805,7 @@ class NewPreTrainedModel(PreTrainedModel):
     config_class = NewConfig
     base_model_prefix = "new"
     supports_gradient_checkpointing = True
+    _supports_sdpa = True
 
     def _init_weights(self, module):
         """Initialize the weights"""
@@ -894,9 +908,7 @@ class NewModel(NewPreTrainedModel):
         )
 
         batch_size, seq_length = input_shape
-
-        if unpad_inputs:
-            assert self.config.use_memory_efficient_attention
+        if unpad_inputs and self.config.use_memory_efficient_attention:
             attention_bias = xops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(length)
         else:
             # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
@@ -906,20 +918,29 @@ class NewModel(NewPreTrainedModel):
                 # Invalid shape for attention bias: torch.Size([48, 1, 1, 512]) (expected (48, 12, 512, 512))
                 attention_bias = attention_bias.expand(-1, self.config.num_attention_heads, seq_length, -1)
 
+        padding_inputs = None
+        if unpad_inputs and (output_padded or not self.config.use_memory_efficient_attention):
+            indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+            if not self.config.use_memory_efficient_attention:
+                padding_inputs = (indices, *input_shape)
+
+        attention_scale = None
         if self.config.logn_attention_scale:
-            # attention scale log_512(input_len)
-            attention_scale = attention_mask.sum(1).log() / torch.tensor(self.config.max_position_embeddings).log()
-            # inference-time logn scale need clip 1
-            if self.config.logn_attention_clip1:
-                attention_scale.clip_(1)
-            attention_scale = attention_scale[:, None, None, None]
-        else:
-            attention_scale = None
+            logger.warning_once("TODO: logn_attention_scale")
+        #     # attention scale log_512(input_len)
+        #     attention_scale = attention_mask.sum(1).log() / torch.tensor(self.config.max_position_embeddings).log()
+        #     # inference-time logn scale need clip 1
+        #     if self.config.logn_attention_clip1:
+        #         attention_scale.clip_(1)
+        #     attention_scale = attention_scale[:, None, None, None]
+        # else:
+        #     attention_scale = None
 
         encoder_outputs = self.encoder(
             embedding_output,
             attention_bias=attention_bias,
             rope_embeds=rope_embeds,
+            padding_inputs=padding_inputs,
             attention_scale=attention_scale,
             subset_indices=subset_indices,
             head_mask=head_mask,
@@ -929,7 +950,6 @@ class NewModel(NewPreTrainedModel):
         )
         sequence_output = encoder_outputs[0]
         if unpad_inputs and output_padded:
-            indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
             sequence_output = pad_input(
                 sequence_output.squeeze(), indices, batch_size, seq_length
             )