fix(transformers): align SDPA with Transformers' logic (#1167)

* fix sdpa to align with Transformers' logic

* speed up execution with `ops.speed_fusion_attention`

* add backprop test

* small fix

Author: hadipash

mindone/transformers/integrations/sdpa_attention.py

@@ -1,9 +1,13 @@
 """Adapted from https://github.com/huggingface/transformers/tree/main/src/transformers/integrations/sdpa_attention.py."""
-
+from math import sqrt
 from typing import Optional
 
 import mindspore as ms
-from mindspore import mint, nn
+from mindspore import mint, nn, ops
+
+from ..utils import logging
+
+logger = logging.get_logger(__name__)
 
 
 def repeat_kv(hidden_states: ms.Tensor, n_rep: int) -> ms.Tensor:
@@ -23,23 +27,53 @@ def sdpa_attention_forward(
     query: ms.Tensor,
     key: ms.Tensor,
     value: ms.Tensor,
-    attention_mask: Optional[ms.Tensor],
+    attention_mask: Optional[ms.Tensor] = None,
     dropout: float = 0.0,
     scaling: Optional[float] = None,
-    is_causal: Optional[bool] = None,  # to align with torch
+    is_causal: Optional[bool] = None,
     **kwargs,
 ) -> tuple[ms.Tensor, None]:
-    key_states = repeat_kv(key, module.num_key_value_groups)
-    value_states = repeat_kv(value, module.num_key_value_groups)
+    if kwargs.get("output_attentions", False) or kwargs.get("head_mask", None) is not None:
+        logger.warning_once(
+            "`sdpa` attention does not support `output_attentions=True` or `head_mask`."
+            " Please set your attention to `eager` if you want any of these features."
+        )
+
+    if hasattr(module, "num_key_value_groups"):
+        key = repeat_kv(key, module.num_key_value_groups)
+        value = repeat_kv(value, module.num_key_value_groups)
+
+    if attention_mask is not None and attention_mask.ndim == 4:
+        attention_mask = attention_mask[:, :, :, : key.shape[-2]]
+
+    # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
+    # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
+    # Note that it is important to check first for the shape, otherwise compile will fail with `argument 'is_causal' must be bool, not SymBool`
+    if is_causal is None:
+        # The last condition is for encoder (decoder) models which specify this by passing their own `is_causal` flag
+        # This is mainly due to those models having mixed implementations for encoder, decoder, and encoder-decoder attns
+        is_causal = query.shape[2] > 1 and attention_mask is None and getattr(module, "is_causal", True)
+
+    if is_causal:
+        if attention_mask is not None:
+            raise ValueError("Causal mode cannot be used with an explicit `attention_mask`")
+        attention_mask = mint.ones((query.shape[-2], key.shape[-2]), dtype=ms.bool_).tril(diagonal=0)
+
+    if attention_mask is not None:
+        attention_mask = mint.logical_not(attention_mask)  # in MindSpore, 0 indicates retain, 1 indicates discard
 
-    attn_weights = mint.matmul(query, key_states.transpose(2, 3)) * scaling
-    if attention_mask is not None and attention_mask.dim() == 4:
-        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
-        attn_weights = attn_weights + causal_mask
+    scaling = 1 / sqrt(query.shape[-1]) if scaling is None else scaling
 
-    attn_weights = mint.nn.functional.softmax(attn_weights, dim=-1, dtype=ms.float32).to(query.dtype)
-    attn_weights = mint.nn.functional.dropout(attn_weights, p=dropout, training=module.training)
-    attn_output = mint.matmul(attn_weights, value_states)
-    attn_output = attn_output.transpose(1, 2).contiguous()
+    attn_output = ops.speed_fusion_attention(
+        query,
+        key,
+        value,
+        head_num=query.shape[1],
+        input_layout="BNSD",
+        atten_mask=attention_mask,
+        scale=scaling,
+        keep_prob=1 - dropout,
+    )[0]
+    attn_output = mint.transpose(attn_output, 1, 2).contiguous()
 
-    return attn_output, attn_weights
+    return attn_output, None

tests/transformers_tests/integrations/__init__.py

@@ -0,0 +1,162 @@
+import numpy as np
+import pytest
+import torch
+from transformers.integrations.sdpa_attention import sdpa_attention_forward as sdpa_attention_forward_transformers
+
+from mindspore import grad, mint, set_seed, tensor
+
+from mindone.transformers.integrations.sdpa_attention import sdpa_attention_forward
+from tests.modeling_test_utils import MS_DTYPE_MAPPING, PT_DTYPE_MAPPING
+
+DTYPE_AND_THRESHOLDS = {"fp32": 1e-6, "fp16": 2e-3, "bf16": 2e-2}
+
+
+class MockAttentionModule:
+    def __init__(self, is_causal):
+        self.is_causal = is_causal
+        self.training = False
+
+
+@pytest.fixture(scope="module")
+def q_k_v_target_mask() -> dict[str, tuple]:
+    # B, H, S, D
+    set_seed(42)
+    q = np.random.uniform(size=(2, 8, 256, 32)).astype(np.float32)
+    k = np.random.uniform(size=(2, 8, 256, 32)).astype(np.float32)
+    v = np.random.uniform(size=(2, 8, 256, 32)).astype(np.float32)
+    target = np.random.uniform(size=(2, 256, 8, 32)).astype(np.float32)
+    attention_mask = np.random.randint(0, 2, (q.shape[0], 1, q.shape[2], k.shape[2]), dtype=bool)
+    return {
+        "ms": (tensor(q), tensor(k), tensor(v), tensor(target), tensor(attention_mask)),
+        "pt": (torch.tensor(q), torch.tensor(k), torch.tensor(v), torch.tensor(target), torch.tensor(attention_mask)),
+    }
+
+
+def cast_inputs(q, k, v, target, attention_mask, dtype):
+    return (
+        q.to(dtype),
+        k.to(dtype),
+        v.to(dtype),
+        target.to(dtype),
+        attention_mask,  # no casting needed for attention_mask
+    )
+
+
+@pytest.mark.parametrize("dtype", ["fp32", "fp16", "bf16"])
+@pytest.mark.parametrize("use_mask", [True, False], ids=["with_mask", "without_mask"])
+@pytest.mark.parametrize("jit", [False, True], ids=["eager", "jit"])
+def test_sdpa_attention_forward(q_k_v_target_mask, use_mask: bool, dtype: str, jit: bool):
+    if jit:
+        pytest.skip("`sdpa_attention_forward` can't be compiled with jit.")
+
+    module = MockAttentionModule(is_causal=False)
+
+    q, k, v, _, attn_mask = cast_inputs(*q_k_v_target_mask["ms"], dtype=MS_DTYPE_MAPPING[dtype])
+    q_pt, k_pt, v_pt, _, attn_mask_pt = cast_inputs(*q_k_v_target_mask["pt"], dtype=PT_DTYPE_MAPPING[dtype])
+
+    output = sdpa_attention_forward(module, q, k, v, attention_mask=attn_mask if use_mask else None, is_causal=False)[0]
+    output_pt = sdpa_attention_forward_transformers(
+        module, q_pt, k_pt, v_pt, attention_mask=attn_mask_pt if use_mask else None, is_causal=False
+    )[0]
+
+    assert output.shape == output_pt.shape, f"Shape mismatch: {output.shape} vs {output_pt.shape}"
+    assert not output.isnan().any(), "Output contains NaNs."
+    assert np.allclose(
+        output.numpy().astype(np.float32), output_pt.to(torch.float32).numpy(), atol=DTYPE_AND_THRESHOLDS[dtype]
+    )
+
+
+@pytest.mark.parametrize("dtype", ["fp32", "fp16", "bf16"])
+@pytest.mark.parametrize("use_mask", [True, False], ids=["with_mask", "without_mask"])
+@pytest.mark.parametrize("jit", [False, True], ids=["eager", "jit"])
+def test_sdpa_attention_backward(q_k_v_target_mask, use_mask: bool, dtype: str, jit: bool):
+    if jit:
+        pytest.skip("`sdpa_attention_forward` can't be compiled with jit.")
+
+    module = MockAttentionModule(is_causal=False)
+
+    # MindONE
+    q, k, v, target, attn_mask = cast_inputs(*q_k_v_target_mask["ms"], dtype=MS_DTYPE_MAPPING[dtype])
+
+    def _forward(q_, k_, v_, target_):
+        output = sdpa_attention_forward(
+            module, q_, k_, v_, attention_mask=attn_mask if use_mask else None, is_causal=False
+        )[0]
+        return mint.nn.functional.mse_loss(output, target_)
+
+    grad_out = grad(_forward, grad_position=(0, 1, 2))(q, k, v, target)
+    grad_out = mint.stack(grad_out, dim=0)
+
+    # Transformers
+    q_pt, k_pt, v_pt, target_pt, attn_mask_pt = cast_inputs(*q_k_v_target_mask["pt"], dtype=PT_DTYPE_MAPPING[dtype])
+    q_pt, k_pt, v_pt = q_pt.clone(), k_pt.clone(), v_pt.clone()
+    q_pt.requires_grad, k_pt.requires_grad, v_pt.requires_grad = (True,) * 3
+
+    output_pt = sdpa_attention_forward_transformers(
+        module, q_pt, k_pt, v_pt, attention_mask=attn_mask_pt if use_mask else None, is_causal=False
+    )[0]
+    loss = torch.nn.functional.mse_loss(output_pt, target_pt)
+    loss.backward()
+    grad_out_pt = torch.stack([q_pt.grad, k_pt.grad, v_pt.grad], dim=0)
+
+    assert grad_out.shape == grad_out_pt.shape, f"Shape mismatch: {grad_out.shape} vs {grad_out_pt.shape}"
+    assert not grad_out.isnan().any(), "Output contains NaNs."
+    assert np.allclose(
+        grad_out.numpy().astype(np.float32), grad_out_pt.to(torch.float32).numpy(), atol=DTYPE_AND_THRESHOLDS[dtype]
+    )
+
+
+@pytest.mark.parametrize("dtype", ["fp32", "fp16", "bf16"])
+@pytest.mark.parametrize("jit", [False, True], ids=["eager", "jit"])
+def test_sdpa_attention_causal_forward(q_k_v_target_mask, dtype: str, jit: bool):
+    if jit:
+        pytest.skip("`sdpa_attention_forward` can't be compiled with jit.")
+
+    module = MockAttentionModule(is_causal=True)
+
+    q, k, v, *_ = cast_inputs(*q_k_v_target_mask["ms"], dtype=MS_DTYPE_MAPPING[dtype])
+    q_pt, k_pt, v_pt, *_ = cast_inputs(*q_k_v_target_mask["pt"], dtype=PT_DTYPE_MAPPING[dtype])
+
+    output = sdpa_attention_forward(module, q, k, v, attention_mask=None)[0]
+    output_pt = sdpa_attention_forward_transformers(module, q_pt, k_pt, v_pt, attention_mask=None)[0]
+
+    assert output.shape == output_pt.shape, f"Shape mismatch: {output.shape} vs {output_pt.shape}"
+    assert not output.isnan().any(), "Output contains NaNs."
+    assert np.allclose(
+        output.numpy().astype(np.float32), output_pt.to(torch.float32).numpy(), atol=DTYPE_AND_THRESHOLDS[dtype]
+    )
+
+
+@pytest.mark.parametrize("dtype", ["fp32", "fp16", "bf16"])
+@pytest.mark.parametrize("jit", [False, True], ids=["eager", "jit"])
+def test_sdpa_attention_causal_backward(q_k_v_target_mask, dtype: str, jit: bool):
+    if jit:
+        pytest.skip("`sdpa_attention_forward` can't be compiled with jit.")
+
+    module = MockAttentionModule(is_causal=True)
+
+    # MindONE
+    q, k, v, target, _ = cast_inputs(*q_k_v_target_mask["ms"], dtype=MS_DTYPE_MAPPING[dtype])
+
+    def _forward(q_, k_, v_, target_):
+        output = sdpa_attention_forward(module, q_, k_, v_, attention_mask=None)[0]
+        return mint.nn.functional.mse_loss(output, target_)
+
+    grad_out = grad(_forward, grad_position=(0, 1, 2))(q, k, v, target)
+    grad_out = mint.stack(grad_out, dim=0)
+
+    # Transformers
+    q_pt, k_pt, v_pt, target_pt, _ = cast_inputs(*q_k_v_target_mask["pt"], dtype=PT_DTYPE_MAPPING[dtype])
+    q_pt, k_pt, v_pt = q_pt.clone(), k_pt.clone(), v_pt.clone()
+    q_pt.requires_grad, k_pt.requires_grad, v_pt.requires_grad = (True,) * 3
+
+    output_pt = sdpa_attention_forward_transformers(module, q_pt, k_pt, v_pt, attention_mask=None)[0]
+    loss = torch.nn.functional.mse_loss(output_pt, target_pt)
+    loss.backward()
+    grad_out_pt = torch.stack([q_pt.grad, k_pt.grad, v_pt.grad], dim=0)
+
+    assert grad_out.shape == grad_out_pt.shape, f"Shape mismatch: {grad_out.shape} vs {grad_out_pt.shape}"
+    assert not grad_out.isnan().any(), "Output contains NaNs."
+    assert np.allclose(
+        grad_out.numpy().astype(np.float32), grad_out_pt.to(torch.float32).numpy(), atol=DTYPE_AND_THRESHOLDS[dtype]
+    )