Text-to-speech data collator exhibits weird batching behavior with Seq2SeqTrainer
Author: GinUTE
Pull request: no

### System Info

- transformers version: 4.37.0.dev0
- platform: Linux-6.1.58+-x86_64-with-glibc2.35 (Colaboratory free accelerated runtime)
- python version: 3.10.12

### Who can help?

_No response_

### Information

- [X] The official example scripts
- [ ] 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

I am currently fine-tuning SpeechT5 on Vietnamese TTS. I followed the official fine-tuning guide [here](https://colab.research.google.com/drive/1i7I5pzBcU3WDFarDnzweIj4-sVVoIUFJ).

The only difference I made is that I changed the tokenizer wrapped in SpeechT5Processor with my own Vietnamese SentencePiece character-level tokenizer. I made sure to add the same special tokens in the original tokenizer, and it is working as expected. I used the following code snippet:  
```
processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
tokenizer = SpeechT5Tokenizer("spm-char.model")
processor.tokenizer = tokenizer
model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
model.resize_token_embeddings(new_num_tokens=len(tokenizer), pad_to_multiple_of=8)
```

The issue arises when I got to the training phase at `trainer.train()`. It throws the following error:  
`Sizes of tensors must match except in dimension 2. Expected size 16 but got size 256 for tensor number 1 in the list.`

I found that the error changes according to batch size. Specifically, the second sentence always throws:  
`Expect size <batch size> but got size <batch size to the power of 2> for tensor number 1 in the list.`

Batch size other than 1 will throw such an error.

I made no change to the original data collator, here is the code snippet:  
```
@dataclass
class TTSDataCollatorWithPadding:
    processor: Any

    def __call__(
        self, features: List[Dict[str, Union[List[int], torch.Tensor]]]
    ) -> Dict[str, torch.Tensor]:
        input_ids = [{"input_ids": feature["input_ids"]} for feature in features]
        label_features = [{"input_values": feature["labels"]} for feature in features]
        speaker_features = [feature["speaker_embeddings"] for feature in features]

        batch = processor.pad(
            input_ids=input_ids, labels=label_features, return_tensors="pt"
        )
        batch["labels"] = batch["labels"].masked_fill(
            batch.decoder_attention_mask.unsqueeze(-1).ne(1), -100
        )
        del batch["decoder_attention_mask"]

        if model.config.reduction_factor > 1:
            target_lengths = torch.tensor(
                [len(feature["input_values"]) for feature in label_features]
            )
            target_lengths = target_lengths.new(
                [
                    length - length % model.config.reduction_factor
                    for length in target_lengths
                ]
            )
            max_length = max(target_lengths)
            batch["labels"] = batch["labels"][:, :max_length]
        batch["speaker_embeddings"] = torch.tensor(speaker_features)
        return batch


data_collator = TTSDataCollatorWithPadding(processor=processor)
```

I checked the batch returned by the data collator with 16 examples and it seems to check out:  
```
{'input_ids': torch.Size([16, 188]),
 'attention_mask': torch.Size([16, 188]),
 'labels': torch.Size([16, 628, 80]),
 'speaker_embeddings': torch.Size([16, 512])}
```

I suspect it must be something to do with the DataLoader, or something else obvious that I just cannot wrap my head around. Any help is appreciated.

### Expected behavior

The fine-tuning should proceed as per usual. I fine-tuned SpeechT5 on Vietnamese TTS once before but not with a custom tokenizer.

Comments:
cc @sanchit-gandhi @ylacombe 

Hello again, I wish to provide more information.

First, I followed the stack trace and landed on this specific line throwing the error:  
```
[/usr/local/lib/python3.10/dist-packages/transformers/models/speecht5/modeling_speecht5.py](https://localhost:8080/#) in forward(self, input_values, speaker_embeddings)
    699             speaker_embeddings = speaker_embeddings.expand(-1, inputs_embeds.size(1), -1)
    700             speaker_embeddings = speaker_embeddings.repeat(inputs_embeds.size(0), 1, 1)
--> 701             inputs_embeds = torch.cat([inputs_embeds, speaker_embeddings], dim=-1)
    702             inputs_embeds = nn.functional.relu(self.speaker_embeds_layer(inputs_embeds))
    703 

RuntimeError: Sizes of tensors must match except in dimension 2. Expected size 16 but got size 256 for tensor number 1 in the list.
```

It seems the error is thrown by the `speaker_embeddings` within the batch. I tested this by not including it in the batch returned by the data collator, and the training proceeds as per usual. As far as I am concerned, however, `speaker_embeddings` is a required component, both during fine-tuning and inference.

I did triple-check the shape of the `speaker_embeddings` tensors within the batch returned by the data collator, and the size is still the same as what I reported in my post.

What should I do now? Any help is much appreciated.

I have the same error. Any solution?

> Hello again, I wish to provide more information.
> 
> First, I followed the stack trace and landed on this specific line throwing the error:
> 
> ```
> [/usr/local/lib/python3.10/dist-packages/transformers/models/speecht5/modeling_speecht5.py](https://localhost:8080/#) in forward(self, input_values, speaker_embeddings)
>     699             speaker_embeddings = speaker_embeddings.expand(-1, inputs_embeds.size(1), -1)
>     700             speaker_embeddings = speaker_embeddings.repeat(inputs_embeds.size(0), 1, 1)
> --> 701             inputs_embeds = torch.cat([inputs_embeds, speaker_embeddings], dim=-1)
>     702             inputs_embeds = nn.functional.relu(self.speaker_embeds_layer(inputs_embeds))
>     703 
> 
> RuntimeError: Sizes of tensors must match except in dimension 2. Expected size 16 but got size 256 for tensor number 1 in the list.
> ```
> 
> It seems the error is thrown by the `speaker_embeddings` within the batch. I tested this by not including it in the batch returned by the data collator, and the training proceeds as per usual. As far as I am concerned, however, `speaker_embeddings` is a required component, both during fine-tuning and inference.
> 
> I did triple-check the shape of the `speaker_embeddings` tensors within the batch returned by the data collator, and the size is still the same as what I reported in my post.
> 
> What should I do now? Any help is much appreciated.

Have you solved your problem?

Hey @GinUTE and @yasamanhbn, thanks for flagging this issue!

@GinUTE did you manage to solve this ? If not, could you send a snippet to reproduce the issue (e.g with a dataset on the hub) ?



Dear Friends,

I have the same problem, and I am fine-tuning the Amharic language.

I can't handle the error when I try to run "trainer.train()."

18 frames
/usr/local/lib/python3.10/dist-packages/transformers/models/speecht5/modeling_speecht5.py in forward(self, input_values, speaker_embeddings)
    699             speaker_embeddings = speaker_embeddings.expand(-1, inputs_embeds.size(1), -1)
    700             speaker_embeddings = speaker_embeddings.repeat(inputs_embeds.size(0), 1, 1)
--> 701             inputs_embeds = torch.cat([inputs_embeds, speaker_embeddings], dim=-1)
    702             inputs_embeds = nn.functional.relu(self.speaker_embeds_layer(inputs_embeds))
    703 

RuntimeError: Sizes of tensors must match except in dimension 2. Expected size 16 but got size 256 for tensor number 1 in the list.

Any help is much appreciated.
Thank you


My apologies, I forgot to leave a comment detailing how I solve the issue after I closed it. To help with understanding the issue, I am devising a short notebook to reproduce the error, the original is a mess.

Following my comment on Dec 26, I knew that the problem was with the speaker embedding batch, but I did not understand why. Thus, I inspected the two code cells involving speaker embedding extraction and batching: my batch-processing function (to use with dataset mapping) and data collator.

I only made modifications to the batch-processing function though because it was more likely to be faulty. The two changes I made are:
- Instead of directly extracting speaker embeddings from audio arrays, I used `torchaudio` to load the audio signals from their paths first
- I converted the labels returned by the processor within the batch from numpy arrays to lists

Honestly, I think the changes are orthogonal to the issue though. I was just blindly changing where I think the issue originated. I will make another comment when I can reproduce the error.

I tried my best to clean up the original inferior notebook. I was able to reproduce the error. I managed to pinpoint the culprit, which is not the batch-processing function or data collator. The issue was, in fact, how I installed the `transformers` package.

I installed `transformers` as follows:
```
!pip3 install git+https://github.com/huggingface/transformers.git # to install transformers
!pip3 install transformers[torch] # to install accelerate>=0.21.0
```

The second `pip3 install` was because of the import error thrown when instantiating `Seq2SeqTrainingArguments` without it:
```
Using the `Trainer` with `PyTorch` requires `accelerate>=0.21.0`: Please run `pip install transformers[torch]` or `pip install accelerate -U`
```

This environment setup will somehow break the trainer and produce the batch size issue. I confirmed this by installing `transformers` using only ```!pip3 install transformers[torch]```, and the error does not persist.

> I tried my best to clean up the original inferior notebook. I was able to reproduce the error [here](https://colab.research.google.com/drive/1I1mMkBTW-Og-ulTJ2XnHAzvYu74qUOsh). You will find the code cells to batch-process the dataset in the first section. I pushed the processed dataset onto my Hugging Face Hub though, so you only need to run the second code section to reproduce the error.
> 
> You might find the only difference between the working and inferior notebooks lies within the batch-processing function.
> 
> ### Edit
> I managed to pinpoint the culprit, which is (unsurprisingly) not the batch-processing function or data collator. The issue was, in fact, how I installed the `transformers` package.
> 
> In the inferior notebook, I installed `transformers` as follows:
> 
> ```
> !pip3 install git+https://github.com/huggingface/transformers.git # to install transformers
> !pip3 install transformers[torch] # to install accelerate>=0.21.0
> ```
> 
> The second `pip3 install` is because of the import error thrown when instantiating `Seq2SeqTrainingArguments` without it:
> 
> ```
> Using the `Trainer` with `PyTorch` requires `accelerate>=0.21.0`: Please run `pip install transformers[torch]` or `pip install accelerate -U`
> ```
> 
> This environment setup will somehow break the trainer and produce the batch size issue. I confirmed this by installing `transformers` using only `!pip3 install transformers[torch]`, and the error does not persist. This behavior is also evidenced by my working fine-tuning notebook, in which I installed `transformers` by the second method.

Thanks a lot. It solves my problem.