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

llama sdpa patching WIP - static class function import

Browse Source
This commit is contained in:
Sunny Liu
2025-01-22 20:45:44 -05:00
parent d7b133dc1f
commit 06f83a54a5
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/axolotl/monkeypatch/llama_patch_multipack.py
View File

@@ -28,11 +28,11 @@ def hijack_llama_prepare_4d_mask():
# )
def llama_patched_prepare_4d_causal_attention_mask_with_cache_position(
attention_mask: Optional[torch.Tensor], *args, **kwargs
self, attention_mask: Optional[torch.Tensor], *args, **kwargs
):
dtype = torch.bfloat16 if is_torch_bf16_gpu_available() else torch.float32
return LlamaModel._prepare_4d_causal_attention_mask_with_cache_position(
mask_2d_to_4d(attention_mask, dtype=dtype), *args, **kwargs
self, mask_2d_to_4d(attention_mask, dtype=dtype), *args, **kwargs
)
LlamaModel._prepare_4d_causal_attention_mask_with_cache_position = ( # pylint: disable=protected-access