Repository: speechbrain/speechbrain

Issue:
`torch.arange` doesn't implemented for `float16`
# Problem 👩‍⚕️

I found that when `x` is in `float16` this line doesn't work:

https://github.com/speechbrain/speechbrain/blob/c7219b5b2b5c1b8a68da23e81eae0046b6769556/speechbrain/nnet/attention.py#L345

It is important because when I train the model in `fp16` (with `autocast` for ex) it throws an exception. 

# Repro 📔

```
 import torch

 seq_len = 10
 x = torch.randn(8, 50, 128).to(torch.float16)

 torch.arange(0, seq_len, dtype=x.dtype)
 ```
 
 Return:
 ```
 ---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
[<ipython-input-28-1c012662a902>](https://localhost:8080/#) in <module>()
      6 print(x.dtype)
      7 
----> 8 torch.arange(0, seq_len, dtype=x.dtype)

RuntimeError: "arange_cpu" not implemented for 'Half'
 ```
 
 It is reproducible in colab too
 
 # Workaround 👀
 
I found that in Nemo `arange` is `float32` by default in dtype

```
# one of examples
positions = torch.arange(length - 1, -length, -1, dtype=torch.float32, device=device).unsqueeze(1)
```
 https://github.com/NVIDIA/NeMo/blob/41f27a5e68934761b05dd373e7b288aa798858e5/nemo/collections/asr/parts/submodules/multi_head_attention.py#L288


Discussion:
Also, float16 has max value of 65536 and quite inaccurate in the higher range, so using it for length calculation (if at all) should be done with extreme caution
That is interesting, @popcornell ?
I also encountered this error. I tried to fix it, and it is quite easy to do so. 
However I got much worse performance for vanilla transformer. 
In my understanding half precision is broken right now. Didn't investigate further. 

> Also, float16 has max value of 65536 and quite inaccurate in the higher range, so using it for length calculation (if at all) should be done with extreme caution

Probably for ASR or NLP will be fine as you will not end up with such long sequences. 
But frankly nemo solution or using int64 is better. 

Also for float16 is badly supported on CPU (https://github.com/pytorch/pytorch/issues/80483#issuecomment-1169404072), it should probably be changed to default int64 or float32 that are supported on most devices. Probably positional encoding computation should always be done at least in fp32 for sin/cos arguments are not overflown/underflown. fairseq/nemo/huggingface can be a reference
@popcornell Do you mean model accuracy was much worse? or computation speed? did you get any nans/infs while training?
Here's Nemo's impl of positional encoding (also includes caching):
https://github.com/NVIDIA/NeMo/blob/633674898132b315b4f2eaf3761de391d3d0c1f4/nemo/collections/asr/parts/submodules/multi_head_attention.py#L231-L244
@vadimkantorov the model didn't converge and around maximum lr by lr_scheduling I got nans
Hello,

any news on this issue, please? 😃 
float16 was implemented for arange in core pytorch. but i think it's still a bad idea to do arange in float16. this implicitly assumes small tensor sizes, better would be to remove `dtype=x.dtype` which would make it use int64
OK. Thanks.

Have you done any tests between float16 and int64 on our transformers recipes? 


No, @grazder just fixed it to be int64 (like in nemo) to be safe of overflow/underflow of fp16
OK. 

I will need to train some transformer recipes to test it. But I agree that it would be better to have int64 or fp32. 
I've looked into this again and I think there is no reason _not_ to use fp32 in the `arange`:

1. Checking out the `dtype` of the input tensor might not be a good idea (here or elsewhere). Even if you are using fp16 AMP, by design, not all tensors will be fp16, depending on the previous operation. The end result may be unnecessary extra casting or lost precision in certain cases.
2. Using `int64` does not seem meaningful as the result will get ultimately converted to floating-point and fp32 is enough precision and range anyway.
3. ~~Since `positions` is immediately multiplied by `self.inv_freq`, which itself is fp32, it is promoted to fp32 anyway in the end.~~ (Nevermind, it _does_ get downcasted into a fp16 buffer at the end).

Extra note: the vanilla positional encodings are not affected and unconditionally use fp32.

Resolving patch:
diff --git a/speechbrain/nnet/attention.py b/speechbrain/nnet/attention.py
--- a/speechbrain/nnet/attention.py
+++ b/speechbrain/nnet/attention.py
@@ -323,9 +323,20 @@ def forward(self, enc_states, enc_len, dec_states):
 
 
 class RelPosEncXL(nn.Module):
-    """ """
+    """Relative positional encoding for the :class:`~RelPosMHAXL`.
 
-    def __init__(self, emb_dim):
+    Arguments
+    ---------
+    emb_dim : int
+        Size of the embedding, which controls the size of the last dimension
+        of the positional embedding as well
+    dtype : torch.dtype, optional
+        If unspecified, defaults to `torch.float32`. Controls the data type of
+        the output embedding (but does not affect the precision of the
+        computations, which remain `torch.float32`).
+    """
+
+    def __init__(self, emb_dim: int, dtype: torch.dtype = torch.float32):
         super().__init__()
         self.emb_dim = emb_dim
 
@@ -335,29 +346,46 @@ def __init__(self, emb_dim):
         )
         self.register_buffer("inv_freq", inv_freq)
 
-    def forward(self, x: torch.Tensor):
+        self.emb_dtype = dtype
+
+    @torch.no_grad()
+    def make_pe(self, seq_len: int):
         """
+        Builds the positional embedding tensor for a given sequence length.
+
         Arguments
         ---------
-        x : torch.Tensor
-            input tensor with shape batch_size, seq_len, embed_dim
+        seq_len : int
+            The length of the sequence to create the position embedding for.
 
         Returns
         -------
-        pos_emb : torch.Tensor
+        torch.Tensor
+            Positional embedding tensor of shape `[1, 2*seq_len-1, embed_dim]`
         """
-        seq_len = x.size(1)
+
+        emb_dtype = self.emb_dtype
+        device = self.inv_freq.device
+
         with torch.no_grad():
-            tot_pe = torch.zeros((2, seq_len, self.emb_dim), dtype=x.dtype).to(
-                x
+            # perform initialization with the same type as `inv_freq`, to enable
+            # migrating the embeddings to fp16 by calling
+            # `posenc.to(torch.float16)`
+
+            tot_pe = torch.empty(
+                (2, seq_len, self.emb_dim),
+                dtype=torch.float32,
+                device=device,
             )
             pe_past = tot_pe[0]
             pe_future = tot_pe[1]
-            positions = (
-                torch.arange(0, seq_len, dtype=x.dtype, device=x.device)
-                .to(x)
-                .unsqueeze(-1)
-            )
+            positions = torch.arange(
+                0,
+                seq_len,
+                dtype=torch.float32,
+                device=device,
+            ).unsqueeze(-1)
+
             sinusoids = torch.sin(positions * self.inv_freq)
             pe_past[:, 0::2] = sinusoids
             pe_past[:, 1::2] = torch.cos(positions * self.inv_freq)
@@ -367,8 +395,28 @@ def forward(self, x: torch.Tensor):
             pe_past = torch.flip(pe_past, (0,)).unsqueeze(0)
             pe_future = pe_future[1:].unsqueeze(0)
             pe = torch.cat([pe_past, pe_future], dim=1)
-            # pe is now 1, 2*seq_len, embed_dim
-            return pe
+            pe = pe.to(emb_dtype)  # convert to type of module
+
+        return pe
+
+    def forward(self, x: torch.Tensor):
+        """
+        Builds the positional embedding tensor. Similar to
+        :meth:`~RelPosEncXL.make_pe` but uses the shape information from the
+        provided tensor.
+
+        Arguments
+        ---------
+        x : torch.Tensor
+            input tensor with shape batch_size, seq_len, embed_dim
+
+        Returns
+        -------
+        pos_emb : torch.Tensor
+            Positional embedding tensor of shape `[1, 2*seq_len-1, embed_dim]`
+        """
+
+        return self.make_pe(seq_len=x.size(1))
 
 
 class RelPosMHAXL(nn.Module):


Test patch: