attempt - mask padding

2025-02-02 20:18:49 -05:00
parent d3ea379a23
commit b0871c8d3b
1 changed files with 2 additions and 1 deletions
--- a/src/axolotl/monkeypatch/flex_attn.py
+++ b/src/axolotl/monkeypatch/flex_attn.py
@@ -131,7 +131,8 @@ def packed_block_causal_mask(
        """
        causal_mask = q_idx >= kv_idx
        document_mask = document_ids[b, q_idx] == document_ids[b, kv_idx]
-        return causal_mask & document_mask & (q_idx < totalseqlens[b])
+        finite_mask = q_idx < totalseqlens[b]
+        return causal_mask & document_mask & finite_mask

    return create_block_causal_mask_flex(
        mask_mod,