Add support for ONNX-TensorRT conversion for GPT-J6B (and possible bug in rotary embedding)
Author: tomerip
Pull request: no

### Who can help
@patil-suraj 

## Information

Model I am using: GPT-J

The problem arises when using:
* [x] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

## Description
I opened this issue for two reasons:
1. This is not strictly a bug report, rather a change that enables converting this model to ONNX and then parsing it using the current TensorRT ONNX parser.
2. Possible implementation bug in GPT-J.

## Details
1. When exporting GPT-J to ONNX using the latest version (v4.16.2), one of the ops that is exported is [SplitToSequence](https://github.com/onnx/onnx/blob/main/docs/Operators.md#SplitToSequence) (along with more Sequence* ops) that is currently not supported in the [TensorRT ONNX parser](https://github.com/onnx/onnx-tensorrt/blob/master/docs/operators.md).
This is entirely due to just 1 line of code that uses `torch.repeat_interleave`. ([relevant line](https://github.com/huggingface/transformers/blob/52d2e6f6e904ef9b75c78716ce77b98196ed837a/src/transformers/models/gptj/modeling_gptj.py#L67))
```
sin, cos = map(lambda t: t[None, offset : x.shape[1] + offset, None, :].repeat_interleave(2, 3), sincos)
```
By replacing `lambda t` with this:
```
lambda t: t.view(-1, 1).repeat(1, 2).view(seq_len, -1)[None, offset : x.shape[1] + offset, None, :]
```
we get the exact same output tensors but now exporting to ONNX doesn't include any Sequence* ops, and TensorRT can parse it successfully.
The suggested function is even faster, although probably not critical in this huge model (benched only on CPU):
```
original: 106 µs ± 20.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
suggested: 32.4 µs ± 6.55 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
```

2. I was following the implementation in EleutherAI for rotary positional embeddings and I'm trying to understand if this is a bug or I'm simply missing something (would love an explanation if you can spare the time) but there (EleutherAI) they implement this function (rotary positional embedding) using `torch.cat` instead of `torch.repeat_interleave`, as can be seen [here](https://github.com/EleutherAI/gpt-neox/blob/b30afd1d0a1d06220be9b5f2c9c9c1523defba96/megatron/model/positional_embeddings.py#L41).

If I'm not missing something, the EleutherAI version transforms a tensor from
```
[[1,2,3],
 [4,5,6]]
```
to 
```
[[1,2,3,1,2,3],
 [4,5,6,4,5,6]]
```
and HF version (using repeat_interleave):
```
[[1,2,3],
 [4,5,6]]
```
to 
```
[[1,1,2,2,3,3],
 [4,4,5,5,6,6]]
```
Can anyone confirm the current implementation is indeed correct? Because otherwise `cat` and `repeat_interleave` are very different, and the rest of the implementation doesn't take it into account.

Comments:
Maybe also of interest to @lewtun @michaelbenayoun 

Hey @tomerip thank you for this very detailed and informative feedback! 

Regarding point (1), I agree that it would be nice to have a work around for the absent `SplitToSequence` op in TensorRT. 

To proceed, I think we'd need to try the following:

1. Validate that the proposed change doesn't have a negative impact on fine-tuning / vanilla inference (i.e. on model metrics like accuracy etc). GPT-J is a popular model, so we need to be sure that we don't introduce a subtle error.
2. Implement an ONNX config for this architecture (see [here](https://huggingface.co/docs/transformers/serialization#exporting-a-model-for-an-unsupported-architecture) for a guide)
3. Validate that the proposed change to the modeling file works for the supported `features` (i.e. tasks) defined by step (2)
4. [Nice to have] Validate that the proposed change also works for ONNX Runtime. They support the [`SplitToSequence` op](https://github.com/microsoft/onnxruntime/blob/df841ee87d0e3c9b814dabaadb4f0490e374a3e6/docs/OperatorKernels.md), so my guess is that the proposed change would also work (although I haven't checked)

I can help out with steps 2-4, but would like to first know what the impact of step (1) is. Regarding your point (2), I'll defer to @patil-suraj who is the expert on the GPT-J models 😃 

Hi @lewtun, thanks for replying.
I agree that the first step should be making sure the change keeps everything the same.
Since the the new `lambda_t` function returns _exactly_ the same outputs, and the underlying ops are just repeats I don't think the grads will be different. In terms of performance, in my original post we can see it's even slightly faster (although probably negligible anyway).

In this short code we can see that it returns the same outputs:
```
import torch
import random

random.seed(42)

def fixed_pos_embedding(x, seq_dim=1, seq_len=None):  # no changes in this function
    dim = x.shape[-1]
    if seq_len is None:
        seq_len = x.shape[seq_dim]
    inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2) / dim))
    sinusoid_inp = torch.einsum("i , j -> i j", torch.arange(seq_len), inv_freq).to(x.device).float()
    return torch.sin(sinusoid_inp), torch.cos(sinusoid_inp)

def test_lambda_t(batch_size, num_heads, seq_len, head_features, num_blocks, offset):
    
    x1 = torch.randn(batch_size, num_heads, seq_len, head_features)  # (batch, head, seq_length, head_features)
    x2 = torch.randn(batch_size, num_blocks, num_heads, seq_len, head_features)  # (batch, blocks, head, block_length, head_features)
    
    for x in [x1, x2]:
        
        sincos = fixed_pos_embedding(x, 1, seq_len=seq_len)
        
        lambda_old = lambda t: t[None, offset : x.shape[1] + offset, None, :].repeat_interleave(2, 3)
        lambda_new = lambda t: t.view(-1, 1).repeat(1, 2).view(seq_len, -1)[None, offset : x.shape[1] + offset, None, :]
        
        sin_old, cos_old = map(lambda_old, sincos)
        sin_new, cos_new = map(lambda_new, sincos)
        
        if ((sin_old == sin_new).all() and (cos_old == cos_new).all()) == False:
            return False
        
    return True

def test_lambda_t_with_params(num_iter=10):
    for _ in range(num_iter):
        params = {
            'batch_size': random.randint(1, 8),
            'num_heads': random.randint(4, 16),
            'seq_len': random.randint(32, 2048),
            'head_features': random.randint(64, 512),
            'num_blocks': random.randint(4, 12),
            'offset': random.randint(0, 4),
        }
        if not test_lambda_t(**params):
            return False
    return True

test_lambda_t_with_params()
```
Regarding steps 2-3, actually since I hit other issues with TensorRT and GPT-J at the moment, I moved to use DeepSpeed (that currently also has some other problems with GPT-J :D), so I will probably not pursue the ONNX path soon, although in the future I intend to.

Still if @patil-suraj can comment on point (2) it would be great!

Will look into this and let you know tomorrow.

Hey @tomerip thanks a lot for the detailed explanation and for showing that the suggested change produces identical outputs! 

If @patil-suraj agrees, would you like to open a PR to implement this? I can then take care of the ONNX export :)

Hi @lewtun,
Sounds great, I'd be happy to open a PR.

Hey @tomerip ! Thanks a lot for this, feel free to open a PR!

> Can anyone confirm the current implementation is indeed correct

It is correct. In GPTNeoX it is different because it uses `[seq, batch, heads, hdim]` format for `qkv` tensors but in transformers it is `[batch, seq, heads, hdim]`

Hey @patil-suraj,
Should we keep this issue open so @lewtun can take care of the ONNX export? i.e. steps 2-4 in his comment above.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

The GPT-J export was added by a community member in https://github.com/huggingface/transformers/pull/16274