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

handle args to drop dropout

Browse Source
This commit is contained in:
Wing Lian
2025-05-18 13:11:56 -07:00
parent 0735454782
commit d6f64a3684
2 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -644,10 +644,16 @@ class ModelLoader:
def flash_attn_func_v3_wrapper(*args, **kwargs): def flash_attn_func_v3_wrapper(*args, **kwargs):
kwargs.pop("dropout_p", None) kwargs.pop("dropout_p", None)
if "softmax_scale" in kwargs and len(args) >= 4:
# if softmax_scale is provided, then the 3rd position is dropout_p that we need to drop
args = (*args[:3],) + args[4:]
return flash_attn_func_v3(*args, **kwargs)[0] return flash_attn_func_v3(*args, **kwargs)[0]
def flash_attn_varlen_func_v3_wrapper(*args, **kwargs): def flash_attn_varlen_func_v3_wrapper(*args, **kwargs):
kwargs.pop("dropout_p", None) kwargs.pop("dropout_p", None)
if "softmax_scale" in kwargs and len(args) >= 4:
# if softmax_scale is provided, then the 3rd position is dropout_p that we need to drop
args = (*args[:3],) + args[4:]
return flash_attn_varlen_func_v3(*args, **kwargs)[0] return flash_attn_varlen_func_v3(*args, **kwargs)[0]
transformers.modeling_flash_attention_utils.flash_attn_func = ( transformers.modeling_flash_attention_utils.flash_attn_func = (

2
src/axolotl/utils/schemas/config.py
View File

@@ -233,7 +233,7 @@ class AxolotlInputConfig(
flash_attn_fuse_qkv: bool | None = None flash_attn_fuse_qkv: bool | None = None
flash_attn_fuse_mlp: bool | None = None flash_attn_fuse_mlp: bool | None = None
flash_optimum: bool | None = None flash_optimum: bool | None = None
use_flash_attention_3: Literal["auto"] | bool | None = "auto" use_flash_attention_3: Literal["auto"] | bool | None = None
eager_attention: bool | None = None eager_attention: bool | None = None