Repository: facebookresearch/xformers

Issue:
Favor breaks with fp16
# 🐛 Bug
The loss nans immediately

## Command
`python3 microGPT.py` when precision is set to 16

## Environment
Unrelated



Discussion:

Resolving patch:
diff --git a/examples/microGPT.py b/examples/microGPT.py
--- a/examples/microGPT.py
+++ b/examples/microGPT.py
@@ -28,7 +28,7 @@ def __init__(
         vocab_size,
         weight_decay=0.1,
         betas=(0.9, 0.95),
-        learning_rate=6e-4,
+        learning_rate=1e-4,
         n_embd=512,
         block_size=128,
         n_layer=4,
@@ -298,7 +298,7 @@ def top_k_logits(logits, k):
     model = GPT(
         vocab_size=train_dataset.vocab_size,
         block_size=train_dataset.block_size,
-        attention="nystrom",
+        attention="favor",
         warmup_tokens=REF_BATCH * WARMUP,
         final_tokens=EPOCHS * len(train_dataset) * BLOCK,
     )
@@ -310,7 +310,7 @@ def top_k_logits(logits, k):
         precision=16,
         gradient_clip_val=1,
         log_every_n_steps=1,
-        terminate_on_nan=True,
+        detect_anomaly=True,
         accumulate_grad_batches=REF_BATCH // BATCH,
     )
 
diff --git a/xformers/components/attention/favor.py b/xformers/components/attention/favor.py
--- a/xformers/components/attention/favor.py
+++ b/xformers/components/attention/favor.py
@@ -10,6 +10,7 @@
 
 import torch
 import torch.nn as nn
+from torch.cuda.amp import autocast
 
 from xformers.components.attention import Attention, AttentionConfig, register_attention
 from xformers.components.attention.feature_maps import (
@@ -99,6 +100,11 @@ def __init__(
         self.feature_map_query: FeatureMap = feature_map_constructor(**feature_settings)  # type: ignore
         self.feature_map_key: FeatureMap = feature_map_constructor(**feature_settings)  # type: ignore
 
+    @staticmethod
+    def _maybe_promote(x: torch.Tensor) -> torch.Tensor:
+        # Only promote fp16 buffers, bfloat16 would be fine for instance
+        return x.float() if x.dtype == torch.float16 else x
+
     @staticmethod
     def _causal_attention(
         k_prime: torch.Tensor, q_prime: torch.Tensor, v: torch.Tensor
@@ -135,17 +141,26 @@ def forward(
         # Project key and queries onto the feature map space
         k_prime = self.feature_map_key(k)
         q_prime = self.feature_map_query(q)
-        if not self.causal:
-            att_normalization = q_prime @ (
-                k_prime.transpose(-2, -1) @ torch.ones_like(v)
-            )
-            att_raw = q_prime @ (k_prime.transpose(-2, -1) @ v)
-        else:
-            # Actually compute attention
-            att_raw, att_normalization = self._causal_attention(k_prime, q_prime, v)
 
-        # Normalize
-        att = att_raw / att_normalization
+        with autocast(enabled=False):
+            # The softmax kernel approximation for Favor will easily overflow
+            # Force the computations here to stay in fp32 for numerical stability
+            # Note that the dimensions are vastly reduced when compared to scaled_dot_product
+            k_prime = self._maybe_promote(k_prime)
+            q_prime = self._maybe_promote(q_prime)
+            v = self._maybe_promote(v)
+
+            if not self.causal:
+                att_normalization = q_prime @ (
+                    k_prime.transpose(-2, -1) @ torch.ones_like(v)
+                )
+                att_raw = q_prime @ (k_prime.transpose(-2, -1) @ v)
+            else:
+                # Actually compute attention
+                att_raw, att_normalization = self._causal_attention(k_prime, q_prime, v)
+
+            # Normalize
+            att = att_raw / att_normalization
 
         if self.attn_drop is not None:
             att = self.attn_drop(att)


Test patch: