Repository: pytorch
Issue: torch.cuda.amp fails with torch.sparse.softmax
Labels: module: sparse,triaged,module: amp (automated mixed precision)

## 🐛 Bug

Models using `torch.sparse.softmax` fail under `torch.cuda.amp.autocast` with `RuntimeError: "softmax" not implemented for 'Half'` even though softmax would be carried out in f32 anyway. Fixing this should merely require blacklisting `torch.sparse.softmax` from autocast f16 ops.

## To Reproduce

```python
import torch
import torch.nn as nn

class NetDense(nn.Module):
    '''
    A baseline dense network. A linear layer followed by softmax
    '''
    def __init__(self):
        super(NetDense, self).__init__()
        
        self.lin = nn.Linear(4, 5)
    
    def forward(self, x):
        assert x.shape == (3, 4)
        y = self.lin(x)
        print(f'{self.__class__.__name__}: dtype before softmax {y.dtype=}')
        y = y.softmax(dim=1)
        print(f'{self.__class__.__name__}: dtype after softmax {y.dtype=}')
        
        return y

class NetSparse(nn.Module):
    '''
    The same as NetDense except softmax is done on a sparse tensor
    '''
    def __init__(self):
        super(NetSparse, self).__init__()
        
        self.lin = nn.Linear(4, 5)

    def forward(self, x):
        assert x.shape == (3, 4)
        y = self.lin(x)
        
        # convert y to sparse. this constructs a sparse tensor equal to the
        # original dense tensor (y_sparse.to_dense() == y)
        y_sparse = torch.sparse_coo_tensor(
            torch.nonzero(y).T,
            y.flatten(),
            y.shape,
        )
        
        print(f'{self.__class__.__name__}: dtype before softmax {y_sparse.dtype=}')
        y = torch.sparse.softmax(y_sparse, dim=1)
        print(f'{self.__class__.__name__}: dtype after softmax {y.dtype=}')

        y = y.to_dense() # go back to dense because otherwise y.sum() will fail
        return y

def loss_fn(y):
    return y.sum()

device = torch.device('cuda')
input = torch.randn(3, 4, device=device)

def try_it(net_cls):
    model = net_cls().to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

    # for input, target in data:
    optimizer.zero_grad()

    # Enables autocasting for the forward pass (model + loss)
    with torch.cuda.amp.autocast():
        output = model(input)
        loss = loss_fn(output)

    # Exits the context manager before backward()
    loss.backward()
    optimizer.step()

try_it(NetDense)
### OUTPUT:
# NetDense: dtype before softmax y.dtype=torch.float16
# NetDense: dtype after softmax y.dtype=torch.float32

try_it(NetSparse)
### OUTPUT:
# NetSparse: dtype before softmax y_sparse.dtype=torch.float16
# throws
```

Stack trace:

```
RuntimeError                              Traceback (most recent call last)
<ipython-input-8-74aa554852fd> in <module>
----> 1 try_it(NetSparse)

<ipython-input-6-b655f9e08017> in try_it(net_cls)
      8     # Enables autocasting for the forward pass (model + loss)
      9     with torch.cuda.amp.autocast():
---> 10         output = model(input)
     11         loss = loss_fn(output)
     12 

/mnt/disk-1/michat/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

<ipython-input-3-1940d9e6e811> in forward(self, x)
     18 
     19         print(f'{self.__class__.__name__}: dtype before softmax {y_sparse.dtype=}')
---> 20         y = torch.sparse.softmax(y_sparse, dim=1)
     21         print(f'{self.__class__.__name__}: dtype after softmax {y.dtype=}')
     22 

/mnt/disk-1/michat/miniconda3/lib/python3.8/site-packages/torch/sparse/__init__.py in softmax(input, dim, dtype)
    184           overflows. Default: None
    185     """
--> 186     return torch._sparse_softmax(input, dim, dtype=dtype)
    187 
    188 

RuntimeError: "softmax" not implemented for 'Half'
```
## Expected behavior

Successful optimizer step in both cases (with y_sparse being autocast to fp32 before proceeding with `torch.sparse.softmax`).
<!-- A clear and concise description of what you expected to happen. -->

## Environment
PyTorch version: 1.9.0
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.8.10 (default, Jun  4 2021, 15:09:15)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.8.0-1038-gcp-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-16GB
GPU 1: Tesla V100-SXM2-16GB
GPU 2: Tesla V100-SXM2-16GB
GPU 3: Tesla V100-SXM2-16GB
GPU 4: Tesla V100-SXM2-16GB
GPU 5: Tesla V100-SXM2-16GB
GPU 6: Tesla V100-SXM2-16GB
GPU 7: Tesla V100-SXM2-16GB

Nvidia driver version: 470.57.02
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.3
[pip3] pytorch-lightning==1.4.0
[pip3] torch==1.9.0
[pip3] torch-dimcheck==0.0.1
[pip3] torchaudio==0.9.0a0+33b2469
[pip3] torchmetrics==0.4.1
[pip3] torchvision==0.10.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.1.74              h6bb024c_0    nvidia
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.3.0           h06a4308_520  
[conda] mkl-service               2.4.0            py38h7f8727e_0  
[conda] mkl_fft                   1.3.0            py38h42c9631_2  
[conda] mkl_random                1.2.2            py38h51133e4_0  
[conda] numpy                     1.20.3           py38hf144106_0  
[conda] numpy-base                1.20.3           py38h74d4b33_0  
[conda] pytorch                   1.9.0           py3.8_cuda11.1_cudnn8.0.5_0    pytorch
[conda] pytorch-lightning         1.4.0                    pypi_0    pypi
[conda] torchaudio                0.9.0                      py38    pytorch
[conda] torchmetrics              0.4.1                    pypi_0    pypi
[conda] torchvision               0.10.0               py38_cu111    pytorch


cc @nikitaved @pearu @cpuhrsch @IvanYashchuk @mcarilli @ptrblck