423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-08-01 11:12:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

tiny fix

Browse Source 
Former-commit-id: 3f24337a8a995b145b1e8075bc23878eaa363844
This commit is contained in:
hiyouga 2024-06-11 01:04:16 +08:00
parent d6632fefc9
commit e3baa5aa08
1 changed files with 1 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/llamafactory/model/model_utils/longlora.py
View File

@ -182,11 +182,9 @@ def llama_flash_attention_2_forward(
query_states, key_states, value_states = shift(query_states), shift(key_states), shift(value_states) query_states, key_states, value_states = shift(query_states), shift(key_states), shift(value_states)
if attention_mask is not None: if attention_mask is not None:
attention_mask = attention_mask[:, :groupsz].repeat(num_groups, 1) attention_mask = attention_mask[:, :groupsz].repeat(num_groups, 1)
else:
groupsz = q_len
attn_output: torch.Tensor = self._flash_attention_forward( attn_output: torch.Tensor = self._flash_attention_forward(
query_states, key_states, value_states, attention_mask, groupsz, dropout=dropout_rate query_states, key_states, value_states, attention_mask, query_states.size(1), dropout=dropout_rate
) )
if getattr(self.config, "group_size_ratio", None) and self.training: # shift back if getattr(self.config, "group_size_ratio", None) and self.training: # shift back