[shardformer]fix flash attention, when mask is casual, just don't unpad it (#5084)

* fix flash attn

* fix

fix

colossalai/shardformer/modeling/opt.py

@@ -581,7 +581,8 @@ def get_opt_flash_attention_forward():
                     f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
                 )
             flash_attention_mask = ~(attention_mask[:, :, -1].squeeze(1).to(torch.bool)).contiguous()
-            attn_mask_type = AttnMaskType.paddedcausal
+            if not torch.all(flash_attention_mask):
+                attn_mask_type = AttnMaskType.paddedcausal
 
         attention = ColoAttention(
             embed_dim=self.embed_dim, num_heads=self.num_heads, dropout=self.dropout, scale=self.scaling