Repository: pytorch
Issue: Inconsistent behavior of cosine_similarity between fp16 and fp32 inputs
Labels: high priority,module: numerical-stability,triaged,module: half,module: distance functions

## 🐛 Bug

`torch.nn.functional.cosine_similarity` outputs `0` for fp16 inputs and outputs normal results if converted to fp32.

## To Reproduce

Steps to reproduce the behavior:

```
import torch
import torch.nn.functional as F
device = torch.device('cuda')
a = torch.randn((1, 512), dtype=torch.float16, device=device)
b = torch.randn((1, 512), dtype=torch.float16, device=device)
print(F.cosine_similarity(a, b))
print(F.cosine_similarity(a.float(), b.float()))
print((a * b).sum(1) / a.norm() / b.norm())
print(F.cosine_similarity(a, a))
print(F.cosine_similarity(a.float(), a.float()))
print((a * a).sum(1) / a.norm() / a.norm())
```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

Actual output:
```
tensor([-0.], device='cuda:0', dtype=torch.float16)
tensor([-0.0136], device='cuda:0')
tensor([-0.0136], device='cuda:0', dtype=torch.float16)
tensor([0.], device='cuda:0', dtype=torch.float16)
tensor([1.], device='cuda:0')
tensor([0.9995], device='cuda:0', dtype=torch.float16)
```

I expect `cosine_similarity` to output correct results for fp16, but it just outputs zeros, while `(a * b).sum(1) / a.norm() / b.norm()` works pretty fine for fp16.

## Environment

```
PyTorch version: 1.9.1+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.10

Python version: 3.7.10 (default, Feb 26 2021, 18:47:35)  [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-5.11.0-41-generic-x86_64-with-debian-buster-sid
Is CUDA available: True
CUDA runtime version: 11.1.105
GPU models and configuration:
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090

Nvidia driver version: 470.86
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.0
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.3
[pip3] pytorch-lightning==1.5.2
[pip3] pytorch3d==0.6.0
[pip3] torch==1.9.1+cu111
[pip3] torchaudio==0.9.1
[pip3] torchmetrics==0.6.0
[pip3] torchvision==0.10.1+cu111
[pip3] vector-quantize-pytorch==0.3.3
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               10.2.89              hfd86e86_1
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.2.0           h06a4308_296
[conda] mkl-service               2.3.0            py37h27cfd23_1
[conda] mkl_fft                   1.3.0            py37h42c9631_2
[conda] mkl_random                1.2.1            py37ha9443f7_2
[conda] numpy                     1.20.2           py37h2d18471_0
[conda] numpy-base                1.20.2           py37hfae3a4d_0
[conda] pytorch                   1.8.1           py3.7_cuda10.2_cudnn7.6.5_0    pytorch
[conda] torchelastic              0.2.0                    pypi_0    pypi
[conda] torchtext                 0.9.1                      py37    pytorch
[conda] torchvision               0.9.1                py37_cu102    pytorch
```

## Additional context

<!-- Add any other context about the problem here. -->


cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @ngimel