Update modeling_cpmbee.py and README.md

README.md

@@ -67,4 +67,39 @@ res = model.generate(
 )
 print(res)
 
-```
+```
+
+We suggest to use `bmtrain` to finetune CPM-Bee. Also, you can use `accelerate` and `deepspeed` to finetune CPM-Bee. Here we will give a brief example of a training loop:
+
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from accelerate import Accelerator
+from torch.utils.data import Dataset, DataLoader
+
+accelerator = Accelerator()
+
+trainset = Dataset()  # Make sure trainset.__getitem__() can get data with correct format like {"input": "...", "<ans>": ""}
+# for details, you can read https://github.com/OpenBMB/CPM-Bee/tree/main/tutorials/basic_task_finetune
+train_loader = DataLoader(trainset, batch_size=1)
+
+tokenizer = AutoTokenizer.from_pretrained("openbmb/cpm-bee-2b", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("openbmb/cpm-bee-2b", trust_remote_code=True).cuda()
+
+optimizer = torch.optim.Adam(model.parameters())
+
+model, optimizer, train_loader = accelerator.prepare(
+    model, optimizer, train_loader
+)
+
+for iter, data in enumerate(train_loader):
+    optimizer.zero_grad()
+
+    # change the data to a trainable format
+    input_encoded = tokenizer.prepare_for_finetune(data, max_length=512).to(model.device)
+
+    outputs = model(**input_encoded)
+    loss = outputs.loss
+    accelerator.backward(loss)
+    optimizer.step()
+```
+You should design your own parallel and mix_precision training strategy on the basis of it.

modeling_cpmbee.py

@@ -569,10 +569,10 @@ class CpmBeeRotaryEmbedding(nn.Module):
         self.inv_freq = inv_freq.to(config.torch_dtype)
 
     def forward(self, x: torch.Tensor, x_pos: torch.Tensor):
-        inv_freq = self.inv_freq.to(device=x.device, dtype=self.dtype)
+        inv_freq = self.inv_freq.to(device=x.device, dtype=x.dtype)
 
         x_pos = x_pos * self.distance_scale
-        freqs = x_pos[..., None].to(self.dtype) * inv_freq[None, :]  # (..., dim/2)
+        freqs = x_pos[..., None] * inv_freq[None, :]  # (..., dim/2)
 
         emb = torch.cat((freqs, freqs), dim=-1)  # (..., dim)
         emb_cos = emb.cos()  # (..., dim)