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

most model types now support flash attention 2 regardless of multipack support (#1854)

Browse Source
This commit is contained in:
Wing Lian
2024-08-22 16:39:23 -04:00
committed by GitHub
parent b33dc07a77
commit fefa95e350
2 changed files with 5 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/axolotl/monkeypatch/multipack.py
View File

@@ -17,6 +17,7 @@ SUPPORTED_MULTIPACK_MODEL_TYPES = [
"qwen2_moe", "qwen2_moe",
"falcon", "falcon",
"phi", "phi",
"phi3",
"gemma", "gemma",
"gemma2", "gemma2",
"gemmoe", "gemmoe",

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

@@ -591,16 +591,10 @@ def load_model(
"flash_attention_2" "flash_attention_2"
) )
else: else:
if model_config.model_type in SUPPORTED_MULTIPACK_MODEL_TYPES: model_kwargs["attn_implementation"] = "flash_attention_2"
model_kwargs["attn_implementation"] = "flash_attention_2" model_config._attn_implementation = ( # pylint: disable=protected-access
model_config._attn_implementation = ( # pylint: disable=protected-access "flash_attention_2"
"flash_attention_2" )
)
else:
model_kwargs["attn_implementation"] = "eager"
model_config._attn_implementation = ( # pylint: disable=protected-access
"eager"
)
elif cfg.sdp_attention: elif cfg.sdp_attention:
model_kwargs["attn_implementation"] = "sdpa" model_kwargs["attn_implementation"] = "sdpa"
model_config._attn_implementation = "sdpa" # pylint: disable=protected-access model_config._attn_implementation = "sdpa" # pylint: disable=protected-access