Generate: should `softmax` be upcasted to `.float()`?
Author: gante
Pull request: no
Labels: name: Good Second Issue

I am looking for a contributor to do some numerical exploration across models! 💛 

### Context
Some operations are notoriously unstable at lower precisions, such as the `softmax`. In fact, in the models' attention layers, we often force the `softmax` operation to run in 32-bit FP precision ([example](https://github.com/huggingface/transformers/blob/90b4adc1f1111f42eada62ea611895646aaee6b6/src/transformers/models/llama/modeling_llama.py#L406)), regardless of the model `dtype`.

### Potential problem
In `generate`, we have a `softmax` operation when `do_sample=True`, and we are not upcasting it. Most models return the `logits` in the same `dtype` as the model itself -- a few notable exceptions being `llama` and `mistral`. Does this mean we are losing performance there?

### What we are looking for
A numerical exploration of the benefits (some metric) versus downsides (execution speed and memory consumption). 

If the returns turn out to be positive in favor of upcasting, you will be making a huge contribution to the whole community!


Comments:
Note: According to this [issue](https://github.com/pytorch/pytorch/pull/103167#issuecomment-1583326602), PyTorch already uses 32 bit accumulation under the hood. 

If true, we can do the opposite -- remove all manual casts to `torch.float32`

Hi @gante I would love to take up this issue.

Hey @nandwalritik 👋  Awesome, let me know how it goes and whether you need further pointers!

Hi @gante some further pointers will greatly help me to proceed further.

@nandwalritik 

We can start with the basics, assuming that what I wrote in [this comment](https://github.com/huggingface/transformers/issues/27341#issuecomment-1798510786) is true. To verify it, you can:
1. generate a random tensor in float16
2. pass it through `softmax` and `softmax(..., dtype=torch.float32)`
3. repeat it a couple of thousand times

If the outcome is exactly the same every single time, we can confirm the comment. In that case, we can close this issue.

If it is False, then report the average of the maximum absolute error, so we can decide how to proceed :)

I tried above steps, but `softmax` gives `RuntimeError: "softmax_lastdim_kernel_impl" not implemented for 'Half'` this error, which means that PyTorch does not have a kernel implementation for the softmax function on half-precision (float16) tensors.
To test if its automatically casting to `float32` I tried this
```
>>> tensor = torch.randn(100,100,dtype=torch.float16)
>>> sf_16 = nn.functional.softmax(tensor,dim=1,dtype=torch.float16)
```
but again it throws out same error.

I may be able to help out and make a test based on those steps as well. It could be useful to have two separate tests to verify or compare the results.

Can confirm the error @nandwalritik found. 

`softmax` does not seem to work with float16 tensors unless you force it to run in 32-bit precision.

```
>>> tensor = torch.randn(5, dtype=torch.float16)
>>> sm_32d = nn.functional.softmax(tensor, dim=-1)
>>> print(sm_32d)
RuntimeError: "softmax_lastdim_kernel_impl" not implemented for 'Half'
```

```
>>> tensor = torch.randn(5, dtype=torch.float16)
>>> sm_32d = nn.functional.softmax(tensor, dim=-1, dtype=torch.float16)
>>> print(sm_32d)
RuntimeError: "softmax_lastdim_kernel_impl" not implemented for 'Half'
```

```
>>> tensor = torch.randn(5, dtype=torch.float16)
>>> sm_32d = nn.functional.softmax(tensor, dim=-1, dtype=torch.float32)
>>> print(sm_32d)
tensor([0.3898, 0.0632, 0.1406, 0.2392, 0.1671])
```

Maybe that's why softmax is often forced to run in 32-bit FP precision? To accept tensors of any precision, it has to cast to a precision that softmax can use.

Hey @Tanman2001 @nandwalritik 👋 

Most FP16 operations are GPU-only, so you have to move the tensors to a CUDA device :) (e.g. `tensor = torch.randn(5, dtype=torch.float16, device="cuda")`)

Hi @gante I would like take up this issue. 

Feel free to report any conclusions here 🤗 

@gante Created a little test to compare the performance of the `softmax` function with and without the `dtype` parameter declared as float32 on a float16 tensor.

It is setup so that each iteration has two loops, where each loop is run 500 times. The first loop generates a random tensor of size 1000 then executes the `softmax` function with `dtype=torch.float32`. The second loop does the same but the `softmax` omits the `dtype` parameter. The time it takes each loop to execute the 500 times is recorded in every iteration.

Here's the results from my machine after 1000 iterations:
```
Default Times (softmax w/ default dtype param):
Median: 22.938966751098633
Mean: 24.283889055252075
Standard Deviation: 5.126754368188971
Minimum: 16.4642333984375
Maximum: 77.31127738952637

Forced Times (softmax w/ float32 dtype param:)
Median: 25.928974151611328
Mean: 26.8714861869812
Standard Deviation: 5.366920780984994
Minimum: 18.945932388305664
Maximum: 70.31035423278809
```

And I had someone else run it on their machine (which is apparently better than mine):
```
Default Times (softmax w/ default dtype param): 
Median: 19.004344940185547
Mean: 22.11885714530945
Standard Deviation: 11.286693376217634
Minimum: 12.00103759765625
Maximum: 46.010732650756836

Forced Times (softmax w/ float32 dtype param):
Median: 21.004796028137207
Mean: 23.721716165542603
Standard Deviation: 11.330417161274704
Minimum: 11.969327926635742
Maximum: 43.00975799560547
```

I've run this a few times and these results are typical. The loops with `softmax` forced to `float32` do appear to take slightly longer to execute. The mean and median times for the `float32` forced `softmax` are consistently higher (though not significantly) than the same times for the `softmax` without the `dtype` parameter.

Curious to see if @nandwalritik and @codeserra can corroborate these results.

@Tanman2001 Interesting, so we can see that `softmax` is faster without `dtype=torch.float32`. Two follow-up questions:
1. Which version of `torch` were you using?
2. Have you confirmed whether the two values (with and without `dtype=torch.float32`) matched for the same input?

@gante 

1. My machine has version 2.1.1+cu121. Other machine was 2.1.0+cu118.
2. I have not yet. I may be able to check that sometime this week.

Hi, followed what @Tanman2001 had described, for setting up 2 loops for described iterations and dtype.

I am getting below results, i am using torch version 1.12.1

Default Times (softmax w/ default dtype param):
Median: 0.015637
Mean: 0.015441
Standard Deviation: 0.006205
Minimum: 0.000000
Maximum: 0.032977

Forced Times (softmax w/ float32 dtype param):
Median: 0.015635
Mean: 0.015082
Standard Deviation: 0.006126
Minimum: 0.000000
Maximum: 0.040869

I am not seeing that much difference though, the differences i am seeing seems to be due precision change between 2 data types or due to random generator.

@codeserra 

How did you record and measure the running time? A 0.0 minimum time does not seem correct.

@Tanman2001 
I am using below code, i am fairly new at this, if any correction is required please let me know. 

```
import torch
import time
import numpy as np

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.manual_seed(42)

num_iterations = 1000
num_runs = 500
tensor_size = 1000

# Results list
default_times = []
forced_times = []

for iteration in range(num_iterations):
    # Record start time for softmax without dtype
    start_time_default = time.time()

    # softmax without dtype
    for _ in range(num_runs):
        input_tensor = torch.rand(tensor_size, dtype=torch.float32, device=device)
        result = torch.nn.functional.softmax(input_tensor.to(torch.float32), dim=0)
    end_time_default = time.time()

    # time for default dtype
    time_default = end_time_default - start_time_default
    default_times.append(time_default)

    start_time_forced = time.time()

    # softmax with float32 dtype 
    for _ in range(num_runs):
        input_tensor = torch.rand(tensor_size, dtype=torch.float32, device=device)
        result = torch.nn.functional.softmax(input_tensor, dim=0, dtype=torch.float32)

    end_time_forced = time.time()
    
    # time for forced time
    time_forced = end_time_forced - start_time_forced
    forced_times.append(time_forced)
```

@codeserra (and @gante correct me if I'm wrong)

The random input tensors should be `dtype=torch.float16` which would also require the device to be CUDA.

For the "without dtype" loop I think you are still manually casting the tensor to float32 with the `.to()` call. For that loop I would have had just `input_tensor` as that first parameter.

@gante 

Correctness test, assuming this is done correctly, fails.

```
import torch
import torch.nn as nn
import numpy as np

mismatchCount = 0
failedTensors = 0
for t in range(5):
    tensor = torch.randn(1000,dtype=torch.float16,device="cuda")
    sm_32d = nn.functional.softmax(tensor, dim=-1, dtype=torch.float32)
    sm_def = nn.functional.softmax(tensor, dim=-1)
    failFlag = 0
    for x in range(1000):
        if sm_32d[x].item() != sm_def[x].item():
            print("Mismatch found. Expected " + str(sm_32d[x].item()) +
                  " in default, got " + str(sm_def[x].item()))
            mismatchCount = mismatchCount + 1
            failFlag = 1
    if failFlag == 1:
        failedTensors = failedTensors + 1
        print("Tensor " + str(t) + " failed.")
    else:
        print("Tensor " + str(t) + " passed.")
    print("Current mismatch count is " + str(mismatchCount))
```

Based on the output, it seems every value in each tensor differs.

Here's a small portion of the output this generated for me:

```
Mismatch found. Expected 0.0004018559120595455 in default, got 0.0004019737243652344
Mismatch found. Expected 0.000258067884715274 in default, got 0.0002579689025878906
Mismatch found. Expected 0.001631039078347385 in default, got 0.0016307830810546875
Mismatch found. Expected 0.0002674324787221849 in default, got 0.0002675056457519531
Mismatch found. Expected 0.0008042750996537507 in default, got 0.0008044242858886719
Mismatch found. Expected 0.0004807825025636703 in default, got 0.00048089027404785156
Mismatch found. Expected 0.00018524506594985723 in default, got 0.00018525123596191406
Mismatch found. Expected 0.0007350771920755506 in default, got 0.0007352828979492188
Mismatch found. Expected 0.0008412283495999873 in default, got 0.0008411407470703125
Mismatch found. Expected 0.00029043699032627046 in default, got 0.0002903938293457031
```

Also of note: the dtype of the default dtype softmax tensor is `float16` while the forced to float32 softmax tensor is `float32`.

Is there anything I missed here or anything I am misunderstanding about this issue?

@Tanman2001 Your script has missing a downcast :) `sm_32d` is a `fp32` tensor, while `sm_def` is a `fp16` tensor. If the hypothesis of the internal computations being in `fp32` is correct, then downcasting `sm_32d` after softmax would make it equal to `sm_def` -- which is NOT the case!

```py
import time
import torch
from tqdm import tqdm

DTYPE = torch.float16
NUM_RUNS = 100000

mismatches = 0
avg_diff = 0
max_diff = 0
time_16bit = 0
time_32bit = 0
for t in tqdm(range(NUM_RUNS)):
    tensor_16bit = torch.randn(32000, dtype=DTYPE, device="cuda")

    time_32bit -= time.time()
    softmax_upcast = torch.nn.functional.softmax(tensor_16bit, dim=-1, dtype=torch.float32).to(DTYPE)
    time_32bit += time.time()

    time_16bit -= time.time()
    softmax = torch.nn.functional.softmax(tensor_16bit, dim=-1)
    time_16bit += time.time()

    if not torch.allclose(softmax, softmax_upcast):
        mismatches += 1
        avg_diff += torch.mean(torch.abs(softmax - softmax_upcast)).item()
        max_diff = max(max_diff, torch.max(torch.abs(softmax - softmax_upcast)).item())

print(f"Mismatches: {mismatches}")
print(f"Avg diff: {avg_diff/NUM_RUNS:2e}")
print(f"Max diff: {max_diff:2e}")
print(f"Time 16bit: {time_16bit/NUM_RUNS:2e} seconds")
print(f"Time 32bit: {time_32bit/NUM_RUNS:2e} seconds")
```

On my end I get

```
Mismatches: 24706
Avg diff: 0.000000e+00
Max diff: 9.536743e-07
Time 16bit: 5.357883e-06 seconds
Time 32bit: 1.586733e-05 seconds
```

The difference is so small that I don't think it's worth the change.

Is there any follow up to this that needs to be done? Should there be another issue opened to modify any softmax calls based on this information?

@Tanman2001 there is no need to change softmax, based on the information we found :)