tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

attempt - mask padding

Browse Source
This commit is contained in:
Sunny Liu
2025-02-02 20:18:49 -05:00
parent d3ea379a23
commit b0871c8d3b
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/axolotl/monkeypatch/flex_attn.py
View File

@@ -131,7 +131,8 @@ def packed_block_causal_mask(
""" """
causal_mask = q_idx >= kv_idx causal_mask = q_idx >= kv_idx
document_mask = document_ids[b, q_idx] == document_ids[b, 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( return create_block_causal_mask_flex(
mask_mod, mask_mod,