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

more test

Browse Source
This commit is contained in:
bursteratom
2025-02-02 00:43:41 -05:00
parent fa7355404c
commit 10de67e8ea
2 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -50,8 +50,10 @@ def create_block_causal_mask(
residue_len = max_seq_len - torch.sum(seq_lens[sample_idx]) residue_len = max_seq_len - torch.sum(seq_lens[sample_idx])
block_attn_masks.append( block_attn_masks.append(
torch.zeros( torch.tril(
residue_len, residue_len, dtype=torch.bool, device=seq_lens[0][0].device torch.zeros(
residue_len, residue_len, dtype=torch.bool, device=seq_lens[sample_idx].device
)
) )
) )

2
src/axolotl/utils/collators/batching.py
View File

@@ -182,7 +182,7 @@ class FlexBatchSamplerDataCollatorForSeq2Seq(DataCollatorForSeq2Seq):
collated_seq_lens, max_seq_len = get_seqlens_from_pos_ids(out["position_ids"]) collated_seq_lens, max_seq_len = get_seqlens_from_pos_ids(out["position_ids"])
# out["attention_mask"] = packed_block_causal_mask(collated_seq_lens) # out["attention_mask"] = packed_block_causal_mask(collated_seq_lens)
out["attention_mask"] = create_block_causal_mask(collated_seq_lens, max_seq_len) out["attention_mask"] = create_block_causal_mask(collated_seq_lens, max_seq_len)
raise ValueError(f"{out['attention_mask'].shape}") # raise ValueError(f"{out['attention_mask'].shape}")
return out return out