Move applying rotary embeddings inside LlamaRotaryEmbedding class
Author: kunal-vaishnavi
Pull request: yes

# What does this PR do?

This PR moves `rotate_half` and `apply_rotary_pos_emb` inside the `LlamaRotaryEmbedding` base class. When all of the rotary embedding computations are in one class, the [export of LLaMA-2 to ONNX](https://github.com/pytorch/pytorch/pull/109759) can be done with the `export_modules_as_functions` option for the `LlamaRotaryEmbedding` class in `torch.onnx.export`. With all of the rotary embedding computations in one function in the exported ONNX model, it is easier to optimize the model and [integrate with Optimum](https://github.com/huggingface/optimum/pull/1289).

## Who can review?

@fxmarty, @ArthurZucker, @younesbelkada

Comments:
Thank you @kunal-vaishnavi, is it for later easier pattern matching on the rotary positional embedding? Could you share a given ONNX exported with the argument `export_modules_as_functions` set (e.g. https://huggingface.co/fxmarty/tiny-llama-fast-tokenizer)?

There are a few other archs that do the same: persimmon, gptj, gpt_neox_japanese, gpt_neox, idefics, esm, codegen.

I am wondering whether removing `rotate_half` and `apply_rotary_pos_emb` is not a breaking change given that those are not private functions, but @ArthurZucker will be a better judge.

> Thank you @kunal-vaishnavi, is it for later easier pattern matching on the rotary positional embedding?

Yes, it is for easier pattern matching on the rotary positional embedding. Without representing the rotary embedding as a function, there is a lot of careful pattern matching to do.

> Could you share a given ONNX exported with the argument `export_modules_as_functions` set (e.g. https://huggingface.co/fxmarty/tiny-llama-fast-tokenizer)?

The `export_modules_as_functions` option will convert an `nn.Module` into an ONNX `FunctionProto` (see the [ONNX docs](https://onnx.ai/onnx/intro/python.html#functions) for more details on functions). In Netron, the function will look something like this at the top level.

![TorchScript Export Modules as Functions Example](https://github.com/huggingface/transformers/assets/115581922/41dee37d-08e5-4fe2-a189-61c54e281d3d)

If you click on the "f" in the top right, you can open the function and examine its contents.

![TorchScript Functions Example](https://github.com/huggingface/transformers/assets/115581922/1f963dd5-66df-4627-9ffc-21c802a53e1e)


> There are a few other archs that do the same: persimmon, gptj, gpt_neox_japanese, gpt_neox, idefics, esm, codegen.

Yes, these changes can be made for the other model architectures as well to enable easier pattern matching after exporting to ONNX.


It's a breaking change and unless you have a strong argument for this in terms of speed of exportation or runtime, it's gonna be a hard to have this one merged 😓 We can however try to have a deprecation cycle if this makes sense. But if it's just for easier pattern matching I'm not sure it's worth the effort wdyt? 

Could you explain a bit more on why the proposed approach would be a breaking change? Is it with regards to the scope of the functions and that they would only be seen in the class now? From my understanding, the two methods are used to apply the rotary embeddings and should be called in the forward pass of the class for the rotations to be performed. Additionally, the methods are only called in the attention forward pass (which has access to the class through `self.rotary_emb`) and don't appear to be used elsewhere.

As a workaround, can we create another subclass similar to the scaling variants that inherit from `LlamaRotaryEmbedding`? The new class can have the two functions in it, and we can still leave the default rotary embedding class as `LlamaRotaryEmbedding` during execution. In the attention layer, there can be a check to see if the new class is being used. If yes, then we can call the forward pass to get the new query and key states. If no, then we can leave the current code as is.

With the subclass approach, we can maintain current behavior and the deprecation can also be done over time to switch from `LlamaRotaryEmbedding` to the new class as the default.

The simplest fix is to keep both functions and also add them in the class. The breaking part is that we can’t just remove it as it’s a « global » function. And the issue is that LlamaRotary is used in other models which would need the same changes. 

my question still applies: is there a strong incentive? 

> The simplest fix is to keep both functions and also add them in the class.

Adding the functions to the class would not be enough. They need to be used in the forward pass to fully capture the operations with `export_modules_as_functions`. I was suggesting another class if modifying the forward pass to `LlamaRotaryEmbedding` is considered a breaking change.

> And the issue is that LlamaRotary is used in other models which would need the same changes.

The `LlamaRotaryEmbedding` code is copied for other models and not imported (as seen [here](https://github.com/search?q=repo%3Ahuggingface%2Ftransformers%20path%3A%2F%5Esrc%5C%2Ftransformers%5C%2Fmodels%5C%2F%2F%20llamarotaryembedding&type=code)) so changes to `LlamaRotaryEmbedding` wouldn't directly impact those models. A similar change could be made for those models to maintain consistency. Given the goal of this proposed PR is to make changes to enable optimizing the LLaMA-2 ONNX model and integrate with Optimum, those changes could be in another PR.

> is there a strong incentive?

Using `export_modules_as_functions` speeds up the export to ONNX and decreases model size. This is because the functions are templated so only one copy is needed for reference. Without it, the current exported model has a copy of the nodes that comprise the function at each layer. Additionally, the ONNX model loads much faster in Netron and is much easier to view because there are less nodes at the top level.

> it's gonna be a hard to have this one merged

Since it seems unlikely that these PR changes will be approved, I will begin switching to another approach. With the new approach, however, I won't be able to integrate the ORTOptimizer class in Optimum with optimizing the LLaMA-2 ONNX model converted with Optimum.

Sounds good. I added a logger warning for deprecation and made two model changes to satisfy the CI pipeline requirements. 

- OpenLLaMA: The changes from LLaMA seem fine to use so I copied them over. 
- Persimmon: The rotary embeddings calculations appear different so I removed the "Copied from transformers.models.llama.modeling_llama" references to satisfy the copies checker.

I removed the "copied from" comments because the copy checker was failing. After further examination, it was failing because a space was needed at the end of the first line in the logger warning. I added the spaces and the "copied from" comments are now back.

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