support flashmask (#8670)

* support flashmask * support flashmask
PaddlePaddle · Jun 27, 2024 · 1182cd5 · 1182cd5
1 parent dae64cc
commit 1182cd5
Showing 1 changed file with 17 additions and 5 deletions.
diff --git a/paddlenlp/transformers/llama/modeling_pp.py b/paddlenlp/transformers/llama/modeling_pp.py
@@ -208,11 +208,23 @@ def forward(self, args):
             alibi = position_ids
             position_ids = attn_mask_startend_row_indices
             attn_mask_startend_row_indices = None
-        elif not self.config.alibi and position_ids is None and attn_mask_startend_row_indices is not None:
-            # hidden_states, attention_mask, position_ids
-            position_ids = attn_mask_startend_row_indices
-            attn_mask_startend_row_indices = None
-            alibi = None
+        elif not self.config.alibi:
+            if get_env_device() in ["gpu"]:
+                if attention_mask is not None and attention_mask.dtype == paddle.int32:
+                    attention_mask, attn_mask_startend_row_indices, position_ids = (
+                        None,
+                        attention_mask,
+                        attn_mask_startend_row_indices,
+                    )
+                elif attention_mask is not None and attention_mask.dtype == paddle.int64:
+                    attention_mask, attn_mask_startend_row_indices, position_ids = None, None, attention_mask
+                elif (
+                    attn_mask_startend_row_indices is not None and attn_mask_startend_row_indices.dtype == paddle.int64
+                ):
+                    attn_mask_startend_row_indices, position_ids = None, attn_mask_startend_row_indices
+            elif position_ids is None and attn_mask_startend_row_indices is not None:
+                position_ids = attn_mask_startend_row_indices
+                attn_mask_startend_row_indices = None
 
         has_gradient = not hidden_states.stop_gradient
         if self.enable_recompute and self.config.recompute_granularity == "full" and has_gradient: