..

mosaicml · Dec 6, 2024 · 77115c5 · 77115c5
1 parent ad601e4
commit 77115c5
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/llmfoundry/models/layers/flex_attn_utils.py b/llmfoundry/models/layers/flex_attn_utils.py
@@ -110,8 +110,7 @@ def _mask_mod_fn(
             )
         sequence_id = sequence_id_info['sequence_id']
         # Check if the query and key belong to the same sequence and the query token is not a padding token.
-        return (sequence_id[b, q_idx]
-                == sequence_id[b, kv_idx]) & (sequence_id[b, kv_idx] != -1)
+        return (sequence_id[b, q_idx] == sequence_id[b, kv_idx])
 
     def __init__(self) -> None:
         super().__init__(mod_type='mask')