Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix SDPA dispatch & make SDPA CI compatible with torch<2.1.1 #27940

Merged
merged 1 commit into from
Dec 11, 2023
+10 −8
Merged

Fix SDPA dispatch & make SDPA CI compatible with torch<2.1.1 #27940

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 8 additions & 7 deletions src/transformers/modeling_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1244,6 +1244,7 @@ def _autoset_attn_implementation(
# Here we use config._attn_implementation_internal to check whether the attention implementation was explicitely set by the user.
# The property `PretrainedConfig._attn_implementation` is never `None`, for backward compatibility (always fall back on "eager").
# The `hasattr` here is used as some Transformers tests for some reason do not call PretrainedConfig __init__ (e.g. test_no_super_init_config_and_model)
requested_attn_implementation = None
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this be "default" instead?

Copy link
Contributor Author

@fxmarty fxmarty Dec 11, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No, the idea here is to check whether the user passed attn_implementation="eager", attn_implementation="sdpa" or attn_implementation="sdpa" explicitly when loading the model from from_pretrained or from_config.

In case attn_implementation is explicitly set, we hard error if a dependency is missing (torch>=2.1.1, model does not support SDPA), otherwise we smoothly fall back on eager.

if hasattr(config, "_attn_implementation_internal") and config._attn_implementation_internal is not None:
if config._attn_implementation != "flash_attention_2" and use_flash_attention_2:
raise ValueError(
Expand All @@ -1260,9 +1261,7 @@ def _autoset_attn_implementation(
raise ValueError(message + ".")

# If a config is passed with a preset attn_implementation, we skip the automatic dispatch and use the user-provided config, with hard checks that the requested attention implementation is available.
hard_check_only = True
else:
hard_check_only = False
requested_attn_implementation = config._attn_implementation_internal

if use_flash_attention_2:
logger.warning_once(
Expand All @@ -1275,13 +1274,15 @@ def _autoset_attn_implementation(
config,
torch_dtype=torch_dtype,
device_map=device_map,
hard_check_only=hard_check_only,
hard_check_only=False,
check_device_map=check_device_map,
)
elif cls._supports_sdpa or config._attn_implementation == "sdpa":
elif requested_attn_implementation in [None, "sdpa"]:
# use_flash_attention_2 takes priority over SDPA, hence SDPA treated in this elif.
config = cls._check_and_enable_sdpa(config, hard_check_only=hard_check_only)
elif not hard_check_only:
config = cls._check_and_enable_sdpa(
config, hard_check_only=False if requested_attn_implementation is None else True
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks better thanks

)
else:
config._attn_implementation = "eager"

return config
Expand Down
3 changes: 2 additions & 1 deletion tests/test_modeling_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@
is_flax_available,
is_tf_available,
is_torch_fx_available,
is_torch_sdpa_available,
)
from transformers.utils.generic import ModelOutput

Expand Down Expand Up @@ -778,7 +779,7 @@ def _create_and_check_torchscript(self, config, inputs_dict):
configs_no_init.torchscript = True
for model_class in self.all_model_classes:
for attn_implementation in ["eager", "sdpa"]:
if attn_implementation == "sdpa" and not model_class._supports_sdpa:
if attn_implementation == "sdpa" and (not model_class._supports_sdpa or not is_torch_sdpa_available()):
continue

configs_no_init._attn_implementation = attn_implementation
Expand Down
Loading