Training RWKV is ~10x slower than GPT2 on GPU
Author: ivnle
Pull request: no

### System Info

Training RWKV is ~10x slower than GPT2 on GPU and ~3x slower on CPU. Both use Huggingface's implementations.

cuda version: 12.1
torch version: 2.1.0.dev20230812+cu121
cuda driver: 8902
huggingface version: 4.30.2
operating system: posix
Linux 5.15.0-78-generic
GPU: 1x A6000

I was able to replicate these results after downgrading to Pytorch 2.01. I was also able to replicate it on a 2080ti on another machine.



### Who can help?

@ArthurZucker @younesbelkada 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```
"""
This is a minimal example of training Huggingface's implementations of
RWKV and GPT2. You should observe that RWKV is 10x slower than GPT2 on
GPU.
"""

import os
import platform
import torch
import torch.nn as nn
import torch.optim as optim
import transformers
import time

input_size = 128
seq_len = 32
hidden_size = input_size
num_layers = 2
num_heads = 2
batch_size = 128
num_epochs = 1000

# set seed
torch.manual_seed(0)


def count_paramters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


gp2_config = transformers.GPT2Config(
    n_positions=input_size,
    n_embd=hidden_size,
    n_layer=num_layers,
    n_head=num_heads,
    n_inner=hidden_size * 4,
)

rwkv_config = transformers.RwkvConfig(
    context_length=input_size,
    hidden_size=hidden_size,
    num_hidden_layers=num_layers,
    intermediate_size=hidden_size * 4,
)


def train(model, device="cuda"):
    head = nn.Linear(hidden_size, 1)
    # move to device
    model = model.to(device)
    head = head.to(device)
    model.train()
    head.train()

    # Define loss function and optimizer
    criterion = nn.MSELoss()  # Mean Squared Error loss
    optimizer = optim.SGD(list(model.parameters()) + list(head.parameters()), lr=0.01)

    # Training loop
    start = time.time()
    for epoch in range(num_epochs):
        # Generate random inputs and labels
        inputs = torch.rand(
            batch_size, seq_len, input_size, device=device
        )  # Batch size: 32, Sequence length: 10, Input size: 10
        labels = torch.rand(
            batch_size, seq_len, 1, device=device
        )  # Batch size: 32, Sequence length: 10, Output size: 1

        # Zero the gradients
        optimizer.zero_grad()

        # Forward pass
        outputs = model(inputs_embeds=inputs).last_hidden_state
        outputs = head(outputs)

        # Compute the loss
        loss = criterion(outputs, labels)

        # Backpropagation
        loss.backward()

        # Update weights
        optimizer.step()

        # Print loss every 100 epochs
        if (epoch + 1) % 100 == 0:
            print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")
    end = time.time()
    return end - start


for device in ["cpu", "cuda"]:
    print(f"Training GPT2 on {device}")
    gpt2_model = transformers.GPT2Model(gp2_config)
    print(f"GPT2 parameters: {count_paramters(gpt2_model)}")
    gpt2_train_time = train(gpt2_model, device=device)
    print(f"GPT2 training time: {gpt2_train_time:.2f} seconds\n")

    print(f"Training RWKV on {device}")
    rwkv_model = transformers.RwkvModel(rwkv_config)
    print(f"RWKV parameters: {count_paramters(rwkv_model)}\n")
    rwkv_train_time = train(rwkv_model, device=device)
    print(f"RWKV training time: {rwkv_train_time:.2f} seconds\n")

    # compare speed up
    print(f"RWKV / GPT2: {rwkv_train_time / gpt2_train_time:.2f}x\n")
    print("-" * 80 + "\n")

# print cuda version
print(f"cuda version: {torch.version.cuda}")
# print torch version
print(f"torch version: {torch.__version__}")
# print what cuda driver is being used
print(f"cuda driver: {torch.backends.cudnn.version()}")
# print huggingface version
print(f"huggingface version: {transformers.__version__}")
# print system information like python version, operating system, etc.
print(f"operating system: {os.name}")
print(platform.system(), platform.release())

```

Running this script produced the following
```
Training GPT2 on cpu
GPT2 parameters: 6846080
Epoch [100/1000], Loss: 0.1093
Epoch [200/1000], Loss: 0.0973
Epoch [300/1000], Loss: 0.0898
Epoch [400/1000], Loss: 0.0887
Epoch [500/1000], Loss: 0.0880
Epoch [600/1000], Loss: 0.0838
Epoch [700/1000], Loss: 0.0844
Epoch [800/1000], Loss: 0.0853
Epoch [900/1000], Loss: 0.0841
Epoch [1000/1000], Loss: 0.0846
GPT2 training time: 60.22 seconds

Training RWKV on cpu
RWKV parameters: 6864768

Epoch [100/1000], Loss: 0.1054
Epoch [200/1000], Loss: 0.0845
Epoch [300/1000], Loss: 0.0842
Epoch [400/1000], Loss: 0.0823
Epoch [500/1000], Loss: 0.0850
Epoch [600/1000], Loss: 0.0823
Epoch [700/1000], Loss: 0.0836
Epoch [800/1000], Loss: 0.0812
Epoch [900/1000], Loss: 0.0836
Epoch [1000/1000], Loss: 0.0846
RWKV training time: 174.93 seconds

RWKV / GPT2: 2.91x

--------------------------------------------------------------------------------

Training GPT2 on cuda
GPT2 parameters: 6846080
Epoch [100/1000], Loss: 0.1088
Epoch [200/1000], Loss: 0.0966
Epoch [300/1000], Loss: 0.0899
Epoch [400/1000], Loss: 0.0898
Epoch [500/1000], Loss: 0.0830
Epoch [600/1000], Loss: 0.0832
Epoch [700/1000], Loss: 0.0830
Epoch [800/1000], Loss: 0.0844
Epoch [900/1000], Loss: 0.0858
Epoch [1000/1000], Loss: 0.0873
GPT2 training time: 8.16 seconds

Training RWKV on cuda
RWKV parameters: 6864768

Epoch [100/1000], Loss: 0.1078
Epoch [200/1000], Loss: 0.0864
Epoch [300/1000], Loss: 0.0832
Epoch [400/1000], Loss: 0.0833
Epoch [500/1000], Loss: 0.0831
Epoch [600/1000], Loss: 0.0827
Epoch [700/1000], Loss: 0.0809
Epoch [800/1000], Loss: 0.0827
Epoch [900/1000], Loss: 0.0845
Epoch [1000/1000], Loss: 0.0847
RWKV training time: 77.16 seconds

RWKV / GPT2: 9.45x

--------------------------------------------------------------------------------

cuda version: 12.1
torch version: 2.1.0.dev20230812+cu121
cuda driver: 8902
huggingface version: 4.30.2
operating system: posix
Linux 5.15.0-78-generic
```


### Expected behavior

The RWKV paper claims that "one of the defining characteristics of RWKV is its ability to offer parallelized training and robust scalability, similar to Transformers". This led me to believe that their training times should be comparable. 

Comments:
The problem was that I wasn't using RWKV's custom kernel. `pip install ninja` resolves the slowdown.