Update modeling_quiet.py

modeling_quiet.py

@@ -432,7 +432,7 @@ class QuietFlashAttention2(QuietAttention):
         super().__init__(*args, **kwargs)
 
         # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
-        # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+        # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
         # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
         self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
 
@@ -533,171 +533,123 @@ class QuietFlashAttention2(QuietAttention):
             if torch.is_autocast_enabled():
                 target_dtype = torch.get_autocast_gpu_dtype()
             # Handle the case where the model is quantized
-            elif hasattr(self.config, "_pre_quantization_dtype"):
-                target_dtype = self.config._pre_quantization_dtype
             else:
-                target_dtype = self.q_proj.weight.dtype
-
-            logger.warning_once(
-                f"The input hidden states seems to be silently casted in float32, this might be related to"
-                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
-                f" {target_dtype}."
-            )
-
+                target_dtype = torch.float16
             query_states = query_states.to(target_dtype)
             key_states = key_states.to(target_dtype)
             value_states = value_states.to(target_dtype)
 
-        # Reashape to the expected shape for Flash Attention
-        query_states = query_states.transpose(1, 2)
-        key_states = key_states.transpose(1, 2)
-        value_states = value_states.transpose(1, 2)
-
-        attn_output = self._flash_attention_forward(
-            query_states,
-            key_states,
-            value_states,
-            attention_mask,
-            q_len,
-            dropout=dropout_rate,
-            use_sliding_windows=use_sliding_windows,
-        )
-
-        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
-        attn_output = self.o_proj(attn_output)
-
-        if not output_attentions:
-            attn_weights = None
-
-        return attn_output, attn_weights, past_key_value
-
-def _flash_attention_forward(
-    self,
-    query_states,
-    key_states,
-    value_states,
-    attention_mask,
-    query_length,
-    dropout=0.0,
-    softmax_scale=None,
-    use_sliding_windows=False,
-):
-    """
-    Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
-    first unpad the input, then computes the attention scores and pad the final attention scores.
-    Args:
-        query_states (`torch.Tensor`):
-            Input query states to be passed to Flash Attention API
-        key_states (`torch.Tensor`):
-            Input key states to be passed to Flash Attention API
-        value_states (`torch.Tensor`):
-            Input value states to be passed to Flash Attention API
-        attention_mask (`torch.Tensor`):
-            The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
-            position of padding tokens and 1 for the position of non-padding tokens.
-        dropout (`int`, *optional*):
-            Attention dropout
-        softmax_scale (`float`, *optional*):
-            The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
-        use_sliding_windows (`bool`, *optional*):
-            Whether to activate sliding window attention.
-    """
-    if not self._flash_attn_uses_top_left_mask:
-        causal = self.is_causal
-    else:
-        # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
-        causal = self.is_causal and query_length != 1
-        
-    # Ensure attention_mask has the correct shape and values
-    if attention_mask is not None:
-        if attention_mask.dim() == 4:
-            # Convert 4D attention mask to 2D
-            attention_mask = attention_mask.squeeze(1).squeeze(1)
-        elif attention_mask.dim() != 2:
-            raise ValueError(
-                f"Invalid attention mask dimension: {attention_mask.dim()}. Expected 2D or 4D mask."
-            )
-        
-        # Ensure attention_mask has values of 0 and 1
-        attention_mask = attention_mask.to(torch.bool).to(torch.int32)
-
-    # Contains at least one padding token in the sequence
-    if attention_mask is not None:
-        batch_size = query_states.shape[0]
-        query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
-            query_states, key_states, value_states, attention_mask, query_length
-        )
+        # Compute the causal mask
+        causal = self.config.causal
+        if causal:
+            if self._flash_attn_uses_top_left_mask:
+                # Compute the causal mask
+                causal_mask = torch.tril(torch.ones((q_len, kv_seq_len), dtype=torch.bool, device=query_states.device))
+                # Invert the mask
+                causal_mask = ~causal_mask
+            else:
+                causal_mask = torch.triu(
+                    torch.ones((q_len, kv_seq_len), dtype=torch.bool, device=query_states.device), diagonal=1
+                )
+        else:
+            causal_mask = None
 
-        cu_seqlens_q, cu_seqlens_k = cu_seq_lens
-        max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+        # Compute the attention mask
+        if attention_mask is not None:
+            if attention_mask.dim() == 2:
+                attention_mask = attention_mask[:, None, :]
+            attention_mask = attention_mask.to(torch.bool)
 
-        # Create the cu_seqlens_q and cu_seqlens_k tensors
-        q_max_s, k_max_s = query_states.shape[1], key_states.shape[1]
-        qkv_max_s = max(q_max_s, k_max_s)
+            if causal:
+                attention_mask = attention_mask & causal_mask
+            else:
+                attention_mask = attention_mask
 
-        q_seqlens = torch.full((batch_size,), q_max_s, dtype=torch.int32, device=query_states.device)
-        k_seqlens = torch.full((batch_size,), k_max_s, dtype=torch.int32, device=key_states.device)
+        # Compute the softmax scale
+        softmax_scale = self.head_dim**-0.5
 
-        # Adjust the attention mask to match the sequence lengths
+        # Compute the attention scores
         if attention_mask is not None:
-            q_seqlens = attention_mask.sum(dim=1).int()
-            k_seqlens = attention_mask.sum(dim=1).int()
-
-        # Convert seqlens to cumulative sequence lengths
-        cu_seqlens_q = torch.cat([torch.zeros(1, dtype=torch.int32, device=q_seqlens.device), q_seqlens.cumsum(dim=0)])
-        cu_seqlens_k = torch.cat([torch.zeros(1, dtype=torch.int32, device=k_seqlens.device), k_seqlens.cumsum(dim=0)])
-
-        if not use_sliding_windows:
-            attn_output_unpad = flash_attn_varlen_func(
-                query_states,
-                key_states,
-                value_states,
-                cu_seqlens_q=cu_seqlens_q,
-                cu_seqlens_k=cu_seqlens_k,
-                max_seqlen_q=qkv_max_s,
-                max_seqlen_k=qkv_max_s,
-                dropout_p=dropout,
-                softmax_scale=softmax_scale,
-                causal=causal,
-            )
-        else:
-            attn_output_unpad = flash_attn_varlen_func(
+            # Unpad the input
+            (
                 query_states,
                 key_states,
                 value_states,
-                cu_seqlens_q=cu_seqlens_q,
-                cu_seqlens_k=cu_seqlens_k,
-                max_seqlen_q=qkv_max_s,
-                max_seqlen_k=qkv_max_s,
-                dropout_p=dropout,
-                softmax_scale=softmax_scale,
-                causal=causal,
-                window_size=(self.config.sliding_window, self.config.sliding_window),
-            )
+                indices_q,
+                cu_seq_lens,
+                max_seq_lens,
+            ) = self._upad_input(query_states, key_states, value_states, attention_mask, q_len)
 
-        attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
-    else:
-        if not use_sliding_windows:
-            attn_output = flash_attn_func(
-                query_states,
-                key_states,
-                value_states,
-                dropout,
-                softmax_scale=softmax_scale,
-                causal=causal,
-            )
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+            # Create the cu_seqlens_q and cu_seqlens_k tensors
+            q_max_s, k_max_s = query_states.shape[1], key_states.shape[1]
+            qkv_max_s = max(q_max_s, k_max_s)
+
+            q_seqlens = torch.full((batch_size,), q_max_s, dtype=torch.int32, device=query_states.device)
+            k_seqlens = torch.full((batch_size,), k_max_s, dtype=torch.int32, device=key_states.device)
+
+            # Adjust the attention mask to match the sequence lengths
+            if attention_mask is not None:
+                q_seqlens = attention_mask.sum(dim=1).int()
+                k_seqlens = attention_mask.sum(dim=1).int()
+
+            # Convert seqlens to cumulative sequence lengths
+            cu_seqlens_q = torch.cat([torch.zeros(1, dtype=torch.int32, device=q_seqlens.device), q_seqlens.cumsum(dim=0)])
+            cu_seqlens_k = torch.cat([torch.zeros(1, dtype=torch.int32, device=k_seqlens.device), k_seqlens.cumsum(dim=0)])
+
+            if not use_sliding_windows:
+                attn_output_unpad = flash_attn_varlen_func(
+                    query_states,
+                    key_states,
+                    value_states,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_q=qkv_max_s,
+                    max_seqlen_k=qkv_max_s,
+                    dropout_p=dropout,
+                    softmax_scale=softmax_scale,
+                    causal=causal,
+                )
+            else:
+                attn_output_unpad = flash_attn_varlen_func(
+                    query_states,
+                    key_states,
+                    value_states,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_q=qkv_max_s,
+                    max_seqlen_k=qkv_max_s,
+                    dropout_p=dropout,
+                    softmax_scale=softmax_scale,
+                    causal=causal,
+                    window_size=(self.config.sliding_window, self.config.sliding_window),
+                )
+
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
         else:
-            attn_output = flash_attn_func(
-                query_states,
-                key_states,
-                value_states,
-                dropout,
-                softmax_scale=softmax_scale,
-                causal=causal,
-                window_size=(self.config.sliding_window, self.config.sliding_window),
-            )
+            if not use_sliding_windows:
+                attn_output = flash_attn_func(
+                    query_states,
+                    key_states,
+                    value_states,
+                    dropout,
+                    softmax_scale=softmax_scale,
+                    causal=causal,
+                )
+            else:
+                attn_output = flash_attn_func(
+                    query_states,
+                    key_states,
+                    value_states,
+                    dropout,
+                    softmax_scale=softmax_scale,
+                    causal=causal,
+                    window_size=(self.config.sliding_window, self.config.sliding_window),
+                )
 
-    return attn_output
+        return attn_output
 
     def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
         batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape
@@ -711,8 +663,7 @@ def _flash_attention_forward(
         indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
 
         key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
-        value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
-
+        value_layer= index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
         if query_length == kv_seq_len:
             query_layer = index_first_axis(
                 query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
@@ -740,8 +691,6 @@ def _flash_attention_forward(
             (cu_seqlens_q, cu_seqlens_k),
             (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
         )
-
-
 # Copied from transformers.models.llama.modeling_llama.LlamaSdpaAttention with Llama->Quiet
 class QuietSdpaAttention(QuietAttention):
     """