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

move flash-attn monkey patch alongside the others

Browse Source
This commit is contained in:
Aman Karmani
2023-08-03 17:20:49 +00:00
parent 248bf90f89
commit 312a9fad07
2 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
src/axolotl/flash_attn.py → src/axolotl/monkeypatch/llama_attn_hijack_flash.py
View File

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

@@ -92,7 +92,9 @@ def load_model(
if cfg.is_llama_derived_model and cfg.flash_attention: if cfg.is_llama_derived_model and cfg.flash_attention:
if cfg.device not in ["mps", "cpu"] and not cfg.inference: if cfg.device not in ["mps", "cpu"] and not cfg.inference:
from axolotl.flash_attn import replace_llama_attn_with_flash_attn from axolotl.monkeypatch.llama_attn_hijack_flash import (
replace_llama_attn_with_flash_attn,
)
LOG.info("patching with flash attention") LOG.info("patching with flash attention")
replace_llama_attn_with_flash_attn() replace_llama_attn_with_flash_attn()