Disable the FA backend for SDPA on AMD GPUs (huggingface#30850)

* disable fa * disable fa * update warning * update warning
sywangyi · May 16, 2024 · 0753134 · 0753134
1 parent 9d889f8
commit 0753134
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
@@ -1479,6 +1479,16 @@ def _autoset_attn_implementation(
                 config,
                 hard_check_only=False if requested_attn_implementation is None else True,
             )
+
+            if (
+                torch.version.hip is not None
+                and config._attn_implementation == "sdpa"
+                and torch.cuda.device_count() > 1
+            ):
+                logger.warning_once(
+                    "Using the `SDPA` attention implementation on multi-gpu setup with ROCM may lead to performance issues due to the FA backend. Disabling it to use alternative backends."
+                )
+                torch.backends.cuda.enable_flash_sdp(False)
         else:
             config._attn_implementation = "eager"