[FLAX] Whisper
Author: kamalkraj
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 # (issue)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] 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).
- [ ] 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:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @LysandreJik

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @n1t0, @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->


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

Hi,

I need little clarification about implementing the `FlaxWhisperDecoder` Module.

What would be the best way to pass `past_key_values_length` to the module?
 
Reference in Pytorch implementation. 
https://github.com/huggingface/transformers/blob/main/src/transformers/models/whisper/modeling_whisper.py#L863-L873 

@patrickvonplaten @ydshieh @patil-suraj 

Whisper on TPU will make :fire: colab demos

![Screenshot 2022-10-16 at 8 29 27 PM](https://user-images.githubusercontent.com/17096858/196042609-971f30b8-a96c-4df8-8524-c3c6ef379ae6.png)


Awesome work here! Feel free to ping me for a review once it is ready 😄 

Hi,

I have finished the model and working on the test cases now. 
The pt<->flax equivalence test is failing, even though the `model.generate` produce the exact speech-to-text like the PyTorch model. 

![Screenshot 2022-10-19 at 10 15 39 PM](https://user-images.githubusercontent.com/17096858/196753624-b0941182-401d-443c-a0b2-fdc32ed57112.png)


I have attached steps to reproduce the issue in this notebook - https://colab.research.google.com/drive/1KmO8OBUpHfs1uYA_eSwamQAXnjsdbkRS?usp=sharing

Any pointers will be helpful.

Thanks

@patrickvonplaten @patil-suraj @ydshieh @ArthurZucker 

Hi @kamalkraj First, thank you for this awesome PR!

Regarding the PT/Flax tests, I probably need to improve that PT/Flax equivalence tests to make it (a bit) easier to find out which layers gives the larger difference.

In the meantime, I have to say there is no easy way to debug such issue. We need patience to find out at which layer(s) we have the first large difference (greater than the tolerance) and see what's wrong inside that layer.

This is usually tedious and involving manually debugging process.

Anyway, I can open a PR to make the process (a bit) easier - if you want to wait a bit. But notice that we still need similar process even that PR is merged.

Will try to get #18420 merged so that we can maybe use the `find_pt_fx_differences(pt_outputs, fx_outputs)` function! But in the mean time, you should set `output_hidden_states=True` and check where the lists differ 🤗 

Hi @kamalkraj Actually that test is quite good enough, but we need to change a bit to debug more.

The last 2 commit in this [branch](https://github.com/huggingface/transformers/commits/temp-debug-whisper-flax) could log more information.

If you run the tests like

```bash
RUN_PT_FLAX_CROSS_TESTS=true python3 -m pytest -v tests/models/whisper/test_modeling_flax_whisper.py -k "test_equivalence_pt_to_flax"
```
it logs something
```
max diff. in outputs.logits: 0.0020506680011749268
```
but it doesn't fail the test -> it continues. So far, I got

```bash
E   AssertionError: <class 'list'> != <class 'tuple'> : outputs.decoder_hidden_states: Output types differ between Flax and PyTorch
```
so you will have to look the output type of `decoder_hidden_states` and make sure the type is the same as the PyTorch one.
Continue this process will eventually show you all the difference, and you can get a better idea where to debug in the modeling code.

Also, it seems when running the tests from `tests/models/whisper/test_modeling_whisper.py`, we have some shape issue. This is another thing to debug.

Hopefully this gives you some idea of how we can debug here 🤗 



Thanks, @ydshieh and @ArthurZucker



To make for a more consistent API across models, couldn't we swap out `past_key_values_length` and instead compute `position_ids` to get the current positional embeddings for the decoder? It feels like this would make it easier to fit Whisper in with other finetuning codebases (no need to create custom logic for computing `past_key_values_length` when dealing with Whisper). As the code currently stands, I think it would actually give incorrect outputs when decoding a batch when each element of the batch has different decoder prefix/prompt tokens. Computing position ids from the attention mask would also allow for either left or right padding. 

I have another Flax Whisper implementation with .from_pretrained(..., from_pt=True) working correctly and it giving correct outputs for variable length prompts that I'd be happy to share (or create a separate PR for). It also adds some stuff to the generation utilities to support prompt tokens to the decoder that already exist in the PyTorch utilities (using prompt tokens instead of `model.config.decoder_start_token_id` if specified).

I haven't look into this. But @andyehrenberg do you suggest a different way of computation in Flax Whisper than the one implemented in our PyTorch/TensorFlow Whisper?

It's also better for @kamalkraj to express if he would like to continue this PR before we go ahead.

@ydshieh @andyehrenberg

If there is already a working implementation, please continue.
I am closing this one. 

Thanks

@ydshieh I guess what I'm suggesting for this could also be helpful for the PyTorch/TF implementations to improve flexibility/compatibility with existing codebases that use `position_ids` for other models (such as when finetuning).

For example, the use-case I'm working on is fine-tuning Whisper with RL (trying to expose it to its own outputs to reduce hallucinations). At each step when collecting rollouts, it is given a batch of audio features and decoder prompts (from previous audio snippets) - these prompts are of varying lengths, so padding/attention masks are needed, and the position embeddings need to adjust accordingly. And then when doing PPO updates on these steps, the position embeddings need to be computed correctly based off of which timesteps (tokens) are padding.

The implementation in this PR wouldn't accommodate this scenario as it assumes the same `past_key_values_length` for each sequence in the batch, whereas the implementation I've worked on uses `position_ids` to keep track of where we are in each sequence of the batch. Earlier I had use a different method that only used the attention mask along with another caching method in the decoder, but using position_ids is much simpler and accommodates multiple padding schemes more simply.