prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Rotary embeddings convert queries and keys to float32 when using mixed precision training
  Not sure, if this is expected behaviour. The problem is that the cos_sin table in the RotaryEmbedding class is stored in float32 format, thus the returned queries and keys get converted to float32 aswell.
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `xformers/components/positional_embedding/rotary.py`
  Content:
  ```
  1 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
  2 #
  3 # This source code is licensed under the BSD license found in the
  4 # LICENSE file in the root directory of this source tree.
  5 
  6 
  7 # CREDITS: This implementation is inspired by GPT-NeoX https://github.com/EleutherAI/gpt-neox
  8 # NOTE: Almost the same right now, moving parts to Triton is the next step
  9 
  10 from typing import Tuple
  11 
  12 import torch
  13 
  14 
  15 def rotate_half(x):
  16     x1, x2 = x.chunk(2, dim=-1)
  17     return torch.cat((-x2, x1), dim=-1)
  18 
  19 
  20 @torch.jit.script
  21 def apply_rotary_pos_emb(x, cos, sin):
  22     # NOTE: This could probably be moved to Triton
  23 
  24     # Handle a possible sequence length mismatch in between q and k
  25     cos = cos[:, :, : x.shape[-2], :]
  26     sin = sin[:, :, : x.shape[-2], :]
  27 
  28     return (x * cos) + (rotate_half(x) * sin)
  29 
  30 
  31 class RotaryEmbedding(torch.nn.Module):
  32     """
  33     The rotary position embeddings from RoFormer_ (Su et. al).
  34     A crucial insight from the method is that the query and keys are
  35     transformed by rotation matrices which depend on the relative positions.
  36 
  37     Other implementations are available in the Rotary Transformer repo_ and in
  38     GPT-NeoX_, GPT-NeoX was an inspiration
  39 
  40     .. _RoFormer: https://arxiv.org/abs/2104.09864
  41     .. _repo: https://github.com/ZhuiyiTechnology/roformer
  42     .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox
  43 
  44 
  45     .. warning: Please note that this embedding is not registered on purpose, as it is transformative
  46         (it does not create the embedding dimension) and will likely be picked up (imported) on a ad-hoc basis
  47     """
  48 
  49     def __init__(self, dim_model: int, *_, **__):
  50         super().__init__()
  51         # Generate and save the inverse frequency buffer (non trainable)
  52         inv_freq = 1.0 / (10000 ** (torch.arange(0, dim_model, 2).float() / dim_model))
  53         self.register_buffer("inv_freq", inv_freq)
  54 
  55         self._seq_len_cached = None
  56         self._cos_cached = None
  57         self._sin_cached = None
  58 
  59     def _update_cos_sin_tables(self, x, seq_dimension=1):
  60         seq_len = x.shape[seq_dimension]
  61 
  62         # Reset the tables if the sequence length has changed,
  63         # or if we're on a new device (possibly due to tracing for instance)
  64         if seq_len != self._seq_len_cached or self._cos_cached.device != x.device:
  65             self._seq_len_cached = seq_len
  66             t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(
  67                 self.inv_freq
  68             )
  69             freqs = torch.einsum("i,j->ij", t, self.inv_freq)
  70             emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
  71 
  72             self._cos_cached = emb.cos()[None, None, :, :]
  73             self._sin_cached = emb.sin()[None, None, :, :]
  74 
  75         return self._cos_cached, self._sin_cached
  76 
  77     def forward(
  78         self, q: torch.Tensor, k: torch.Tensor
  79     ) -> Tuple[torch.Tensor, torch.Tensor]:
  80         self._cos_cached, self._sin_cached = self._update_cos_sin_tables(
  81             k, seq_dimension=-2
  82         )
  83 
  84         return (
  85             apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),
  86             apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),
  87         )
  88
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/xformers/components/positional_embedding/rotary.py b/xformers/components/positional_embedding/rotary.py
  --- a/xformers/components/positional_embedding/rotary.py
  +++ b/xformers/components/positional_embedding/rotary.py
  @@ -61,16 +61,20 @@
   
           # Reset the tables if the sequence length has changed,
           # or if we're on a new device (possibly due to tracing for instance)
  -        if seq_len != self._seq_len_cached or self._cos_cached.device != x.device:
  +        if (
  +            seq_len != self._seq_len_cached
  +            or self._cos_cached.device != x.device
  +            or self._cos_cached.dtype != x.dtype
  +        ):
               self._seq_len_cached = seq_len
  -            t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(
  -                self.inv_freq
  +            t = torch.arange(
  +                x.shape[seq_dimension], device=x.device, dtype=torch.float32
               )
  -            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
  +            freqs = torch.einsum("i,j->ij", t, self.inv_freq.to(x.dtype))
               emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
   
  -            self._cos_cached = emb.cos()[None, None, :, :]
  -            self._sin_cached = emb.sin()[None, None, :, :]
  +            self._cos_cached = emb.cos()[None, None, :, :].to(x.dtype)
  +            self._sin_cached = emb.sin()[None, None, :, :].to(x.dtype)
   
           return self._cos_cached, self._sin_cached
  
verification_info: {"golden_diff": "diff --git a/xformers/components/positional_embedding/rotary.py b/xformers/components/positional_embedding/rotary.py\n--- a/xformers/components/positional_embedding/rotary.py\n+++ b/xformers/components/positional_embedding/rotary.py\n@@ -61,16 +61,20 @@\n \n         # Reset the tables if the sequence length has changed,\n         # or if we're on a new device (possibly due to tracing for instance)\n-        if seq_len != self._seq_len_cached or self._cos_cached.device != x.device:\n+        if (\n+            seq_len != self._seq_len_cached\n+            or self._cos_cached.device != x.device\n+            or self._cos_cached.dtype != x.dtype\n+        ):\n             self._seq_len_cached = seq_len\n-            t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(\n-                self.inv_freq\n+            t = torch.arange(\n+                x.shape[seq_dimension], device=x.device, dtype=torch.float32\n             )\n-            freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq)\n+            freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq.to(x.dtype))\n             emb = torch.cat((freqs, freqs), dim=-1).to(x.device)\n \n-            self._cos_cached = emb.cos()[None, None, :, :]\n-            self._sin_cached = emb.sin()[None, None, :, :]\n+            self._cos_cached = emb.cos()[None, None, :, :].to(x.dtype)\n+            self._sin_cached = emb.sin()[None, None, :, :].to(x.dtype)\n \n         return self._cos_cached, self._sin_cached\n", "issue": "Rotary embeddings convert queries and keys to float32 when using mixed precision training\nNot sure, if this is expected behaviour. The problem is that the cos_sin table in the RotaryEmbedding class is stored in float32 format, thus the returned queries and keys get converted to float32 aswell.\n", "before_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\n\n# CREDITS: This implementation is inspired by GPT-NeoX https://github.com/EleutherAI/gpt-neox\n# NOTE: Almost the same right now, moving parts to Triton is the next step\n\nfrom typing import Tuple\n\nimport torch\n\n\ndef rotate_half(x):\n    x1, x2 = x.chunk(2, dim=-1)\n    return torch.cat((-x2, x1), dim=-1)\n\n\n@torch.jit.script\ndef apply_rotary_pos_emb(x, cos, sin):\n    # NOTE: This could probably be moved to Triton\n\n    # Handle a possible sequence length mismatch in between q and k\n    cos = cos[:, :, : x.shape[-2], :]\n    sin = sin[:, :, : x.shape[-2], :]\n\n    return (x * cos) + (rotate_half(x) * sin)\n\n\nclass RotaryEmbedding(torch.nn.Module):\n    \"\"\"\n    The rotary position embeddings from RoFormer_ (Su et. al).\n    A crucial insight from the method is that the query and keys are\n    transformed by rotation matrices which depend on the relative positions.\n\n    Other implementations are available in the Rotary Transformer repo_ and in\n    GPT-NeoX_, GPT-NeoX was an inspiration\n\n    .. _RoFormer: https://arxiv.org/abs/2104.09864\n    .. _repo: https://github.com/ZhuiyiTechnology/roformer\n    .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox\n\n\n    .. warning: Please note that this embedding is not registered on purpose, as it is transformative\n        (it does not create the embedding dimension) and will likely be picked up (imported) on a ad-hoc basis\n    \"\"\"\n\n    def __init__(self, dim_model: int, *_, **__):\n        super().__init__()\n        # Generate and save the inverse frequency buffer (non trainable)\n        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim_model, 2).float() / dim_model))\n        self.register_buffer(\"inv_freq\", inv_freq)\n\n        self._seq_len_cached = None\n        self._cos_cached = None\n        self._sin_cached = None\n\n    def _update_cos_sin_tables(self, x, seq_dimension=1):\n        seq_len = x.shape[seq_dimension]\n\n        # Reset the tables if the sequence length has changed,\n        # or if we're on a new device (possibly due to tracing for instance)\n        if seq_len != self._seq_len_cached or self._cos_cached.device != x.device:\n            self._seq_len_cached = seq_len\n            t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(\n                self.inv_freq\n            )\n            freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq)\n            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)\n\n            self._cos_cached = emb.cos()[None, None, :, :]\n            self._sin_cached = emb.sin()[None, None, :, :]\n\n        return self._cos_cached, self._sin_cached\n\n    def forward(\n        self, q: torch.Tensor, k: torch.Tensor\n    ) -> Tuple[torch.Tensor, torch.Tensor]:\n        self._cos_cached, self._sin_cached = self._update_cos_sin_tables(\n            k, seq_dimension=-2\n        )\n\n        return (\n            apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),\n            apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),\n        )\n", "path": "xformers/components/positional_embedding/rotary.py"}], "after_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.\n#\n# This source code is licensed under the BSD license found in the\n# LICENSE file in the root directory of this source tree.\n\n\n# CREDITS: This implementation is inspired by GPT-NeoX https://github.com/EleutherAI/gpt-neox\n# NOTE: Almost the same right now, moving parts to Triton is the next step\n\nfrom typing import Tuple\n\nimport torch\n\n\ndef rotate_half(x):\n    x1, x2 = x.chunk(2, dim=-1)\n    return torch.cat((-x2, x1), dim=-1)\n\n\n@torch.jit.script\ndef apply_rotary_pos_emb(x, cos, sin):\n    # NOTE: This could probably be moved to Triton\n\n    # Handle a possible sequence length mismatch in between q and k\n    cos = cos[:, :, : x.shape[-2], :]\n    sin = sin[:, :, : x.shape[-2], :]\n\n    return (x * cos) + (rotate_half(x) * sin)\n\n\nclass RotaryEmbedding(torch.nn.Module):\n    \"\"\"\n    The rotary position embeddings from RoFormer_ (Su et. al).\n    A crucial insight from the method is that the query and keys are\n    transformed by rotation matrices which depend on the relative positions.\n\n    Other implementations are available in the Rotary Transformer repo_ and in\n    GPT-NeoX_, GPT-NeoX was an inspiration\n\n    .. _RoFormer: https://arxiv.org/abs/2104.09864\n    .. _repo: https://github.com/ZhuiyiTechnology/roformer\n    .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox\n\n\n    .. warning: Please note that this embedding is not registered on purpose, as it is transformative\n        (it does not create the embedding dimension) and will likely be picked up (imported) on a ad-hoc basis\n    \"\"\"\n\n    def __init__(self, dim_model: int, *_, **__):\n        super().__init__()\n        # Generate and save the inverse frequency buffer (non trainable)\n        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim_model, 2).float() / dim_model))\n        self.register_buffer(\"inv_freq\", inv_freq)\n\n        self._seq_len_cached = None\n        self._cos_cached = None\n        self._sin_cached = None\n\n    def _update_cos_sin_tables(self, x, seq_dimension=1):\n        seq_len = x.shape[seq_dimension]\n\n        # Reset the tables if the sequence length has changed,\n        # or if we're on a new device (possibly due to tracing for instance)\n        if (\n            seq_len != self._seq_len_cached\n            or self._cos_cached.device != x.device\n            or self._cos_cached.dtype != x.dtype\n        ):\n            self._seq_len_cached = seq_len\n            t = torch.arange(\n                x.shape[seq_dimension], device=x.device, dtype=torch.float32\n            )\n            freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq.to(x.dtype))\n            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)\n\n            self._cos_cached = emb.cos()[None, None, :, :].to(x.dtype)\n            self._sin_cached = emb.sin()[None, None, :, :].to(x.dtype)\n\n        return self._cos_cached, self._sin_cached\n\n    def forward(\n        self, q: torch.Tensor, k: torch.Tensor\n    ) -> Tuple[torch.Tensor, torch.Tensor]:\n        self._cos_cached, self._sin_cached = self._update_cos_sin_tables(\n            k, seq_dimension=-2\n        )\n\n        return (\n            apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),\n            apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),\n        )\n", "path": "xformers/components/positional_embedding/rotary.py"}]}