Bug Fixed GPTNeoX Flax supports
Author: HeegyuKim
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 #22950:
- previous FlaxGPTNeoX support PR #22950 contains a error in the cached generation process. I resolved it.
- And I inserted a test code from #24002 
- There were 7 failures in the course of testing, I'm not sure why. It doesn't seem like a very fatal issue, but I'd appreciate it if you could check it out.
- The output with the pytorch model is very similar and the model works fine.
- k/v cache is already implemented.

@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:
I suffering from test issue. Can you help me? @sanchit-gandhi 

summary
- In pytorch, my GPTNeoX test failed to both test_equivalence_flax_to_pt and test_equivalence_pt_to_flax tests
- But in flax, GPTNeoX doesn't failed because my test code overrides it for not using `check_pt_flax_outputs`
- It's same to GPT Neo, Flax GPTNeo don't use `check_pt_flax_outputs` but Pytorch GPTNeo use it
- However, GPT Neo test in pytorch do not fail.
- Flax GPTNeo fails if it uses `check_pt_flax_outputs`

I don't think this is a problem with my model implementation. I wonder why pytorch's test fails.


This PR failed two tests below
```
FAILED tests/models/gpt_neox/test_modeling_gpt_neox.py::GPTNeoXModelTest::test_equivalence_flax_to_pt - AssertionError: 1.0483556 not less than or equal to 1e-05 : outputs.last_hidden_state: Difference between PyTorch and Flax is 1.0483555793762207 (>= 1e-05).
FAILED tests/models/gpt_neox/test_modeling_gpt_neox.py::GPTNeoXModelTest::test_equivalence_pt_to_flax - AssertionError: 1.8777691 not less than or equal to 1e-05 : outputs.last_hidden_state: Difference between PyTorch and Flax is 1.877769112586975 (>= 1e-05).
```
But two flax tests in tests/models/gpt_neox/test_modeling_flax_gptneox.py are fine.

the test code which was copied from #24002 override both test_equivalence_pt_to_flax and test_equivalence_flax_to_pt methods with this comment.
```
    # overwrite from common since `attention_mask` in combination
    # with `causal_mask` behaves slighly differently
```

and they use below assert code
```
# test_modeling_flax_gptneox.py line 267
for fx_output, pt_output in zip(fx_outputs, pt_outputs):
    self.assert_almost_equals(fx_output[:, -1], pt_output[:, -1].numpy(), 4e-2)
```

instead of `self.check_pt_flax_outputs(fx_outputs, pt_outputs, model_class)` in test_equivalence_pt_to_flax method in test_modeling_common.py

This overrides are equal to `tests/models/gpt_neo/test_modeling_flax_gptneo.py` but GPTNeo doesn't fail to pytorch test. I don't know what is different and 


Hey @HeegyuKim - could you confirm that you get the same logits out from the Flax model when you generate as with the PyTorch model? i.e. that the generation scores are the same in both cases. If this is indeed the case, then we can know for certain that the Flax implementation is correct, and that we need to override the PT-FX cross tests. Otherwise, there's a divergence that we need to fix!

We can check this with the **full** GPT NeoX model to ensure we have the right logits here

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.

Hey @HeegyuKim! I thought a little bit more about the PT-FX cross tests in the [WIP] Flax LLaMA port with @vvvm23, and suggested that probably the reason for the failing tests is the **random** attention mask: https://github.com/huggingface/transformers/pull/24587#issuecomment-1703144279

If we switch to using a causal attention mask, we are able to get PT-FX equivalence for Flax LLaMA **without** overriding the tests. Since Flax LLaMA is heavily based off Flax GPT-Neo, I'm fairly certain we'll observe similar behaviour for Flax GPT-NeoX

Would you like to try running the tests using a causal attention mask? E.g. as per https://github.com/huggingface/transformers/pull/24587#discussion_r1315128975

Hi! Thanks @HeegyuKim for the PR. I am wondering if there is any update on this? It would be really cool if we could use GPTNeoXForCausalLM in flax!

Hello @liutianlin0121 I'm trying to solve the problem whenever I have time. However, even if causal masking is applied, the error in the model output is still larger than 1e-5. The current error is around 0.02-0.03. I'm going to try again this weekend.

Even though there are errors, the model works better than expected. I trained several models with this code.

- [heegyu/WizardVicuna-pythia-1.4b-deduped](https://huggingface.co/heegyu/WizardVicuna-pythia-1.4b-deduped)
- [heegyu/RedTulu-Uncensored-3B-0719](https://huggingface.co/heegyu/RedTulu-Uncensored-3B-0719)
- [heegyu/polyglot-ko-5.8b-chat](https://huggingface.co/heegyu/polyglot-ko-5.8b-chat)

I want to contribute to huggingface but it's not as easy as I thought.

Ohhhhh I finally pass the equivalence issue! 🎉🎉

- I use FlaxGPTNeoXRotaryEmbedding class for RoPE and implement caching. This is a problem of the equivalence failure
- I remove overrides in tests/models/gpt_neox/test_modeling_flax_gpt_neox.py and it works!
 
But there are CI failures...
- [check_repository_consistency](https://app.circleci.com/pipelines/github/huggingface/transformers/73610/workflows/4216a793-3879-4e27-85de-9aefe7e75998/jobs/931100): I copied it identically, but why does this message appear?
- [tests_pr_documentation_tests](https://app.circleci.com/pipelines/github/huggingface/transformers/73610/workflows/00dfecbe-027f-41a7-93ff-7161660486f3/jobs/931111): Huggingface's GPT-NeoX 20B model does not have a flax weight. How should we solve this problem?

@sanchit-gandhi 

Well done @HeegyuKim, that's excellent news! Regarding the two failing tests:
* You can run `make fix-copies` to update the modelling code with any copied functions? The linter will copy all the code so that it is one-for-one the same
* Could you open a pull request on the Hugging Face Hub to add the Flax weights to [EleutherAI/gpt-neox-20b](https://huggingface.co/EleutherAI/gpt-neox-20b)? You can convert first load the PyTorch weights into Flax:
```python
from transformers import FlaxAutoModelForCausalLM

model = FlaxAutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neox-20b", from_pt=True)
```

And then push the converted Flax weights to the Hub:
```python
model.push_to_hub("EleutherAI/gpt-neox-20b", create_pr=True)
```

I think we're almost at the end of our work but there are small issues.

### Suddenly wav2vec2 test fails??
- Suddenly wav2vec model tests are failed ([tests_torch CI link](https://app.circleci.com/pipelines/github/huggingface/transformers/74617/workflows/f61c5c58-d39a-4394-8ebc-8b551d74da3b/jobs/945928)). It seems to have something to do with gelu_fast that I added. I don't think wav2vec2 uses gelu_fast, but I don't know why.
- GPT-NeoX-20B uses gelu_fast activation. I converted pytorch implemented to flax version in [src/transformers/modeling_flax_utils.py](src/transformers/modeling_flax_utils.py) and added gelu_fast. 

### Flax weights
- I opened the PR of GPT-NeoX 20B Flax weights - https://huggingface.co/EleutherAI/gpt-neox-20b/discussions/24
- the Flax tests use pythia-410m models but there is no flax weights. So I use `from_pt=True` is it Ok?

### `Copied from` issue
- `make fix-copies` changes every `GPTNeoXBlahBlah` -> `GPTNeoBlahBlah` (config, comments) even there is a `Copied from ... with GPTNeo->GPTNeoX` mark. https://github.com/huggingface/transformers/pull/25334/commits/a670443b2e4fdd88262c93eb98a9dc0330c64812#r1287443977 
- So I moved `copied from` comments to the exactly same method.

@sanchit-gandhi 

Hey @HeegyuKim! Nice job on iterating here! Answering your questions in-line below:
1. For the failing Wav2Vec2 issues, you can rebase onto the main branch of `transformers`, where the tests have been fixed:
```
git fetch upstream
git rebase upstream main
git push -f origin main
```
Note that it's important you force push (`-f` flag) after a rebase to preserve the correct commit history of this PR!
2. Thanks for converting the PyTorch GELU Fast activation to JAX - this looks great!
3. Thanks also for pushing the Flax weights - we can merge them once this PR is approved by the next reviewer and prior to merging this PR. Using `from_pt=True` is ok for the Flax tests - just make sure you decorate the tests with `@is_pt_flax_cross_test` since we need both PyTorch and Flax when we load from pre-trained with `from_pt=True`
4. The issue you had before was that there was an extra space after the start and end of the right arrow: `GPTNeo -> GPTNeoX` should be `GPTNeo->GPTNeoX`. Can you try adding this copied from before `FlaxGPTNeoXPreTrainedModel`? This should allow you to copy the entire module:
```
# Copied from transformers.models.gpt_neo.modeling_flax_gpt_neo.FlaxGPTNeoPreTrainedModel with GPTNeo->GPTNeoX, GPT_NEO->GPT_NEO_X, "transformer"->"gpt_neox"
```

Finally documentation is left, how can I make a documentation for it? @sanchit-gandhi 
```
Exception: The following objects are in the public init so should be documented:
 - FlaxGPTNeoXForCausalLM
 - FlaxGPTNeoXModel
```

You can do so with `make repo-consistency`!

I may passed necessary CI tests! @sanchit-gandhi 

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.

The PR is almost finished! Would you like to make the last remaining changes @HeegyuKim such that we can get this one merged? Let us know if you have any questions, more than happy to help here

Thank you for your comment! I'll check it this weekend

Hello @sanchit-gandhi, I rebased this PR to main branch and pushed again.

There are two CI failures - First is a [documentation issue](https://app.circleci.com/pipelines/github/huggingface/transformers/82163/workflows/26395aa8-e716-4cc0-a97b-8ad4c7ffe89b/jobs/1056756). 
```
OSError: Can't load the model for 'EleutherAI/gpt-neox-20b'. If you were trying to load it from 'https://huggingface.co/models', make sure you don't have a local directory with the same name. Otherwise, make sure 'EleutherAI/gpt-neox-20b' is the correct path to a directory containing a file named flax_model.msgpack or pytorch_model.bin.
```
This problem can be fixed when this [GPT-NeoX model PR](https://huggingface.co/EleutherAI/gpt-neox-20b/discussions/24) is merged. Alternatively, we can add from_pt=True to the example. 

As for the [second issue](https://app.circleci.com/pipelines/github/huggingface/transformers/82163/workflows/216462aa-cd92-40d8-ab8e-3e4103bc29b5/jobs/1056753), I don't know why. I would appreciate it if you could tell me the cause and solution to this problem.
```
ValueError: The main __init__ has objects that are not present in transformers.utils.dummy_flax_objects.py. Run `make fix-copies` to fix this.
```

I ran `make fix-copies` but following error occurs.
```
> make fix-copies

python utils/check_copies.py --fix_and_overwrite
Traceback (most recent call last):
  File "/home/heegyu/transformers/utils/check_copies.py", line 1129, in <module>
    check_copies(args.fix_and_overwrite, args.file)
  File "/home/heegyu/transformers/utils/check_copies.py", line 778, in check_copies
    new_diffs = is_copy_consistent(filename, overwrite, buffer)
  File "/home/heegyu/transformers/utils/check_copies.py", line 736, in is_copy_consistent
    diff_index = check_codes_match(observed_code, theoretical_code)
  File "/home/heegyu/transformers/utils/check_copies.py", line 549, in check_codes_match
    theoretical_name = re_pattern.search(theoretical_code_header).groups()[0]
AttributeError: 'NoneType' object has no attribute 'groups'
```



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.