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

llama sdpa patching WIP

Browse Source
This commit is contained in:
Sunny Liu
2025-01-22 20:15:23 -05:00
parent d1be6e228d
commit cee310dcfa
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -2,10 +2,6 @@
Patched LlamaAttention to use torch.nn.functional.scaled_dot_product_attention Patched LlamaAttention to use torch.nn.functional.scaled_dot_product_attention
""" """
from typing import Optional
import torch
from transformers.utils import is_torch_bf16_gpu_available
from axolotl.monkeypatch.utils import ( from axolotl.monkeypatch.utils import (
mask_2d_to_4d, mask_2d_to_4d,
@@ -15,11 +11,15 @@ from axolotl.monkeypatch.utils import (
def hijack_llama_prepare_4d_mask(): def hijack_llama_prepare_4d_mask():
from typing import Optional
import torch
from transformers import modeling_attn_mask_utils from transformers import modeling_attn_mask_utils
from transformers.models.llama.modeling_llama import LlamaModel from transformers.models.llama.modeling_llama import LlamaModel
from transformers.models.llama.modeling_llama.LlamaModel import ( from transformers.models.llama.modeling_llama.LlamaModel import (
_prepare_4d_causal_attention_mask_with_cache_position, _prepare_4d_causal_attention_mask_with_cache_position,
) )
from transformers.utils import is_torch_bf16_gpu_available
# modeling_llama._prepare_4d_causal_attention_mask_for_sdpa = ( # pylint: disable=protected-access # modeling_llama._prepare_4d_causal_attention_mask_for_sdpa = ( # pylint: disable=protected-access
# patched_prepare_4d_causal_attention_mask_for_sdpa # patched_prepare_4d_causal_attention_mask_for_sdpa