Add OLMo model family
Author: 2015aroras
Pull request: yes
Labels: name: New model

# 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.
-->

This PR adds the OLMo model family to transformers. A base `OLMoModel` and a casual LM `OLMoForCausalLM` are implemented. The models are already present in HF Hub (e.g. [allenai/OLMo-7B](https://huggingface.co/allenai/OLMo-7B)), and this implementation is compatible with the checkpoints in the Hub.

UPDATE: The current version of the PR is not compatible with the old OLMo models in HF Hub. New models have been uploaded to HF Hub (e.g. [allenai/OLMo-7B-hf](https://huggingface.co/allenai/OLMo-7B-hf)) to support this PR.

<!-- Remove if not applicable -->

Fixes #29885


## 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#create-a-pull-request),
      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?

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:
@ArthurZucker and @younesbelkada for general guidance and/or review, since this is a PyTorch text model.

> Hey! Thanks for submitting the PR! The code seems in a very research state, let's try to get it to `transformers` philosophy. My main query would be: what are the main differences with `MixtralMoe` ? (Llama + MoE) ? Only include differences for trained checkpoints (I see both alibi and rotary, all of them use both? Or only one was kept?)
> 
> The fast tokenizer also looks very similar to GptNeoX / Bloom one, any reason to have a new one?

Thanks for having a look!

In a general sense, I tried to stay as faithful to OLMo as I could while following general `transformers` principles (all model code in 1 file, use "Copied from" with re-usable logic, etc.).

Regarding tokenization, OLMo adds an eos token to the end of inputs that don't have them. The GptNeoX and Bloom ones don't. Llama's tokenizer does, but it also has a slow tokenizer and I couldn't figure out how to convert the fast OLMo tokenizer to a slow `sentencepiece` one. My current implementation copies a lot from the Llama fast tokenizer.

In terms of the model, OLMo doesn't have MoE whereas `MixtralMoe` does. I would have to figure out if I could use `MixtralMoe` while avoiding the MoE functionality. OLMo is more directly map-able to Llama; the model architectures mostly map 1 to 1, and the main discrepancies in output come from the 2 models using different precision for some operations.

The official OLMo models use rope and at most 1 of rope and alibi is used in a run (i.e. no alibi), but we hope to release checkpoints from our experiments at some point. Some of those would use alibi. I think it would be nice if transformers could support those checkpoints too.

@dirkgr FYI

Hey! Sorry I might have misunderstood (thought I saw some MoE things here). 
Alright when I mean research I mean there are a lot of things that we never do in transformers. 
1. Tokenizer, we can update bloom / Gpneox to support this option as IMO it makes a lot of sense! No need for a slow one 😉 
2. The modeling code does not seem to include a lot of `# Copied from`, would recommend you to look at `Gemma` modeling for a good example! Basically we want to avoid all the autocasting code when it is not necessary. we do checkpointing a bit differently and we want to have the least number of code paths as possible. 
3. For the config, all of the different enums for layer norms and act fn, we already have the `from ...activations import ACT2FN
` and the `ALL_LAYER_NORMS` which should be used. 
4. `BufferCache` is not really something we use, and I am not sure we had issue with the community after the last Llama fixes!
5. `k_norm` and `q_norm` are used in `Persimmon` so a good place to copy from as well!
6. `_scaled_dot_product_attention` should follow the classing we use! So xxxxSdpaAttention
7. if `clip_qkv` is always true let's not have a if else!
8. attention mask logic should use the `_update_causal_mask` from Llama and have a similar logic
9. We don't use `reset_parameters` or `device` let's try to get closer to Llama!
10. We don't incude `get_fsdp_wrap_policy` and fsdp specifics in transformers
11. Let's add some generation tests!


Just want to clarify my thoughts before I try making changes, so I'm expressing my understanding of what I should do below. Please correct me where I am wrong.

Regarding the tokenizer, if possible I should update one of the existing tokenizers and then use that tokenizer directly (similar to how `mamba` appears to use `GPTNeoXTokenizerFast`).

The comments regarding modeling code suggest that I should aim to re-use existing code using `#Copied from` and any existing functions. More generally, I should aim to be more "transformers-like" than faithful to the original code. I should try to make the code run like OLMo (produce the same output) but look like transformers.

Assuming it were possible, the ideal solution would be for me to use some existing model and convert OLMo checkpoints to work with that model.

----

Regarding copying existing code vs being faithful, what should I do when some operation (say, RoPE) exists in transformers but is performed in a different precision (e.g. float32 vs bfloat16) in OLMo compared to the transformer code? Do I use the transformers code and disregard the "small" error, or write separate code that matches the original OLMo behavior?

> Assuming it were possible, the ideal solution would be for me to use some existing model and convert OLMo checkpoints to work with that model.

Yes! This way it's easier for all the community to use the model, and keeps the diff very clear and clean! 

For the tokenizer, if `GPTNeoXTokenizerFast` does not have `add_eos` or `add_bos`, we can add it but in a backward compatible way! 

> The comments regarding modeling code suggest that I should aim to re-use existing code using #Copied from and any existing functions. More generally, I should aim to be more "transformers-like" than faithful to the original code. I should try to make the code run like OLMo (produce the same output) but look like transformers.

Exactly 💯 

For ROPE, we should always compute it in `float32` that is from experiences and long debugging. 
Otherwise you can use the code to match original Olmo behavior where it is needed of course! We aim for 1e-4 logits equivalence!

Due to the above discussion, I have redone the implementation. The key points are:

- For tokenizer, OLMo now uses `GPTNeoXTransformerFast`. I have modified the tokenizer to make adding `bos_token` or `eos_token` an option. Unfortunately, that tokenizer doesn't come with any tests...
- For model, I have reverted mostly to the code generated from [add-new-model-like with Llama](https://github.com/huggingface/transformers/pull/29890/commits/0afbd4a35925394ee4c9010e6bf985d87adac764). The main changes are changing the norm into the non-parametric version used by OLMo and removing auto-generated model classes like `OLMoForQuestionAnswering`. Should I add those classes back (I have not spent time trying to understand those classes yet)?
- The config class is almost identical to that of Llama. The config and the model do not support all non-official checkpoints (for example, alibi is not supported but is used in some non-official OLMo runs).

The existing checkpoints on HF Hub are not compatible with this implementation, so integration tests are currently broken. I can work on fixing that once we are happy with the approach/implementation. On a 1B checkpoint that I converted using the updated conversion script, the maximal logit difference is `~2.5e-5`.

Thanks! Pinging @molbap for the review! 

@molbap Gentle ping

Do you have any recommendations for how to transition our checkpoints in HF Hub from the old style that relied on remote code to the new checkpoints from this PR? I can accordingly work on updating HF Hub and fixing the integration tests.

Other than integration tests that rely on a HF Hub checkpoint, everything else should be passing now.



@2015aroras For checkpoint adaptation, I'd take a look at [conversion scripts for gemma ](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma/convert_gemma_weights_to_hf.py). Overall not too far from what you are doing, load from config with correctly shaped parameters. Usually how I go about it is running the two implementations side-by-side in 2 different debuggers and look where the activations start to diverge, fixing the conversion script as I go.
For the tests, with the rebase it looks better now indeed! When the integration tests run we'll be able to ask for a core maintainer final approval :) 

@molbap Sorry, I should have been more specific about my concerns regarding transitioning our checkpoints in HF hub. My problem is that if I load the new-style OLMo checkpoints (i.e. OLMo checkpoints converted using `convert_olmo_weights_to_hf.py`), then it would break HF OLMo for people without the code from this PR. OLMo is already consumable using `OLMoForCausalLM.from_pretrained("OLMo-7B", trust_remote_code=True)`, but once I upload the new-style checkpoint their code will break until they update the HF library. How can I upload the new-style OLMo checkpoints without breaking people that are using older `transformers` versions?

It looks like Llama got around this issue by having separate `meta-llama/Llama-2-7b` and `meta-llama/Llama-2-7b-hf`, where the latter is "converted for the Hugging Face Transformers format". Is this approach acceptable? EDIT: It seems like a bunch of models do this. it would be nice to have this in the [contributing guide](https://huggingface.co/docs/transformers/v4.38.2/en/contributing) somewhere if it's the right approach to take.

Ah, I understand better now! Yes, we usually recommend to go for one repo on the hub = one model. In that case, it would be acceptable to add a -hf suffix to the new checkpoints on the hub!

I have now put up "-hf" models as _private_ on the allenai HF Hub organization. Once I get the all clear internally, we will make them public. All the integration tests pass for me locally (they fail here because the repo is private).

We are hoping to release some new model checkpoints using the same architecture early next week, so it would be great if we could get this PR in quickly. Then we would be able to release models natively integrated with HF along with the rest of the release!

I think the PR is in a good state. Thanks for adding this! pinging @ArthurZucker for a final review :) 

@molbap @ArthurZucker All checks are passing! Now it's just a matter of reviews

Addressed all PR comments, and hopefully all checks will pass.