Flax mistral
Author: kiansierra
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 #26809
Implements Mistral Models in Flax, used many components from Llama PR https://github.com/huggingface/transformers/pull/24587


## 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.

@sanchit-gandhi , @ArthurZucker 

<!-- 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: @muellerzr 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 @sanchit-gandhi , @ArthurZucker 
This is the first draft of the PR, no need to review yet, as I still got to implement the slow tests and documentation changes.
I did want you're thoughts on the change on src/transformers/modeling_flax_pytorch_utils.py, as when loading from_pt weights they were in bfloat16, and I'm a bit weary of changing core functionality

Hi @sanchit-gandhi , @ArthurZucker I think this is ready to be reviewed now.
Some of the implementation details were done with the intention to pass some of the tests, so I think I'll highlight them.
* In `MISTRAL_START_DOCSTRING`.
```dtype (`jax.numpy.dtype`, *optional*, defaults to `float32`)```, instead of ```dtype (`jax.numpy.dtype`, *optional*, defaults to `jax.numpy.float32`)```
There seemed to be a change that happened in `utils/check_docstrings.py`
* In `FlaxMistralPreTrainedModel` set `_seed` and `_input_shape` instead of  `seed` and `input_shape` for the same `utils/check_docstrings.py` as these appeared in one of the docstrings to compared to
* Converted checkpoint to Flax and stored in `_CHECKPOINT_FOR_DOC = "ksmcg/Mistral-7B-v0.1"`, this allows to pass the [testing docstring](https://github.com/huggingface/transformers/blob/0baa9246cb1ddac355db1df7824a521426599eb7/src/transformers/utils/doc.py#L1022), since it doesn't allow for the `from_pt=True` flag 
* As far as I can tell these are many of the issues that are being faced in Llama PR https://github.com/huggingface/transformers/pull/24587

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hi @ArthurZucker I was wondering if you could please review this PR
I've seen in https://github.com/huggingface/transformers/pull/24587 that you've been on holidays, hope you had a good time.
I believe many of the issues are the same that occur in that PR, so you would like already know why they're failing.
Since the doc test doesn't allow the load from_pt flag, I created a checkpoint in https://huggingface.co/ksmcg/Mistral-7B-v0.1/tree/main to load it from, but this fails since it takes a long time to execute

Sure I'll have a look 😉 The checkpoints need to be merged to add the doc, I'll open a PR to support a revision I think! 

#27645 will allow you to use the revision in the docstring 😉 I'll finish this asap

Thanks for the wait! Should be usable 😉 

Yes the revision fix is working, now the issue is it times out

Hi @ArthurZucker I've passed all the tests.
If you could please review and let me know your feedback.
It is quite similar, and I think we should wait for https://github.com/huggingface/transformers/pull/24587 to be merged so I can add the fix copies from that PR

Hi @ArthurZucker, I've updated the PR and added the different # Copied from transformers.models.llama.modeling_flax_llama now that it has been merged

This looks awesome! Any chance it can be merged? Really looking forward to it.

Oups sorry @kiansierra let's just rebase now and I'll merge! 

Hey, no worries @ArthurZucker, I've merged main back in

Is this solved? Or is it deprecated in favour of a v0.2 one?

Overall, does huggingface's mistral 7B have a flax implementation?

Hi @sanchit-gandhi ,
First of all thank you for your review and feedback
I believe all of the comments you've raised have been resolved.
Currently Github doesn't let me comment on the `convert_pytorch_sharded_state_dict_to_flax` I've added reconversion back to bf16 as was implemented in `convert_pytorch_state_dict_to_flax`
A lot more `#Copied from` were added.
I think the change I feel most confused about was the needing to freeze the params to have the masking work, that I expand on [here](https://github.com/huggingface/transformers/pull/26943#discussion_r1427962523) 
Also I made a 1 character change to remove a deprecation warning from pytest in `tests/test_modeling_flax_common.py`
Let me know you're thoughts


Hi @sanchit-gandhi, Could you please let me know if there is anything more to add to the PR?
Happy holidays, and thanks to the HF team for the great work you're all doing

Hi @sanchit-gandhi , I implemented the above changes.
The ModelTester can't be directly copied because of the different configurations.
I believe the separate PR you're refering to is this one https://github.com/huggingface/transformers/pull/28367 right?


Hi @sanchit-gandhi, @ArthurZucker just wanted to make sure if there is any additional task on my side to complete this PR

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_26943). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Hi @ArthurZucker, @sanchit-gandhi 
I've implemented the latest changes and merged main branch in.

Awesome! Requesting final review from @ArthurZucker - once we get his approval, we can merge this one @kiansierra 

Hi @ArthurZucker I run the tests with `RUN_SLOW=1 RUN_PT_FLAX_CROSS_TESTS=1 pytest tests/models/mistral`

```
================================================================== short test summary info ===================================================================
FAILED tests/models/mistral/test_modeling_mistral.py::MistralModelTest::test_eager_matches_sdpa_generate - RuntimeError: "addmm_impl_cpu_" not implemented for 'Half'
FAILED tests/models/mistral/test_modeling_mistral.py::MistralIntegrationTest::test_speculative_generation - RuntimeError: "addmm_impl_cpu_" not implemented for 'Half'
============================================ 2 failed, 139 passed, 45 skipped, 92 warnings in 1813.70s (0:30:13) =============================================
```
I believe these fails are due to running the tests without GPU.
For some reason the `tests_pr_documentation_tests ` is failing due to 
```
ERROR src/transformers/models/mistral/modeling_flax_mistral.py - AttributeError: 'HfDoctestModule' object has no attribute 'getini'
ERROR dummy.py - AttributeError: 'HfDoctestModule' object has no attribute 'getini'
```
Do you know why this might be?

Yep no worries for the PR documentation test. We can merge without it / rebasing should also help! Left a last nit about `key = jnp.repeat(key, repeats=self.num_key_value_groups, axis=2)` and let's merge 🚀 

Hi @ArthurZucker that's it done, changed to repeat and merged main in

Thanks for your hard works and congrats on the merge 🤗 

Thank you both for your time reviewing this, it's been a great learning experience