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 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.
+        # 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.
         # 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,72 +533,104 @@ 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 = torch.float16
+                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}."
+            )
+
             query_states = query_states.to(target_dtype)
             key_states = key_states.to(target_dtype)
             value_states = value_states.to(target_dtype)
 
-        # 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
+        # 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)
 
-        # 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)
+        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,
+        )
 
-            if causal:
-                attention_mask = attention_mask & causal_mask
-            else:
-                attention_mask = attention_mask
+        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
 
-        # Compute the softmax scale
-        softmax_scale = self.head_dim**-0.5
+        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)
 
-        # Compute the attention scores
+        # Contains at least one padding token in the sequence
         if attention_mask is not None:
-            # Unpad the input
-            (
-                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, q_len)
+            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
+            )
 
             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,
@@ -606,8 +638,8 @@ class QuietFlashAttention2(QuietAttention):
                     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,
+                    max_seqlen_q=max_seqlen_in_batch_q,
+                    max_seqlen_k=max_seqlen_in_batch_k,
                     dropout_p=dropout,
                     softmax_scale=softmax_scale,
                     causal=causal,
@@ -619,8 +651,8 @@ class QuietFlashAttention2(QuietAttention):
                     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,
+                    max_seqlen_q=max_seqlen_in_batch_q,
+                    max_seqlen_k=max_seqlen_in_batch_k,
                     dropout_p=dropout,
                     softmax_scale=softmax_scale,
                     causal=causal,
@@ -663,7 +695,8 @@ class QuietFlashAttention2(QuietAttention):
         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
@@ -691,6 +724,8 @@ class QuietFlashAttention2(QuietAttention):
             (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):
     """
@@ -768,7 +803,7 @@ class QuietSdpaAttention(QuietAttention):
             attn_mask=attention_mask.to(query_states.device) if attention_mask is not None else None,
             dropout_p=self.attention_dropout if self.training else 0.0,
             # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
-            causal=self.is_causal and attention_mask is None and q_len > 1,
+            is_causal=self.is_causal and attention_mask is None and q_len > 1,
         )
 
         attn_output = attn_output.transpose(1, 2).contiguous()
@@ -1814,7 +1849,6 @@ class QuietForCausalLM(QuietPreTrainedModel):
                             dim=-1
                         )
                         
-                        
                         # print((new_rm_tokens > self.vocab_size - 1).any().item())
                         new_rm_tokens = torch.clamp(new_rm_tokens, 0, self.vocab_size - 1)