Repository: hpcaitech/ColossalAI

Issue:
 [inference] Adapted to Rotary Embedding and RMS Norm
## 📌 Checklist before creating the PR

- [ ] I have created an issue for this PR for traceability
- [x] The title follows the standard format: `[doc/gemini/tensor/...]: A concise description`
- [x] I have added relevant tags if possible for us to better distinguish different PRs


## 🚨 Issue number

> Link this PR to your issue with words like fixed to automatically close the linked issue upon merge
>
> e.g. `fixed #1234`, `closed #1234`, `resolved #1234`



## 📝 What does this PR do?

> Summarize your work here.
> if you have any plots/diagrams/screenshots/tables, please attach them here.
pytest:
![image](https://github.com/hpcaitech/ColossalAI/assets/26091437/7da782ae-3bfc-45e2-b7d4-f2726ed72498)
Benchmark:
batchsize 16 input_len 256 output_len 128
Caiinference:
![image](https://github.com/hpcaitech/ColossalAI/assets/26091437/36554b0b-115a-420d-b5f3-3c016e128695)
transformers(flash attn):
![image](https://github.com/hpcaitech/ColossalAI/assets/26091437/2bc82d19-6cd8-4a76-b78a-a0b83b580111)
batchsize 32 input_len 256 output_len 128
Caiinference:
![image](https://github.com/hpcaitech/ColossalAI/assets/26091437/a72e2dc7-97d8-41f7-925c-10f9775c45f7)
transformers(flash attn):
![image](https://github.com/hpcaitech/ColossalAI/assets/26091437/5a750743-bae9-47f9-b5ad-ae66f90d7c56)
batchsize 64 input_len 256 output_len 128
Caiinference:
![image](https://github.com/hpcaitech/ColossalAI/assets/26091437/511d75b6-ff95-4806-813c-f4d4f8d74c8c)
transformers(flash attn):
![image](https://github.com/hpcaitech/ColossalAI/assets/26091437/1c7e4c45-de80-4ab0-814d-5ed4b64110aa)
batchsize 16 input_len 1024 output_len 128
Caiinference:
![image](https://github.com/hpcaitech/ColossalAI/assets/26091437/fc4b3e96-9063-4e54-a03f-7531d522ca87)
transformers(flash attn):
![image](https://github.com/hpcaitech/ColossalAI/assets/26091437/861837f1-532d-4658-b2fa-be8d57ee65d8)
batchsize 16 input_len 1024 output_len 128
Caiinference:
![image](https://github.com/hpcaitech/ColossalAI/assets/26091437/6fd346c8-b51c-435f-a889-0946a5bf7c9c)
transformers(flash attn):
![image](https://github.com/hpcaitech/ColossalAI/assets/26091437/78f61220-4329-4255-8fd4-9131afe60717)
batchsize 16 input_len 1024 output_len 128
Caiinference:
![image](https://github.com/hpcaitech/ColossalAI/assets/26091437/705c1509-a1bf-48ca-b01f-8ef9839a4dbb)
transformers(flash attn):
OOM
## 💥 Checklist before requesting a review

- [ ] I have linked my PR to an issue ([instruction](https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue))
- [ ] My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
- [x] I have performed a self-review of my code
- [x] I have added thorough tests.
- [x] I have added docstrings for all the functions/methods I implemented

## ⭐️ Do you enjoy contributing to Colossal-AI?

- [x] 🌝 Yes, I do.
- [ ] 🌚 No, I don't.

Tell us more if you don't enjoy contributing to Colossal-AI.

[tensor] fix some unittests

[tensor] fix some unittests



Discussion:

Resolving patch:
diff --git a/colossalai/inference/utils.py b/colossalai/inference/utils.py
new file mode 100644
--- /dev/null
+++ b/colossalai/inference/utils.py
@@ -0,0 +1,51 @@
+"""
+Utils for model inference
+"""
+import os
+
+import torch
+
+
+def init_to_get_rotary(self, base=10000, use_elem=False):
+    """
+    This function initializes the rotary positional embedding, it is compatible for all models and is called in ShardFormer
+    Args:
+        self : Model that holds the rotary positional embedding
+        base : calculation arg
+        use_elem : activated when using chatglm-based models
+    """
+    self.config.head_dim_ = self.config.hidden_size // self.config.num_attention_heads
+    if not hasattr(self.config, "rope_scaling"):
+        rope_scaling_factor = 1.0
+    else:
+        rope_scaling_factor = self.config.rope_scaling.factor if self.config.rope_scaling is not None else 1.0
+
+    if hasattr(self.config, "max_sequence_length"):
+        max_seq_len = self.config.max_sequence_length
+    elif hasattr(self.config, "max_position_embeddings"):
+        max_seq_len = self.config.max_position_embeddings * rope_scaling_factor
+    else:
+        max_seq_len = 2048 * rope_scaling_factor
+    base = float(base)
+
+    # NTK  ref: https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
+    ntk_alpha = os.environ.get("INFER_NTK_ALPHA", None)
+
+    if ntk_alpha is not None:
+        ntk_alpha = float(ntk_alpha)
+        assert ntk_alpha >= 1, "NTK alpha must be greater than or equal to 1"
+        if ntk_alpha > 1:
+            print(f"Note: NTK enabled, alpha set to {ntk_alpha}")
+        max_seq_len *= ntk_alpha
+        base = base * (ntk_alpha ** (self.head_dim_ / (self.head_dim_ - 2)))  # Base change formula
+
+    n_elem = self.config.head_dim_
+    if use_elem:
+        n_elem //= 2
+
+    inv_freq = 1.0 / (base ** (torch.arange(0, n_elem, 2, device="cpu", dtype=torch.float32) / n_elem))
+    t = torch.arange(max_seq_len + 1024 * 64, device="cpu", dtype=torch.float32) / rope_scaling_factor
+    freqs = torch.outer(t, inv_freq)
+
+    self._cos_cached = torch.cos(freqs).to(torch.float16).cuda()
+    self._sin_cached = torch.sin(freqs).to(torch.float16).cuda()


Test patch: