RuntimeError: The size of tensor a (16) must match the size of tensor b (16000) at non-singleton dimension 2
Author: Tylersuard
Pull request: no
Labels: name: Good Difficult Issue

### System Info

- `transformers` version: 4.30.0.dev0
- Platform: Linux-5.10.147+-x86_64-with-glibc2.31
- Python version: 3.10.11
- Huggingface_hub version: 0.14.1
- Safetensors version: not installed
- PyTorch version (GPU?): 2.0.0+cu118 (False)
- Tensorflow version (GPU?): 2.12.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.6.9 (cpu)
- Jax version: 0.4.8
- JaxLib version: 0.4.7
- Using GPU in script?: YES
- Using distributed or parallel set-up in script?: NO

### Who can help?

@ArthurZucker 

### Information

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

### Tasks

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

### Reproduction

Run this notebook: https://colab.research.google.com/drive/1TFI84P9W4VPhNLgEngxPN57RwzS0C4bG?usp=sharing

### Expected behavior

Expected the model to train successfully.  Instead it gives a tensor mismatch error.

Comments:
Hi @Tylersuard, thanks for reporting this issue. 

So that we can best try and help you, could you update the notebook so that it contains the minimal logic to replicate the error and can be run out-of-the-box? As it stands, there's many blocks with comments; references to loading / processing data we don't have access to;  doesn't currently have the reported error shown but does have many other errors. 

Sorry @amyeroberts , Here is the updated version: https://colab.research.google.com/drive/1TFI84P9W4VPhNLgEngxPN57RwzS0C4bG?usp=sharing

I think you're splitting your input sequence into chunks of length 16: https://github.com/huggingface/transformers/blob/v4.29.1/src/transformers/models/mega/modeling_mega.py#L1063

@OllieBroadhurst That is correct.  As per the documentation (https://huggingface.co/docs/transformers/main/model_doc/mega) , I set the chunk_size equal to 16 and use_chunking to true, and the context length is a multiple of the chunk size.  My problem is not solved.

What I mean is have you tried turning chunking off?

@OllieBroadhurst Thank you for your suggestion.  I would likely run into out-of-memory errors, but I will try it.

Ok I tried it without chunking and I got out-of-memory errors.

This should still be adressed! Mega's forward pass might need some debugging. I can't do this fast, but keeping an eye on it! 

Did not have time to dive into this. Marking as good second issue in case community want to have a go! 

I would like to have a go at this @ArthurZucker!

Sure! 😉 

I ran the notebook provided by @Tylersuard on an A6000 with the following settings:
- With `chunk_size=32`: The RuntimeError still persists (I tried this to see if some other multiple of 16 would produce any different of a result)
- With `use_chunking=False`: In this case, the forward pass appears to work fine, but another error is thrown because of the labels.

 Here is that error:
 ```Traceback (most recent call last):
  File "/root/hf_trial/copy_of_hf_mega_music_for_issue.py", line 166, in <module>
    trainer.train()
  File "/opt/conda/lib/python3.10/site-packages/transformers/trainer.py", line 1555, in train
    return inner_training_loop(
  File "/opt/conda/lib/python3.10/site-packages/transformers/trainer.py", line 1837, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/trainer.py", line 2682, in training_step
    loss = self.compute_loss(model, inputs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/trainer.py", line 2707, in compute_loss
    outputs = model(**inputs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/mega/modeling_mega.py", line 1772, in forward
    lm_loss = loss_fct(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/loss.py", line 1174, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/functional.py", line 3029, in cross_entropy
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
RuntimeError: "nll_loss_forward_reduce_cuda_kernel_2d_index" not implemented for 'Int'
```

Now this error is perhaps out of the scope of this issue so I will proceed to debug the forward pass with `use_chunking=True`

cc @ArthurZucker, @amyeroberts 

Hi @ArthurZucker, I found what was causing the error as well as a 'potential' fix (I am not too sure about it so need some feedback on it!).

## Cause
The error in Line [#866](https://github.com/huggingface/transformers/blob/015f8e110d270a0ad42de4ae5b98198d69eb1964/src/transformers/models/mega/modeling_mega.py#L866) was being caused because the `torch.matmul(query, key.transpose(2, 3))` was being divided by `lengths`, which was created in Line [#854-55](https://github.com/huggingface/transformers/blob/015f8e110d270a0ad42de4ae5b98198d69eb1964/src/transformers/models/mega/modeling_mega.py#L854-L855) from the `causal_mask`. 

When using a `chunk_size`, the `causal_mask` dimensions are of the form `[batch_size, 1, target_length, target_length]` which, after being summed using `causal_mask.sum(dim=-1, keepdim=True)` makes the `lengths` to be of the form: `[batch_size, 1, target_length, 1]`.

In reality, with a provided `chunk_size`, the `causal_mask` should be of the dim: `[batch_size, 1, chunk_size, chunk_size]` and in turn the `lengths` should be of the dim: `[batch_size, 1, chunk_size, 1]` for the aforementioned division to work.

## Potential Fix
I have created a PR that adds a simple condition to check if chunking is enabled, in which case it will make the `input_shape` to be `[batch_size, chunk_size]`.

PS: Apologies for this rather odd-looking comment but there were just too many technicalities to explain here!

No worries and sounds good actually!