Repository: huggingface/transformers (Python)
Task: Bug Fix

Issue:
RuntimeError: The size of tensor a (16) must match the size of tensor b (16000) at non-singleton dimension 2
### 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.


Discussion:
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 

Patch:
diff --git a/src/transformers/models/mega/modeling_mega.py b/src/transformers/models/mega/modeling_mega.py
--- a/src/transformers/models/mega/modeling_mega.py
+++ b/src/transformers/models/mega/modeling_mega.py
@@ -1542,6 +1542,9 @@ def forward(
         else:
             raise ValueError("You have to specify either input_ids or inputs_embeds")
 
+        if self.config.use_chunking:
+            input_shape = torch.tensor([input_shape[0], self.config.chunk_size])
+
         batch_size, sequence_length = input_shape
 
         if self.config.use_chunking and (sequence_length > self.config.chunk_size):


Test patch:
diff --git a/tests/models/mega/test_modeling_mega.py b/tests/models/mega/test_modeling_mega.py
--- a/tests/models/mega/test_modeling_mega.py
+++ b/tests/models/mega/test_modeling_mega.py
@@ -313,6 +313,34 @@ def create_and_check_decoder_model_past_large_inputs(
         # test that outputs are equal for slice
         self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))
 
+    def create_and_check_decoder_model_with_chunking(
+        self,
+        config,
+        input_ids,
+        token_type_ids,
+        input_mask,
+        sequence_labels,
+        token_labels,
+        choice_labels,
+        encoder_hidden_states,
+        encoder_attention_mask,
+    ):
+        config.use_chunking = True
+        config.output_attentions = True
+        config.attention_activation = "laplace"
+        config.chunk_size = input_ids.size(1) * 2
+
+        model = MegaForCausalLM(config).to(torch_device).eval()
+
+        input_ids = input_ids.repeat(1, 8)
+        # multiply the sequence length by 8 since we repeat the same ids 8 times in input_ids
+        input_mask = random_attention_mask([self.batch_size, self.seq_length * 8])
+
+        result = model(input_ids, attention_mask=input_mask)
+
+        # test if the sequence length of attentions is same provided chunk_size
+        self.parent.assertEqual(result["attentions"][0].shape[-1], config.chunk_size)
+
     def create_and_check_for_masked_lm(
         self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
     ):
@@ -547,6 +575,10 @@ def test_decoder_model_past_with_large_inputs(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_decoder_model_past_large_inputs(*config_and_inputs)
 
+    def test_decoder_model_with_chunking(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs_for_decoder()
+        self.model_tester.create_and_check_decoder_model_with_chunking(*config_and_inputs)
+
     def test_for_masked_lm(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*config_and_inputs)