Add speecht5 batch generation and fix wrong attention mask when padding
Author: Spycsh
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes [# (25908)](https://github.com/huggingface/transformers/issues/25908)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerz and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
Hi @Spycsh, thanks for this PR ! I'll look at it ASAP

For future reference, here is your [comment](https://github.com/huggingface/transformers/issues/25908#issuecomment-1704922386) explaining your code choices.

Hi @sanchit-gandhi @ylacombe , I have made another commit, which includes:

- [x] add an optional parameter `attention_mask` to the TTS and SpeechToSpeech generate API, where users can decide whether to pass that
- [x] add batched generation capabilities, and the inputs are not limited to one single text, but batched texts
- [ ] recheck the return type

Specifically, different texts in the batch will need different time of while loops for generation to meet the probability threshold, so my implementation is to use a `result_spectrogram` dict to hold the ones that have already met stop thresholds.

However, the return type are changed. That's what I want to recheck with you. For example, I pass batched texts like 


```python
texts = ["I have a dream, do you?", "He likes drinking coffee.", "Hello!"]
inputs = processor(text=texts, padding='max_length', max_length=128, return_tensors="pt")
spectrograms = model.generate_speech(inputs["input_ids"], speaker_embedding, inputs['attention_mask'])
```

What is returned now should be three spectrograms corresponding to three texts, not a single one. This will change the usage in the document and maybe conflict some legacy code in some users' repo. Should I make it compatible to the old usage (e.g. check if len(inputs) ==1 then fallback to old return type)?

I write a quick test here https://github.com/Spycsh/minimal-speecht5-pad-bug/blob/main/main_batch_inputs.py , which should validate the latest implementation. 

Welcome to any suggestions :)

Stacking the outputs into a torch tensor is preferable! We should try and keep the modelling code take the format `{tensors in} -> {tensors out}`. This is the API familiar to the user, and makes the modelling code compatible with torch ops like `torch.compile`. Having ragged lists makes this much more difficult, and means the user can't move the outputs back to the CPU easily with the torch `.cpu()` method

Thanks @ylacombe @sanchit-gandhi for detailed suggestions. I'm working on batching the postnet inputs, return types (with spectrogram and waveform lengths) and the modification will be consistent to old APIs (Will check with the previous tests). Will get back to you soon.

Hi @ylacombe , @sanchit-gandhi , I've just made another commit. It includes:

- [x] batched inputs for the postnet, not the original one-by-one logics
- [x]  consistent code to old user API when the batch_size is 1
- [x] redesign the return type when the batch_size>1 (explained below)
- [x] add another test for batch generation

The return values now should now be (spectrograms, spectrogram_lengths) or (waveforms, waveform_lengths) if using vocoder, when the first element of the tuple is the stacked spectrograms/waveforms values, and the second element is a list of spectrogram/waveform lengths.

I pass the tests `RUN_SLOW=1 pytest tests/models/speecht5/test_modeling_speecht5.py::SpeechT5ForTextToSpeechIntegrationTests`.

Welcome to any suggestions:)

Hi @ylacombe , thanks for suggestions, I've added the test, cross attention and docstring. Welcome to further advices!

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_25943). All of your documentation changes will be reflected on that endpoint.

Hi @ylacombe , is there anything else I need to add or modify to this PR?

Hi @ylacombe , I've adapted most of the code based on your reviews. Welcome to any other questions! I am still confused about the random inconsistency between batch generation and single generation.

One quick reproduction is to edit the `test_modeling_speecht5.py`
1. comment `self.assertEqual(spectrograms.shape, (3, 266, model.config.num_mel_bins))` (i.e. Comment Line 1050, This is irrelevant in the following reproduction)
2. run `RUN_SLOW=1 pytest tests/models/speecht5/test_modeling_speecht5.py::SpeechT5ForTextToSpeechIntegrationTests::test_batch_generation`

You will get everything passed. First sentence spectrogram consistent w/wo batch generation. But when it comes to the second or third input text, it will fail. Also, by simply changing the first sentence "mister quilter is the apostle of the middle classes and we are glad to welcome his gospel" to "mister quilter is the apostle of the middle classes and we are glad to welcome his ", it also fails. Change to "mister quilter is the apostle of the middle classes and we" will pass. It seems quite random and confusing.

@ylacombe I normally use breakpoint() simply in Python. Do you have any other recommendations? :)

Will root cause the inconsistency problem and get back to you soon.

Hi @Spycsh, thanks again for your involvement and motivation. The PR is getting more extensive than expected but IMO you're doing an excellent job!

There are two possible ways to go here:
1. either refocus on the primary focus, which was to fix the attention mask, and fix the batching in another PR
2. Or go the extra mile and fix it now!

In any case, I've encountered similar batching problems in #25693. 
Usually the cause of batching issues in audio models can come from 3 different origins:
1. Wrong attention mask - in practice here, probably making sure that [_get_feature_vector_attention_mask](https://github.com/huggingface/transformers/blob/e3a4bd2bee212a2d0fd9f03b27fe7bfc1debe42d/src/transformers/models/speecht5/modeling_speecht5.py#L612) has the right behavior
2. Wrong masking in convolutional layers,  i.e making sure that it is correctly padded with 0. Here convolutional layers are in SpeechT5FeatureEncoder and in the hifigan.
3. Errors when passing from non-batching to batching: it might be errors in transpose, mean or similar operations

I'll also be extra-careful to the hifigan batching behavior. Notably in two aspects:
1. before passing to the vocoder, the inputs should be padded with 0.
2. the normalization [here](https://github.com/huggingface/transformers/blob/e3a4bd2bee212a2d0fd9f03b27fe7bfc1debe42d/src/transformers/models/speecht5/modeling_speecht5.py#L3273C23-L3274) doesn't care about batching, it might not work when the input is padded with 0.

In terms of debugging, I advise using [vscode debugging tools](https://code.visualstudio.com/docs/python/debugging) or something similar, especially since it allows to easily add breakpoints and logpoints wherever you feel you need it !

This debugging can take some time, and I'd totally understand if you thought it is too much work atm. 
In that case, what do you think of point 1. at the beginning of this comment? in other words, open a PR focusing only on the problem you were initially trying to solve. You already have done the job here so it's only a matter of copy and paste. 
It might be a good thing to do this work step by step and in any case the enormous amount of work you've put in here won't be in vain as it's a very good start to proper batching functionality!

Hi @ylacombe , thanks for the information and guidance. Firstly in this answer, I will talk about  the inconsistency.

Regarding to the inconsistency, I narrow down a lot to find the first occurrence w/wo batching that has a difference. I finally lock on line [963](https://github.com/huggingface/transformers/blob/e3a4bd2bee212a2d0fd9f03b27fe7bfc1debe42d/src/transformers/models/speecht5/modeling_speecht5.py#L963). Here the hidden_states (input) w/wo batching are exactly the same. However the key_states (output) are different. This phenomenon also exists for query states and value states. An evidence is here:

* hidden_states w/wo batching are **the same** (save hidden_states[0].numpy() to compare)
<img width="1204" alt="image" src="https://github.com/huggingface/transformers/assets/39623753/343be0b5-858a-4991-8bf8-e5c3081d4a39">

* key_states w/wo batching are **different** (save key_states[0] to compare)
<img width="1168" alt="image" src="https://github.com/huggingface/transformers/assets/39623753/396459aa-5b41-494a-b1a5-99a1cf983aec">

You can see very slight difference on the second image.

And then I narrowed down further and lock on `self.k_proj(hidden_states)`, which is also different w/wo batching. `hidden_states`, like I said before, are exactly same. So I checked and found that the weight of the Linear layers (k_proj, v_proj...) are exactly the same. Then the only reason is that behind the `nn.Linear` forward, the batching/not batching goes to different kinds of matrix multiplication.

I am not very familiar with BLAS GEMM operations, but based on the observing I have the reason to suspect there is something that lead to the slight difference here. Maybe it is because they invokes different ops like `addmm` or `addbmm` (https://pytorch.org/docs/stable/generated/torch.addbmm.html#torch.addbmm). Maybe it is because the `addmm` has different kinds of computation for faster optimization underneath.

My inference is that these slight differences, during the generation once and again will possibly lead to the different final possibilities to meet the thresholds and affect the spectrogram values and shape, which seems hard for us to write tests that check the consistency with `torch.allclose` at a `atol=1e-4`.

However, I would say that these inconsistencies are not really causing the final audios' quality. I actually listen to the final audios w/wo batching. I cannot feel the difference.

Of course, it may also be the hardware factors. I will check that with the CircleCI's hardware (seems to be skipped because it is a slow test).

Next let us discuss the another problem: the vocoder issue. With batching padded inputs, Vocoder outputs waveforms that are padded with nonzero values, which is hard for us to know which are the useful part or which are not. So I write the following code and I think what I wrote should be robust
```
waveform_lengths = [int(waveforms.size(1) / max(spectrogram_lengths)) * i for i in spectrogram_lengths]
```

This is because I found that the padded waveform lengths (namely `waveforms.size(1)`) should be the integer times of padded spectrograms lengths (namely `max(spectrogram_lengths)`), which is 256 here. By providing the  `waveform_lengths `, users can filter out the effective concrete part of the batched output waveforms.

I actually also listen to the final audios' quality and it works totally like without batching.

I notice that you mention `spectrogram = (spectrogram - self.mean) / self.scale` is not considering batching. I think `self.mean` and `self.scale` are both zeros and ones? Not sure whether it matters.

Welcome to further discussion!

Hi @Spycsh,

Thanks for your efforts here ! Actually I believe that the (really) small numerical differences that you found on [L963](https://github.com/huggingface/transformers/blob/e3a4bd2bee212a2d0fd9f03b27fe7bfc1debe42d/src/transformers/models/speecht5/modeling_speecht5.py#L963) are negligible and should not influence that much the final output! (correct me if I'm wrong)

Maybe you are already doing this, but something that help me a lot to debug is searching the (noticeable) difference by dichotomy!   


Hi @ylacombe , you are right, after a complete re-examining, the small numerical differences that I found in previous comment are not the major factor.

The major reason that behind the scene is this [line](https://github.com/huggingface/transformers/pull/25943/files#diff-a5f5680267f18af447101950cbcadd4559e301efc6d1b700c3e0994d2fcfc28cL727).

When doing dropout with setting `training=True`, torch **randomly zeroes some of the elements of the input tensor with probability p using samples from a Bernoulli distribution. Each channel will be zeroed out independently on every forward call.** (https://pytorch.org/docs/stable/generated/torch.nn.Dropout.html#torch.nn.Dropout). That means during batch inference, the random Bernoulli dropout mask is independently generated, and should be different for each `inputs_embeds` in a batch. However, the random Bernoulli dropout mask should be the same for each `inputs_embeds` when we do N times single run on each text. And that cause the inconsistency.

One example is here:

```
>>> input = torch.randn(3,2)
>>> input
tensor([[-0.9113, -0.9378],
        [ 0.1110,  1.4664],
        [-1.0628,  0.6064]])
>>> set_seed(555)
>>> res1=torch.nn.functional.dropout(input, 0.5, training=True)[0]
>>> set_seed(555)
>>> res2=torch.nn.functional.dropout(input[0], 0.5, training=True)
>>> torch.allclose(res1, res2, atol=1e-4)
True
>>> set_seed(555)
>>> res1=torch.nn.functional.dropout(input, 0.5, training=True)[1]
>>> set_seed(555)
>>> res2=torch.nn.functional.dropout(input[1], 0.5, training=True)
>>> torch.allclose(res1, res2, atol=1e-4)
False
```

We will always get `Dropout(x)[i] != Dropout(x[i])` when i > 0, which is caused by the same reason of inconsistency w/wo batching.

To simply understand that, let's say we have three texts, and we call them x1, x2 and x3, we do a dropout to them for three runs with the same random seed. And also, I think you can agree that I just say dropout is like applying an 1,0 **"mask"** on the origin input, where 1 means the element is kept, and 0 means the element is overwritten as 0.

```
Run1
x1 -> dropout -> y1;  Assume here our mask in this run is 1,0,1,1

Run2
x2 -> dropout -> y2;  Here our mask generated by random seed should also be 1,0,1,1

Run3
x3 -> dropout -> y3; Here our mask generated by random seed should also be 1,0,1,1
```

As you see, for each run, the **"mask"** is the same because the random seed sequence generated from the beginning should be the same.

However, if we do batching on them, the **"mask"** for each one is sequentially generated one by one, but not the same, like the following:

```
Run1
[x1, x2, x3] -> dropout -> [y1, y2, y3]; the mask generated here should be like [1,0,1,1,0,1,0,1,1,0,0,1]
```

The random number generator does not always generate from the beginning like what we do in three single runs, but just keep generating new **"masks"** (different than 1,0,1,1)! And that's why `Dropout(x)[i] != Dropout(x[i])` when i > 0.

Therefore, I create a method called `_consistent_dropout` to create a consistent mask of the first instance and apply it to other inputs in the same batch to fix this inconsistency problem (Just like applying 1,0,1,1 to x2 and x3). As mentioned in the document https://pytorch.org/docs/stable/generated/torch.nn.Dropout.html#torch.nn.Dropout, I scale it by multiplying 1/(1-p). The shape is now correct (Still may have a tiny inconsistency because of the reason in my last answer).

Welcome to further questions!

Hey @Spycsh, nice work on finding out where the problem actually is!

According to your analysis, the problem seems to arise when `training=True`. Did you thus check that you get consistent results using `torch.inference_mode()` which should solve the issue ? Note that we are looking for consistent results when doing inference, it's ok to have difference when `training=True` due to probabilistic behaviour in dropout.

What I would do is to first check if `torch.inference_mode()` resolves our issue. In any case, tests should be run with `inference_mode` to ensure that everything work under the correct setup! Although the `_consistent_dropout` method is elegant, I don't think it is necessary here!

Again, let me know if you have any further questions or if that's not clear enough!





Hi @ylacombe , `torch.inference_mode` does not work here. With using `torch.inference_mode` and setting `training=True`, the shape in the test w/wo batching are not consistent. As I said before, the reason lies in different random processes of the random number generator, which  `inference_mode` cannot control. Here is a hands-on example to show that `inference_mode` not work.

```
import torch
from transformers.trainer_utils import set_seed

input = torch.randn(3,2)

with torch.inference_mode():
    set_seed(555)
    res1=torch.nn.functional.dropout(input, 0.5, training=True)[1]
    set_seed(555)
    res2=torch.nn.functional.dropout(input[1], 0.5, training=True)

print(torch.allclose(res1, res2, atol=1e-4))    # False!
```


Hi @Spycsh, sorry for not being clear enough, and thanks for your response! You're not supposed to test if the results are the same with `training=True`, but when doing `inference_mode` and with `training=False`. You can set `training=False` with `model.eval()`. Does that solve your masking issue?

Hello @ylacombe  still a bit unclear. So should we manually set training=False in modeling_speecht5.py when we do the tests? I suspect both inference mode and model.eval() still cannot yield consistent mask when we explicitly keep training=True in dropout. Will check that.

Hi @ylacombe , I set both model.eval() and with torch.inference_mode(), and it does **NOT** fix the inconsistency. Could you please take a look at the comment [here](https://github.com/huggingface/transformers/blob/main/src/transformers/models/speecht5/modeling_speecht5.py#L721) ? The paper suggests to set `training=True` not only at training, but also at inference! In paper it says "In order to introduce output variation at **inference time**, dropout with probability 0.5 is applied only to layers in the pre-net of the autoregressive decoder".

I think, if my understanding is correct, the methods that you gave me like `model.eval()` and `torch.inference_mode()` are ways to deactivate **nn.Dropout** rather than **nn.functional.dropout** during inference mode. This [example](https://stackoverflow.com/questions/53419474/pytorch-nn-dropout-vs-f-dropout/53452827#53452827) proves that.

That means we cannot use any tricks (e.g. `model.eval()` and `torch.inference_mode()` ) to deactivate `training=True` in our tests, right? So here, from my perspective, are only two choices:

1. we need a test to check consistency, but with a `_consistent_dropout` that I wrote to keep the same bernoulli process for each instance in one batch, which is consistent to batch_size=1. **`_consistent_dropout` not only keep consistent with/without batching, but also introduce output variation at inference time, which have the same effect of setting training=True**! That's why I recommend to do that.
2. we do not need a test to check the consistency, so I just delete `_consistent_dropout` in the modeling_speecht5.py and also delete the consistency check in the test.

Welcome to further advices!

Thanks @ArthurZucker for detailed suggestions! Will get back to you soon:)

Hi @ArthurZucker , I've made some changes and added some comments base on your review. Could you review on those? Thanks.

> Ok sounds good, but we would still return this for a single batch. Sometimes people pad a single input to a multiple of X

Hi @ArthurZucker , padding or not padding to the a single batch will not affect the returned spectrogram's length. If you have only one record in your batch, no matter how long you pad it, the spectrogram you generate will be the same. So here I do not think it is necessary to return the spectrogram's length when batch_size=1, because the spectrogram can be used directly!

> Answered your questions 😉 let's add a `return_spectrogram_length` to make it optional and have similar outputs for batched and non batched!

Hi @ArthurZucker , I do not think a `return_spectrogram_length` boolean parameter is necessary if the usages and returned values are simple and deterministic here. 

Here are all three cases of user inputs:

1. single text
2. single text with padding
3. batched texts with padding

For the first two cases,  users will get their expected spectrograms/waveforms, and they do not need to know anything about the returned lengths. The returned lengths are just identical to the len(spectrogram)/(waveform)

For the third case , users will get the spectrograms/waveforms with returned lengths, because the output spectrograms/waveforms in a batch have different concrete lengths.

I think what you understand is that we can have a `return_spectrogram_length` to distinguish between case 1 and case 2. However the spectrogram lengths is irrelevant to whether an input text is padded or not. Actually, the spectrogram/waveform lengths we return here should only be utility to limit the concrete range in the output spectrograms batch to solve the different lengths of spectrogram when batch_size>1.

I think an extra `return_spectrogram_length` will confuse users and not that necessary here.

What do you think?

> An updated snippet in the readme of speech T5 is also super welcome!
Hi @ArthurZucker . Sure! Could you please tell me where can I edit the README of the SpeechT5?

Hey! 
Basically I don't think we need to go overboard for this. 1,2 and 3 can all be merged into:
1. `return_spectrogram_length = True` (or any other name, but something explicit and documented with what you mentioned, the output might not be the same if padded etc)
2. `return_spectrogram_length = False`
The goal here is just to be consistent with what we return whether we have a batch or not.

> the spectrogram/waveform lengths we return here should only be utility to limit the concrete range in the output spectrograms batch to solve the different lengths of spectrogram when batch_size>1 

I am not sure I completely understand what you mean here, but if it's not something that will always be used, I'd rather we have a simple logic applicable for batched and un batched, and just control the return arg. 

I'll let @sanchit-gandhi make the call both are alright anyway since we are adding the feature 😉 

Hi @ArthurZucker @sanchit-gandhi , I made the change followed @ArthurZucker's request. An example of feed batched inputs should be like following, and also you can review it in the test:

```
        spectrograms, spectrogram_lengths = model.generate_speech(
            input_ids=inputs["input_ids"],
            speaker_embeddings=speaker_embeddings,
            attention_mask=inputs["attention_mask"],
            return_concrete_lengths=True,
        )
```

By setting a `return_concrete_lengths=True`, we can always return the lengths, even if the lengths are unnecessary when batch size==1. However it keeps a uniform form whether batch size=1 or not.

Also, `return_concrete_lengths` is set to False by default for backward compatibility. It will not affect any previous code that call this method.

I here name the parameter as `return_concrete_lengths` instead of `return_spectrogram_lengths` because we need to consider that it is the waveform but not the spectrogram lengths that are returned if a vocoder is given. Therefore, I think `return_concrete_lengths` is a general term.

@sanchit-gandhi , @ArthurZucker could you please review this change? Welcome any suggestions:)

Hi @ArthurZucker @sanchit-gandhi , is this PR still on track?

Looks good to me, I already approved so feel free to merge @sanchit-gandhi if this looks alright with you! Specifically the latest changes

Hi @ylacombe , I've updated the code based on your review, could you please review it?

Hello @sanchit-gandhi, do you think this PR is ready for a merge? Welcome to any suggestions!

@ylacombe @ArthurZucker still need this PR? If there are no more problems could you also consider merging this :)

For sure thanks for your efforts and for this great feature addition! 😉 

Merging! Thanks @Spycsh for the great feature addition and the great work here!