[performance/precision] adding `jit.script` to activation functions
Author: stas00
Pull request: no
Labels: name: WIP

# 🚀 Feature request

### Switch our activation functions to use `@torch.jit.script`

Over at BigScience we have been trying to figure out mismatches between Megatron-LM and HF Transformers when it comes to inference under fp16. There are several mismatches, this one discusses activation functions. And proposes to improve HF's models based on that.

So Megatron uses `@torch.jit.script` for its activation functions, which leads to 2 things:

1. faster performance
2. more correct math under fp16 or amp/fp16. Quoting @ngimel:

> ... that’s due to fusion. Fuser does intermediate operations in fp32, and thus produces more accurate results than simple function that truncates each intermediate to half.

(I need to double check on amp/fp16 - I'm making an assumption here)

So perhaps we should switch our activation functions to use `@torch.jit.script` too?

Caveats: 

1. it appears that when using `@torch.jit.script` one may have to write out the `bwd` part explicitly, see:

https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/megatron/model/fused_bias_gelu.py#L27-L56

2. This will change the results slightly (back-compat OK?) but it should produce more correct results!

You can see how the 2 functions diverge:

```
import torch
import random

seed = 42
random.seed(seed)    # python RNG
torch.manual_seed(seed)          # cpu + cuda
torch.cuda.manual_seed_all(seed) # multi-gpu
torch.backends.cudnn.enabled = True

width = 128
input = torch.rand((1,5,width*4)).cuda().half()

@torch.jit.script
def gelu_megatron_fwd_jit(x):
    return  x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))

def gelu_megatron_fwd(x):
    return  x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
    
output = gelu_megatron_fwd(input)
output_jit = gelu_megatron_fwd_jit(input)
# have to run 2nd time for jit to kick in!
output_jit = gelu_megatron_fwd_jit(input)

torch.testing.assert_equal(output, output_jit, check_stride=False)    
```

gives:

```
AssertionError: Tensors are not equal!

Mismatched elements: 800 / 2560 (31.2%)
Greatest absolute difference: 0.00048828125 at (0, 0, 1)
Greatest relative difference: 0.0009828009828009828 at (0, 0, 2)
```

@patrickvonplaten, @patil-suraj, @sgugger, @LysandreJik 

Comments:
For backward compatibility reasons, I would put this as opt-in only: we can add a new key in the activation function for the jitted version that some models can opt-in to use.
This would also make it easier for the support of older torch versions.

Thanks for the write-up! The difference is very minimal but it could still make an unexpected difference. We had the recent change in the way that GeLU is computed with torch (which is changing between torch 1.9 and torch 1.10) and even with a 1e-8 difference it affected us and our tests, so I would expect a 1e-4 difference to affect some downstream users. 

I tend to agree with Sylvain regarding the activation function: It's a configuration argument that is supported on most models (see BERT for example)

https://github.com/huggingface/transformers/blob/7604557e4470822754c5658a19d81aa2ae7de934/src/transformers/models/bert/configuration_bert.py#L78-L80

which would make it simple to switch while keeping the current behavior identical. WDYT?

Thank you for your feedback, Sylvain and Lysandre

So practically, do you propose to replicate the activation functions verbatim, with the only change of adding `@torch.jit.script` and adding a suffix `_jit` to those?

So there will be `gelu_fast` and `gelu_fast_jit`?

note that under fp32 the two will be identical in their outputs, but the jit one will be slightly faster.

Yes, that's exactly what I was suggesting!