Update modeling_motif.py

modeling_motif.py

@@ -545,9 +545,7 @@ class MotifFlashAttention2(MotifAttention):
 
         bsz = query_states.shape[0]
 
-        return map(
-            lambda x: x.float32(), 
-            _flash_attention_forward(query_states.bfloat16(),
+        return _flash_attention_forward(query_states.bfloat16(),
                                 key_states.bfloat16(),
                                 value_states.bfloat16(),
                                 attention_mask,
@@ -557,7 +555,6 @@ class MotifFlashAttention2(MotifAttention):
                                 sliding_window=sliding_window,
                                 is_causal=self.is_causal,
                                 use_top_left_mask=self._flash_attn_uses_top_left_mask)
-        )
 
     def forward(
             self,
@@ -642,7 +639,7 @@ class MotifFlashAttention2(MotifAttention):
                             self._compute_attention(q2, k2, v2, attention_mask, q_len, position_ids, dropout_rate, sliding_window, self.batch_num)
 
         
-        attn1, attn2 = torch.cat([attn11, attn12], dim=-1), torch.cat([attn21, attn22], dim=-1)
+        attn1, attn2 = torch.cat([attn11, attn12], dim=-1).float(), torch.cat([attn21, attn22], dim=-1).float()
 
         lambda_q1 = self.lambda_q1.unsqueeze(0).expand([bsz, self.lambda_q1.shape[0]])  # bsz, num_head
         lambda_q2 = self.lambda_q2.unsqueeze(0).expand([bsz, self.lambda_q2.shape[0]])  # bsz, num_head
@@ -661,10 +658,10 @@ class MotifFlashAttention2(MotifAttention):
             raise ValueError(f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
                              f" {attn_output.size()}")
 
-        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).bfloat16()
         attn_output = self.o_proj(attn_output) * self.o_proj_alpha
 
-        return attn_output, None, past_key_value
+        return attn_output.float(), None, past_key_value
 
 
 # @log_timing