Add flag for TF32 self-attention layers in Psiformer. Thanks to Seb Bodenstein for investigating.

PiperOrigin-RevId: 831808752
Change-Id: Id5d61e4d47adb5ceddb8354ac5bf5bd14a2d1d2d

Author: dpfau

ferminet/base_config.py

@@ -270,6 +270,7 @@ def default() -> ml_collections.ConfigDict:
               'heads_dim': 64,
               'mlp_hidden_dims': (256,),
               'use_layer_norm': True,
+              'tf32': False,
           },
           # Config common to all architectures.
           'determinants': 16,  # Number of determinants.

ferminet/psiformer.py

@@ -38,13 +38,15 @@ class PsiformerOptions(networks.BaseNetworkOptions):
     mlp_hidden_dims: Tuple of sizes of hidden dimension of the MLP. Note that
       this does not include the final projection to the embedding dimension.
     use_layer_norm: If true, include a layer norm on both attention and MLP.
+    tf32: If true, use TF32-precision matmuls where appropriate.
   """
 
   num_layers: int = 2
   num_heads: int = 4
   heads_dim: int = 64
   mlp_hidden_dims: Tuple[int, ...] = (256,)
   use_layer_norm: bool = False
+  tf32: bool = False
 
 
 def make_layer_norm() ->...:
@@ -68,12 +70,15 @@ def apply(params: networks.ParamTree,
   return init, apply
 
 
-def make_multi_head_attention(num_heads: int, heads_dim: int) ->...:
+def make_multi_head_attention(num_heads: int,
+                              heads_dim: int,
+                              tf32: bool = False) ->...:
   """FermiNet-style version of MultiHeadAttention."""
+  prec = jax.lax.DotAlgorithmPreset.TF32_TF32_F32 if tf32 else None
 
   # Linear layer plus reshape final dimensions to num_heads, heads_dim.
   def linear_projection(x: jnp.ndarray, weights: jnp.ndarray) -> jnp.ndarray:
-    y = jnp.dot(x, weights)
+    y = jnp.dot(x, weights, precision=prec)
     return y.reshape(*x.shape[:-1], num_heads, heads_dim)
 
   def init(key: chex.PRNGKey,
@@ -123,13 +128,13 @@ def apply(params: networks.ParamTree, query: jnp.ndarray, key: jnp.ndarray,
     k = linear_projection(key, params['k_w'])
     v = linear_projection(value, params['v_w'])
 
-    attn_logits = jnp.einsum('...thd,...Thd->...htT', q, k)
+    attn_logits = jnp.einsum('...thd,...Thd->...htT', q, k, precision=prec)
     scale = 1. / np.sqrt(heads_dim)
     attn_logits *= scale
 
     attn_weights = jax.nn.softmax(attn_logits)
 
-    attn = jnp.einsum('...htT,...Thd->...thd', attn_weights, v)
+    attn = jnp.einsum('...htT,...Thd->...thd', attn_weights, v, precision=prec)
 
     # Concatenate attention matrix of all heads into a single vector.
     # Shape [..., q_index_dim, num_heads * heads_dim]
@@ -174,10 +179,11 @@ def make_self_attention_block(num_layers: int,
                               num_heads: int,
                               heads_dim: int,
                               mlp_hidden_dims: Tuple[int, ...],
-                              use_layer_norm: bool = False) ->...:
+                              use_layer_norm: bool = False,
+                              tf32: bool = False) ->...:
   """Create a QKV self-attention block."""
   attention_init, attention_apply = make_multi_head_attention(
-      num_heads, heads_dim)
+      num_heads, heads_dim, tf32)
   if use_layer_norm:
     layer_norm_init, layer_norm_apply = make_layer_norm()
   mlp_init, mlp_apply = make_mlp()
@@ -253,6 +259,7 @@ def make_psiformer_layers(
       heads_dim=options.heads_dim,
       mlp_hidden_dims=options.mlp_hidden_dims,
       use_layer_norm=options.use_layer_norm,
+      tf32=options.tf32,
   )
 
   def init(key: chex.PRNGKey) -> Tuple[int, networks.ParamTree]:
@@ -342,6 +349,7 @@ def make_fermi_net(
     heads_dim: int,
     mlp_hidden_dims: Tuple[int, ...],
     use_layer_norm: bool,
+    tf32: bool,
 ) -> networks.Network:
   """Psiformer with stacked Self Attention layers.
 
@@ -365,6 +373,7 @@ def make_fermi_net(
     heads_dim: Embedding dimension per-head for self-attention.
     mlp_hidden_dims: Tuple of hidden dimensions of the MLP.
     use_layer_norm: If true, use layer_norm on both attention and MLP.
+    tf32: If true, use TF32-precision matmuls where appropriate.
 
   Returns:
     Network object containing init, apply, orbitals, options, where init and
@@ -404,6 +413,7 @@ def make_fermi_net(
       heads_dim=heads_dim,
       mlp_hidden_dims=mlp_hidden_dims,
       use_layer_norm=use_layer_norm,
+      tf32=tf32,
   )  # pytype: disable=wrong-keyword-args
 
   psiformer_layers = make_psiformer_layers(nspins, charges.shape[0], options)

ferminet/tests/psiformer_test.py

@@ -79,6 +79,7 @@ def test_antisymmetry(self, jastrow):
         heads_dim=32,
         mlp_hidden_dims=(64, 128),
         use_layer_norm=True,
+        tf32=False,
     )
 
     key, subkey = jax.random.split(key)
@@ -145,6 +146,7 @@ def test_psiformer(self, **network_options):
         'heads_dim': 128,
         'mlp_hidden_dims': (64, 32),
         'use_layer_norm': True,
+        'tf32': False,
     }
 
     network = psiformer.make_fermi_net(