Add Llama Flax Implementation
Author: vvvm23
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.
-->

Fixes #26809. This is a work-in-progress port of Llama to Flax, leaving it as a draft PR for now.

The implementation is based heavily off the GPT-Neo and GPT-J Flax implementations.

Currently, the submodules are ready, I just need to assemble into a full model, check weight loading, add tests, and update the documentation.

## 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. [mentioned in this issue comment](https://github.com/huggingface/transformers/issues/22647#issuecomment-1579154174)
- [ ] 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.

@sanchit-gandhi 

<!-- 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: @sgugger

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam

Documentation: @sgugger, @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: @sgugger
- TensorFlow: @Rocketknight1

 -->


Comments:
Very cool @vvvm23! Scanned through the PR and it looks very nice already - happy to do a full review when it's close to completion. Just drop me a line and I'll have a look! 🚀 Likewise if you have any questions or queries, I'm on hand to help :)

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

Hi @vvvm23 and @sanchit-gandhi, do you guys have a timeline for this effort? Asking because I would love to import FlaxLlama from Hugging Face, but if it is going to take a while, I will probably build my own pipeline to import the model.

Not sure if this helps at all, but [here](https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/llama_model.py) you find an implementation of Llama in Flax (plus some other library-specific methods that you probably won't need).

Hi @gianlucadetommaso, I haven't had the time to work on this since this draft PR went live, but I am blocking time out this weekend to continue.

Cool to see community interest around running Flax Llama! Feel free to ping me here when you need a review @vvvm23!

Thanks @sanchit-gandhi I found a little time to continue today.

One issue I am noticing is that the tolerance when comparing the ground truth PyTorch implementation (in `modeling_llama.py`) and my own implementation, is a lot higher than I'd like. For three hidden layers in the decoder stack, I have to raise it to `atol=1e-2, rtol=1e-2`, with one hidden layer being at `atol=1e-3, rtol=1e-3` in order to pass. You can see the scratch test I am using at the bottom of `modeling_flax_llama.py`

I think some numerical differences are expected, but not sure to what degree. I am also testing with `float32` so that made me even more suspicious. Would you expected the results to be identical? This is my first time porting a PyTorch model to Flax. Thanks~

Update: I now have a full model working. I haven't checked if the pretrained weight loading wrappers (provided by the Flax GPTNeo implementation) work yet, but once they are it will be ready for review. I'll simultaneously clean it up and add some missing features whilst it is being reviewed.

Hey! Thanks for the progress update here @vvvm23 and great questions regarding numerical equivalence between models.

Generally, for any model less than 1B params we should be able to get equivalence to within 1e-5 between Flax and PyTorch. It's quite likely that you won't get this equivalence running the matmuls in bfloat16 on TPU. But you should be able to running the matmuls in float32, see https://github.com/huggingface/transformers/issues/15754 and https://github.com/google/jax/issues/10413#issue-1212211265 for details

Here's a script that I used previously for checking PT / Flax equivalence for BLOOM: https://github.com/sanchit-gandhi/codesnippets/blob/main/check_flax_bloom_jit_small_testing.ipynb You can ignore the bits about JIT'ing the forward pass for the time being. You can also uncomment the check to run it on CPU to force the highest precision, or use the decorator as provided

If we don't get 1e-5 precision, it's usually an indicator that we have a divergence in our model. Here, going through layer-by-layer and checking the hidden-states might be required to pinpoint it

Okay, thanks for the guidance and helper scripts 🔥 I expected that this lack of precision was not normal 😅 

I'll get the pretrained wrappers working first and then focus on debugging the numerical divergence.

I'm aiming for end of this week to fix those numerical issues, but my responsibilities elsewhere are pulling me a lot, so fingers crossed 🤞 

I've begun my hunt for numerical bugs 🐛

The first I squashed was rather strange. It seems `torch.rsqrt` and `jax.lax.rsqrt` do not match. This is used in the RMSNorm layers. Simple test to reproduce:
```
In [19]: a = np.asarray(a, dtype=np.float32)

In [20]: a
Out[20]:
array([1.16661310, 1.46686172, 0.13794081, 1.22346771, 1.17509305],
      dtype=float32)
In [21]: torch.rsqrt(torch.from_numpy(a))
Out[21]: tensor([0.92584139, 0.82566792, 2.69248700, 0.90407354, 0.92249471])

In [22]: jax.lax.rsqrt(a)
Out[22]: Array([0.92584133, 0.82566792, 2.69248700, 0.90407354, 0.92249471],      dtype=float32)

In [23]: 1 / torch.sqrt(torch.from_numpy(a))
Out[23]: tensor([0.92584139, 0.82566792, 2.69248700, 0.90407354, 0.92249471])

In [24]: 1 / jax.numpy.sqrt(a)
Out[24]: Array([0.92584139, 0.82566792, 2.69248700, 0.90407354, 0.92249471],      dtype=float32)
```
So the fix there was just to replace the `jax.lax.rsqrt` calls with `1 / jax.numpy.sqrt(...)`

Models still mismatches so I'll keep digging.

@sanchit-gandhi The model now numerically matches in fp32 on CPU. The issue was my backend has changed from CPU to GPU since fixing the `rsqrt` issue. I don't think we can expect a perfect match on GPU as the two models use fundamentally different backends. If there is anything you know of that could help remedy this, let me know.

What are the next steps to take? I am guessing some model tests, as well as trying it out on a real model checkpoint rather than random weights. However, my dev machine goes OOM when attempting to load the checkpoint on CPU.

Hey @vvvm23! Excellent work on pinpointing the difference between torch and jax.lax `rsqrt` and glad to hear we're within numerical precision using fp32 on CPU - we can be pretty confident we have an accurate Flax implantation based on these results. For GPU, there will be differences between PyTorch and JAX. This is expected since JAX fundamentally works differently to PyTorch with how it computes the matmuls, and is OK since the JAX model will typically generate predictions that are 'as good' as the PyTorch one.

Adding some tests and updating the docs would be the most sensible next steps! Again, you can refer to the Flax GPT Neo model to see the relevant tests to add: https://github.com/huggingface/transformers/blob/main/tests/models/gpt_neo/test_modeling_flax_gpt_neo.py

> However, my dev machine goes OOM when attempting to load the checkpoint on CPU.

That's interesting - are we loading the weights on GPU by accident? There shouldn't be any GPU OOM if running on CPU. We might see our RAM get full if loading extremely large weights, but the GPU memory shouldn't be affected. What model size are you loading? We can try the smallest 7b checkpoint: https://huggingface.co/meta-llama/Llama-2-7b

Awesome thanks, tests and docs it is! I am currently on leave so won't be progressing on this until the 31st.

> That's interesting - are we loading the weights on GPU by accident?
 
Actually, in the end no. By OOM on my dev machine, I meant out of CPU memory. Switching to a GPU backend meant I could load the model without running out of memory. So, nothing to worry about 😅 

Awesome - thanks for the update @vvvm23. Looking forward to doing a full review of the PR on your return! 

How's it looking @vvvm23? Let me know if I can help in anyway! Otherwise feel free to ping me here as soon as we're ready for a review, very excited to add this Flax model for the community!

Hi, currently been pretty split responsibility wise (moving house and job !!) so have only made a small bit of progress.

Most of the tests pass, however, there seems to be some matmul shape mismatch in the `generate_*` tests. Guessing I didn't implement the KV cache correctly, so I'll need to look at that. I also added back some missing docstrings.

I'll have some time to work on this Thursday, Friday (10th and 11th) but then probably nothing for another week :exploding_head: If you are in a rush and fancy trying to get the remaining tests to pass, please try! Sorry for the slowness on my part also! 

The final tests ended up being easy to fix: I had simply forgotten to swap the attention mask and position ids in the pretrained model wrapper.

@sanchit-gandhi I haven't retested the slow tests locally (as my laptop is slow) but later today I can run them, then tidy the code a bit. If all goes well, should be good for review later today or early tomorrow 👍 

@sanchit-gandhi all tests pass locally :tada: And I've also ran the model using the `generate` API to see if the outputs make sense:
```
In [23]: inputs = tokenizer('Aloha, World!', return_tensors='np')

In [24]: tokenizer.decode(model.generate(**inputs, generation_config=model.generation_config, max_length=100).sequences[0])
Out[24]: '
 Aloha, World!\nI’m back from my trip to Hawaii and I’m feeling great! I’m still trying to get back into the swing of things, but I’m getting there. I’m going to be posting a lot of pictures from my trip, so stay tuned!\nI’m also going to be posting a lot of pictures from my trip to Hawaii, so stay tuned!\nI’m also going to be posting a lot of pictures'
```
Seems good to me!

---

I think this is ready for review. I would like to draw your attention to a few points I was unsure about:

Firstly, the model currently throws a warning when loading pretrained checkpoints:
```
Some weights of the model checkpoint at openlm-research/open_llama_3b_v2 were not used when initializing FlaxLlamaForCausalLM: 
{('model', 'layers', '9', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '1', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '24', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '11', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '7', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '23', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '13', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '5', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '6', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '20', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '21', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '16', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '10', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '4', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '0', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '25', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '12', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '3', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '19', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '14', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '18', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '22', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '8', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '15', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '17', 'self_attn', 'rotary_emb', 'inv_freq'), ('model', 'layers', '2', 'self_attn', 'rotary_emb', 'inv_freq')}
```
This has no effect on the outputs, just simply the Flax version of the model does not store the `inv_freq` tensor for rotary embeddings within the state dictionary, so these just get discarded. Is there a way to suppress this warning so not to scare any users?

Secondly, please double check the licensing. I just copied this from the PyTorch version of Llama and updated the year.

Third, I use the checkpoint `openlm-research/open_llama_3b_v2` as it was the smallest, fully open Llama checkpoint I could find. The 'official' Llama checkpoints have gated access, so I am unsure if they are appropriate for testing / documentation purposes. This also means I haven't been able to test the model with the official Llama checkpoint as I still haven't managed to get permission from Meta :cry: 

Fourth, as we discussed a lot of the code is copied from the Flax implementation of GPT-Neo. There may be some leftover parts from there that we don't need in Llama, and I may have missed some best practices for Llama as GPT-Neo is (relatively) old now. In particular, see the following code block in `FlaxLlamaPreTrainedModel.__call__`:

```
        # TODO: can this handle input tensors being passed as kwargs? I copied GPT-Neo directly here
        outputs = self.module.apply(
            inputs,
            jnp.array(input_ids, dtype="i4"),
            jnp.array(attention_mask, dtype="i4"),
            jnp.array(position_ids, dtype="i4"),
            not train,
            False,
            output_attentions,
            output_hidden_states,
            return_dict,
            rngs=rngs,
            mutable=mutable,
        )
```

Finally, the tests pass but please check if they have sufficient coverage :hugs: 

---

Generally speaking, I have a lot of experience writing model code but no experience making large contributions to the Huggingface ecosystem, so there is almost certainly a lot wrong! Apologies in advance and I will do my best to help you bring this model to the finish line :muscle: Thanks for your work so far!

Thanks for your additional comments, I have some time to work on the more involved points today 🤗 

@sanchit-gandhi I think everything except the missing weight issue is resolved now (see my comment).

Trying to resolve some remaining CI issues, I noticed that the line `# Copied from transformers.models.gpt_neo.modeling_flax_gpt_neo.FlaxGPTNeoPreTrainedModel with GPTNeo->Llama` will change the line `@add_start_docstrings_to_model_forward(GPT_NEO_INPUTS_DOCSTRING)` , and overwrite `LLAMA_INPUTS_DOCSTRNG`. Any idea how to stop this happening? Otherwise the CI won't pass :thinking: 

That's correct behaviour @vvvm23! What we need to do is create the variable `LLAMA_INPUTS_DOCSTRNG` in the modelling file that contains the necessary docstring info for LLAMA (more or less copied one-for-one from Flax GPT Neo, but adapted for any different inputs)

Yeah, it is correct behaviour - what I meant though is that I did have a `LLAMA_INPUTS_DOCSTRING` in a previous commit, but running `make fix-copies` overwrote this docstring with the GPT-Neo version (as you suggested we add that at a class level). I guess my question was, how can we copy everything else in the class but somehow exclude the docstring line?

I get that we need the docstring itself, just currently the CI won't pass with both that docstring and the `# Copied from transformers.models.gpt_neo.modeling_flax_gpt_neo.FlaxGPTNeoPreTrainedModel with GPTNeo->Llama` line. Does the issue make sense?

@sanchit-gandhi fixed the CI issue (I think) by just adding more `Copied from ...` comments and deleting the class level comment. I also fixed the merge conflict. We should be good to go once CI passes I think 🙂 

@sanchit-gandhi the CI still fails, this is for two reasons. Could you assist me with resolving this?

1. The documentation test fails as it tries to load the checkpoint `_CHECKPOINT_FOR_DOC` however it needs `from_pt=True` to be set.
2. Flax Llama is based off the GPT Neo implementation. GPT Neo uses special tests to test equivalence between the flax and pytorch implementations. This overrides the common `test_equivalence_pt_to_flax` test. I copy these special tests (to make my flax tests pass). However, changing the tests for the flax version will cause the pytorch version to fail as it is using the flax version incorrectly.

> edit: for the second, please see `test_modeling_llama.py:308`. These tests need to be overriden somehow, for now I just return directly to get the CI to pass.

All the tests pass and the model is pretty much ready. Just not sure how to get through these last two blockers. Help would be much appreciated!

Awesome - thanks for the progress updates here @vvvm23. Agree with you that this is nearly in a position for a final review!

To fix your problems:
1. Could you open a pull request on the HF Hub to add the Flax model weights to the checkpoint? See this page for opening a PR: https://huggingface.co/openlm-research/open_llama_3b_v2/discussions (just click 'new pull request' and open a PR). You can ping me here once you've done this and I'll merge the Hub PR! Once this is done, we won't need the `from_pt` argument
2. Does the Flax-PyTorch test `test_equivalence_flax_to_pt` fail if we don't override it? I think it should be possible to use the model without this hack? Shall we just construct sensible attention masks in the model tester? I think it's the random attention mask that might be causing the problem: https://github.com/huggingface/transformers/blob/1fa2d89a9bb98a15e9720190e07d272a42f03d28/tests/models/llama/test_modeling_llama.py#L99

In practice, we would never have a 'random' attention mask like this. It would contain ones up to a certain point, then zeros after that. Could you try using a more realistic attention mask like this and seeing if that fixes the test?

> Could you open a pull request on the HF Hub to add the Flax model weights to the checkpoint?

PR is open 🙂 Git LFS is a great thing.

> You can ping me here once you've done this and I'll merge the Hub PR!

@sanchit-gandhi 

> Shall we just construct sensible attention masks in the model tester? I think it's the random attention mask that might be causing the problem

You are a genius! 😅 That was exactly the issue. I changed it to a fixed lower triangular tensor of ones (to test different lengths), padded with ones (to match size of the random one). This should give good enough coverage 👍 

The only test that fails now is the PR documentation one, which should be fixed once the checkpoint PR is complete.

That's awesome - happy to hear that fixed the issue! Let's wait until one of the OpenLM authors sees the open PR to add the flax weights: https://huggingface.co/openlm-research/open_llama_3b_v2/discussions/4#64f5b59fbce84cd8b1c80f38

Hey @sanchit-gandhi, thanks for the second review~

I can try with Llama 2 weights, however last time I requested access I never got a response. I will try again when I get the chance.

Just a heads up, I have basically no time to work on this until the start of next week. I am imminently moving which is basically sucking up all my time currently 🙃 I am also without GPU for a bit. I ask for your continued patience 🙏 or if you have the time to perhaps tackle a couple of the smaller changes yourself? Thanks for your understanding~

Hey @vvvm23, thanks for the update! Cool - let's try again with LLaMA v2 weights. Otherwise the Open Assistant ones are good ✅ (as you have currently)

Best of luck with the move, I hope it goes well! No worries about the delay - I'll be on hand when you return and can help with any issues that we're still facing! Unfortunately, I'm also a bit tied up right now, so can't complete this myself. But from my last review, it looks like you're in a strong position to finish the PR with some small final changes!

Hey @sanchit-gandhi, thanks for bearing with me. I have addressed your comments. They were quite small but I didn't have the headspace to think about this earlier 😅 

Regarding Llama 2 weights, I got the Meta approval, so I can test with the Llama 2 weights. Unfortunately, I am still without a GPU workstation and will likely OOM on CPU. So if you want results there you will have to bear with me more – but we should be good to go as far as a "Llama 1" implementation goes.

edit: also rebased the branch with `main` as we were pretty out of date 😅 

@sanchit-gandhi should be good to go now~ If there is anything else let me know, but I'll be at ICCV this week so maybe slow to respond :+1: 

Failing tests are related to missing Flax weights for the checkpoint [openlm-research/open_llama_3b_v2](https://huggingface.co/openlm-research/open_llama_3b_v2).

Could you load the PyTorch weights in Flax, and then push them to the org in a pull request @vvvm23?
```python
# load PT weights in FX
model = FlaxLlamaForCausalLM.from_pretrained("openlm-research/open_llama_3b_v2", from_pt=True)
# push FX weights to Hub
model.push_to_hub("openlm-research/open_llama_3b_v2", create_pr=True)
```
Once the PR is finalised by @ArthurZucker's review, we can merge this PR on the Hub and thus fix the failing doc tests! Thanks @vvvm23!

@sanchit-gandhi I forgot to say but the PR is up for the weights. Should be good to go once those PRs are merged.

@sanchit-gandhi @ArthurZucker a new issue has arisen. The tests fail when running `utils/check_docstrings.py` which prompts to run `make fix-copies`. However, this does not resolve the error. Any advice?

Hey @sanchit-gandhi @ArthurZucker just nudging on the above point, keen to wrap things up here!

To be more clear on my last, the error is:
```
Traceback (most recent call last):
  File "utils/check_docstrings.py", line 1239, in <module>
    check_docstrings(overwrite=args.fix_and_overwrite)
  File "utils/check_docstrings.py", line 1231, in check_docstrings
    raise ValueError(error_message)
ValueError: There was at least one problem when checking docstrings of public objects.
The following objects docstrings do not match their signature. Run `make fix-copies` to fix this.
- FlaxLlamaForCausalLM
- FlaxLlamaModel

Exited with code exit status 1
```

But running `make fix-copies` does not resolve the error.

Hey! I was off for a week let me try to help you here! 😉 

@ArthurZucker I added two lines that made the docstring checks ignore the Llama Flax models. This matches what GPTNeo did for their Flax models, so I assuming it is okay.

This (unsurprisingly) fixes that failing test. A few others have cropped up, but they seem unrelated to my changes. I will sync the branch with main and see if that fixes them. Other than that, and on the PRs in the individual model checkpoint hub pages, we should be good to go.

Merging with main fixed all the tests apart from the model checkpoint one. Should be good to go.

For the last test you can use `from_pretrained("path", revision = "refs/pr/4")`

Hi @ArthurZucker , the test that fails is the pr documentation test. The checkpoint can only be specified via `CHECKPOINT_FOR_DOC` which doesn't allow for setting revisions. Any other ideas?

Good point. In that case let's just merge the PRs on the hub, or I can open a PR to support a revision in the `add_docstring` 

PRs on the hub would be better, however I cannot merge these PRs :sweat_smile: 

We could always create our own internal copy and point the test at that?

I'll open a PR to support revisions anyway, another contributor is also stuck because of this 

Sorry for the delay, #27645 will help 😉 

nice! left some comments on that PR

Merged 😉 

The PR documentation test seems to time out. Is this a intermittent or known issue? Or something wrong with the code 🤔

@ArthurZucker @sanchit-gandhi could you advise on the above? Not intimately familiar with how your test workers operate.

Hi @vvvm23 I believe that in order to pass the tests you should update the https://huggingface.co/afmck/testing-llama-tiny/tree/main to have the flax ckpts and  also the tokenizer, I created this https://huggingface.co/ksmcg/Mistral-tiny/tree/main for Mistral and it has passed the tests.
Also I belive the revision should be removed to pass

Ah good catch, I will fix that tomorrow 

Finally, green light 😁

Let's just resolve the confilcts and good to go IMO!

done!

Thanks for the PR 🤗