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

Fix patching via import instead of hijacking

Browse Source
This commit is contained in:
NanoCode012
2023-06-09 14:27:24 +09:00
parent 55b8542de8
commit e44c9e0b3e
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/axolotl/utils/models.py
View File

@@ -20,7 +20,9 @@ from transformers import ( # noqa: F401
) )
try: try:
from transformers import LlamaForCausalLM from transformers import ( # pylint: disable=unused-import # noqa: F401
LlamaForCausalLM,
)
except ImportError: except ImportError:
logging.warning( logging.warning(
"This version of transformers does not support Llama. Consider upgrading." "This version of transformers does not support Llama. Consider upgrading."
@@ -115,15 +117,15 @@ def load_model(
logging.info("patching with sdp attention") logging.info("patching with sdp attention")
hijack_llama_sdp_attention() hijack_llama_sdp_attention()
elif cfg.is_llama_derived_model and cfg.landmark_attention: elif cfg.is_llama_derived_model and cfg.landmark_attention:
from axolotl.monkeypatch.llama_landmark_attn import ( from axolotl.monkeypatch.llama_landmark_attn import ( # pylint: disable=redefined-outer-name # noqa: F811
MEM_TOKEN, MEM_TOKEN,
hijack_llama_landmark_attn, LlamaForCausalLM,
) )
logging.info("patching with landmark attention") logging.info("patching with landmark attention")
hijack_llama_landmark_attn()
tokenizer.add_special_tokens({"mem_token": MEM_TOKEN}) # TODO: Check if this would overwrite previous additional_special_tokens
tokenizer.add_special_tokens({"additional_special_tokens": [MEM_TOKEN]})
if cfg.bf16: if cfg.bf16:
torch_dtype = torch.bfloat16 torch_dtype = torch.bfloat16