Repository: google/flax

Issue:
Force `fp32` in `attention.MultiHeadDotProductAttention` for softmax operator
Hi,

As I was trying to implement mixed precision training under Flax for my project, I noticed that the `force_fp32_for_softmax` flag defined in `attention.MultiHeadDotProductAttention` does not get to pass into `dot_product_attention` (the default attention function). 
<img width="464" alt="Screenshot 2024-06-18 at 11 18 38 PM" src="https://github.com/google/flax/assets/90777911/81ee2c74-2dc9-4fd7-a76c-2d9c921df8ff">
<img width="461" alt="Screenshot 2024-06-18 at 11 17 55 PM" src="https://github.com/google/flax/assets/90777911/93d2c348-3837-4fa4-b74d-30778c7e16f6">
<img width="441" alt="Screenshot 2024-06-18 at 11 19 21 PM" src="https://github.com/google/flax/assets/90777911/b2daa7c9-45ef-41ca-aa2d-db594484dd35">

I think this might lead to loose control over the softmax operator and result in some stability issues under `bf16` or `fp16` precision, so I wonder if there's an alternate? Thanks!


Discussion:

Patch:
diff --git a/flax/linen/attention.py b/flax/linen/attention.py
index b7e9de6c..6e91be99 100644
--- a/flax/linen/attention.py
+++ b/flax/linen/attention.py
@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import functools
+import inspect
 import warnings
 from typing import Any, Callable, Optional, Union, overload
 
@@ -574,33 +575,28 @@ class MultiHeadDotProductAttention(Module):
       m_deterministic = True
 
     # apply attention
+    attn_args = (query, key, value)
+    # This kwargs list match the default nn.dot_product_attention.
+    # For custom `attention_fn`s, invalid kwargs will be filtered.
+    attn_kwargs = dict(
+      mask=mask,
+      dropout_rng=dropout_rng,
+      dropout_rate=self.dropout_rate,
+      broadcast_dropout=self.broadcast_dropout,
+      deterministic=m_deterministic,
+      dtype=self.dtype,
+      precision=self.precision,
+      force_fp32_for_softmax=self.force_fp32_for_softmax,
+    )
+    attn_kwargs = {
+        k: v
+        for k, v in attn_kwargs.items()
+        if k in inspect.signature(self.attention_fn).parameters
+    }
     if sow_weights:
-      x = self.attention_fn(
-        query,
-        key,
-        value,
-        mask=mask,
-        dropout_rng=dropout_rng,
-        dropout_rate=self.dropout_rate,
-        broadcast_dropout=self.broadcast_dropout,
-        deterministic=m_deterministic,
-        dtype=self.dtype,
-        precision=self.precision,
-        module=self,
-      )  # pytype: disable=wrong-keyword-args
+      x = self.attention_fn(*attn_args, **attn_kwargs, module=self)
     else:
-      x = self.attention_fn(
-        query,
-        key,
-        value,
-        mask=mask,
-        dropout_rng=dropout_rng,
-        dropout_rate=self.dropout_rate,
-        broadcast_dropout=self.broadcast_dropout,
-        deterministic=m_deterministic,
-        dtype=self.dtype,
-        precision=self.precision,
-      )
+      x = self.attention_fn(*attn_args, **attn_kwargs)
     # back to the original inputs dimensions
     out = DenseGeneral(
       features=features,


Test patch:
diff --git a/tests/linen/linen_attention_test.py b/tests/linen/linen_attention_test.py
index 7dd07878..e7d69a43 100644
--- a/tests/linen/linen_attention_test.py
+++ b/tests/linen/linen_attention_test.py
@@ -14,7 +14,6 @@
 
 """Tests for flax.linen.attention."""
 
-import functools
 from absl.testing import absltest, parameterized
 from flax import errors, jax_utils
 from flax import linen as nn
@@ -565,9 +564,7 @@ class AttentionTest(parameterized.TestCase):
         qkv_features=4,
         kernel_init=initializers.lecun_normal(),
         bias_init=initializers.uniform(),
-        attention_fn=functools.partial(
-            nn.dot_product_attention, force_fp32_for_softmax=force_fp32
-        ),
+        force_fp32_for_softmax=force_fp32,
         deterministic=False,
         dtype=jnp.bfloat16,
     )