[performance/precision] adding `jit.script` to activation functions

# 🚀 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

Start by locating the Transformers activation-function implementations and compare them with Megatron-LM's fused_bias_gelu.py, especially lines 27-56. Measure fp16 or AMP precision and performance before deciding which functions need scripted implementations; done means the selected activations use the approach consistently without breaking supported behavior.

Written by the indexing model from the issue text.