Add NucleusX Model
Author: syncdoth
Pull request: yes
Labels: name: WIP

# 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 -->

This PR adds a new model named NucleusX. This model is contributed by [Sehyun Choi](https://github.com/syncdoth) and [NucleusAI](https://www.withnucleus.ai/). The model is based on the [Retentive Network](https://arxiv.org/abs/2307.08621) architecture, and the code is largely adapted from [this repo](https://github.com/syncdoth/retnet.git), which again borrows core implementations from [torchscale](https://github.com/microsoft/torchscale). We are planning to release our paper and weights soon.


## 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?
- [ ] 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.
- [x] 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?

We kindly request the review of this new model from @ArthurZucker and @younesbelkada!

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: @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:
The current test failure at `tests_pr_documentation_tests` is due to the incorrect repo_id and links, namely `NucleusAI/NucleusX-7B` and https://huggingface.co/NucleusAI/NucleusX-7B used in examples and model docs. These checkpoints are not released yet; we plan to release them soon.

cc: @sippycoder and also @LysandreJik!

Hey! Thanks for opening the PR, I'll let @Rocketknight1 do a first review as he is more familiar with this kind of models! 

Hi all! RetNets seem like a really interesting architecture, so I'm quite excited to take a look - I'll try to review this in the next day or two.

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

@Rocketknight1 Thanks for reviewing this PR! I have gone through the comments and resolved them. There are also some other updates:

- `dtype` handling for tensors created in `NucleusXRelPos` to have the same `dtype` as the model weights
- rename some `*layer_norm` modules to `*rms_norm` for conformity.
- removed `subln` option (Sub-LayerNorm), which is not applicable to our choice of FFN (SwiGLU).

There are other minor changes, which can be found in the commit logs. 

As per weight release, we are working hard to make that happen :) We'll ping here when the weights are ready for public release.

Thanks again!

Also @syncdoth, while we're waiting for the checkpoints, there are some tests failing in this PR that are unrelated. If you pull the latest version from main and rebase your PR, that should fix them.

> Also @syncdoth, while we're waiting for the checkpoints, there are some tests failing in this PR that are unrelated. If you pull the latest version from main and rebase your PR, that should fix them.

This may be a beginner question, but should I rebase main and (force?) push or merge main and push?

Probably the easiest way to do it is to pull the latest version of main, then rebase your branch onto main, and then force push.

Hi @syncdoth, do you know what happened to Nucleus AI? The website is now down

> Hi @syncdoth, do you know what happened to Nucleus AI? The website is now down

This is unrelated to this PR but there's some maintenance going on with the website. Hang tight :)

btw @syncdoth if you're still getting test failures, try 'sync upstream' on the `main` branch of your forked repo, then on your development machine, pull the latest main branch, change to the `add_nucleus_x` branch, rebase and finally force push. Should resolve everything!

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.

Don't stale, please! This looks quite close to being ready! (cc @syncdoth - let me know if you need any help with the last bit!)

We are on the verge of releasing the weight! There's been a bit of delay in the schedule 🥲

The last bit should be updating the weight links in the docs and writing the integration tests; We are working on it hard!

Hi @Rocketknight1, I’m seeing test failure related to document building, and testing the run of NucleusXForCausalLM.forward example. It seems that it might be due to `.from_pretrained` from a 7B checkpoint killing the worker, like the previous example in the configuration. Do you think I should change the example to a smaller one?

Does it require some tinkering to use `generate` in not parallel mode?  (I don't have RAM for processing 16KB prompt in parallel)

I dumped source to model folder, edited config to treat it as `trusted_remoted_code=True` thingy, parallel works fine, as in [test](https://github.com/huggingface/transformers/blob/b624d0518776d1761d29a448b83de4b2ef4d6658/tests/models/nucleus_x/test_modeling_nucleus_x.py#L465):

```python
In [7]: print(tokenizer.decode(model.generate(**tokenizer("Hello my name is", return_tensors="pt").to("cuda"), max_new_tokens=20, do_sample=False, forward_mod
   ...: e="parallel").ravel()))
Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.
/home/fella/src/llama/text-generation-webui/models/NucleusAI_Nucleus-X/modeling_nucleus_x.py:370: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  cache = (current_kv, scale, torch.tensor(prev_seqlen + 1, dtype=torch.long))

 Hello my name is Tina and I am a 25 year old female. I am a very outgoing person
```

but  recurrent no

```python
In [8]: print(tokenizer.decode(model.generate(**tokenizer("Hello my name is", return_tensors="pt").to("cuda"), max_new_tokens=20, do_sample=False, forward_mod
   ...: e="recurrent").ravel()))
Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.
/home/fella/src/llama/text-generation-webui/models/NucleusAI_Nucleus-X/modeling_nucleus_x.py:370: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  cache = (current_kv, scale, torch.tensor(prev_seqlen + 1, dtype=torch.long))

 Hello my name is the most of the world.
The first thing I noticed was the size of the room. It
```

(Even if I say in config.json to use recurrent forward mode, 16KB prompt fails to pass through model.generate unless I use forward_mode='recurrent')

Hi @syncdoth, sorry for the Christmas delay! You're correct, though - the issue is almost certainly caused by the docstring trying to load a model too big for the test runner. Is there any smaller checkpoint we can use? You could also try `torch_dtype=torch.bfloat16`.

Haha plz don't stale this! We are still working hard to put out the model. We are working on a small model to pass the PR requirement, but it has been a lower priority unfortunately :(  will finish to finish this within mid Feb!

No worries 🤗