We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG]: LoRA does not support the training of reward models
### 🐛 Describe the bug

![image](https://user-images.githubusercontent.com/34190033/220584377-0467d526-d1c4-46a6-b4c8-c56a088bff96.png)


### Environment

### env
OS：ubuntu 20.04
GPU：4 x A10
python==3.9.0
torch==1.13.1-cu116
colossalai==0.2.5

### command
`python train_reward_model.py --pretrain "bigscience/bloom-560m" --lora_rank 16`


[tensor] fix some unittests

[tensor] fix some unittests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `applications/ChatGPT/chatgpt/trainer/rm.py`
Content:
```
1 from abc import ABC
2 
3 import loralib as lora
4 import torch
5 from chatgpt.dataset import RewardDataset
6 from chatgpt.nn import PairWiseLoss
7 from torch.optim import Adam, Optimizer
8 from torch.utils.data import DataLoader
9 from tqdm import tqdm
10 
11 from .strategies import Strategy
12 from .utils import is_rank_0
13 
14 
15 class RewardModelTrainer(ABC):
16     """
17         Trainer to use while training reward model.
18 
19     Args:
20         model (torch.nn.Module): the model to train
21         strategy (Strategy): the strategy to use for training
22         optim(Optimizer): the optimizer to use for training
23         train_dataset (RewardDataset): the dataset to use for training
24         eval_dataset (RewardDataset): the dataset to use for evaluation
25         batch_size (int, defaults to 1): the batch size while training
26         max_epochs (int, defaults to 2): the number of epochs to train
27         optim_kwargs (dict, defaults to {'lr':1e-4}): the kwargs to use while initializing optimizer
28     """
29 
30     def __init__(
31         self,
32         model,
33         strategy: Strategy,
34         optim: Optimizer,
35         train_dataset: RewardDataset,
36         eval_dataset: RewardDataset,
37         batch_size: int = 1,
38         max_epochs: int = 2,
39     ) -> None:
40         super().__init__()
41         self.strategy = strategy
42         self.epochs = max_epochs
43         self.train_dataloader = DataLoader(train_dataset, batch_size=batch_size)
44         self.eval_dataloader = DataLoader(eval_dataset, batch_size=batch_size)
45 
46         self.model = strategy.setup_model(model)
47         self.loss_fn = PairWiseLoss()
48         self.optimizer = strategy.setup_optimizer(optim, self.model)
49 
50     def fit(self, use_lora):
51         epoch_bar = tqdm(range(self.epochs), desc='Train epoch', disable=not is_rank_0())
52         for epoch in range(self.epochs):
53             step_bar = tqdm(range(self.train_dataloader.__len__()),
54                             desc='Train step of epoch %d' % epoch,
55                             disable=not is_rank_0())
56             # train
57             if use_lora > 0:
58                 print("Using Lora")
59                 lora.mark_only_lora_as_trainable(self.model.body)
60 
61             else:
62                 self.model.train()
63             for chosen_ids, c_mask, reject_ids, r_mask in self.train_dataloader:
64                 chosen_ids = chosen_ids.squeeze(1).cuda()
65                 c_mask = c_mask.squeeze(1).cuda()
66                 reject_ids = reject_ids.squeeze(1).cuda()
67                 r_mask = r_mask.squeeze(1).cuda()
68                 chosen_reward = self.model(chosen_ids, attention_mask=c_mask)
69                 reject_reward = self.model(reject_ids, attention_mask=r_mask)
70                 loss = self.loss_fn(chosen_reward, reject_reward)
71                 self.strategy.backward(loss, self.model, self.optimizer)
72                 self.strategy.optimizer_step(self.optimizer)
73                 self.optimizer.zero_grad()
74                 step_bar.update()
75                 step_bar.set_postfix({'loss': loss.item()})
76 
77             # eval
78             self.model.eval()
79             with torch.no_grad():
80                 dist = 0
81                 loss_sum = 0
82                 for chosen_ids, c_mask, reject_ids, r_mask in self.eval_dataloader:
83                     chosen_ids = chosen_ids.squeeze(1).cuda()
84                     c_mask = c_mask.squeeze(1).cuda()
85                     reject_ids = reject_ids.squeeze(1).cuda()
86                     r_mask = r_mask.squeeze(1).cuda()
87                     chosen_reward = self.model(chosen_ids, attention_mask=c_mask)
88                     reject_reward = self.model(reject_ids, attention_mask=r_mask)
89                     dist += (chosen_reward - reject_reward).mean().item()
90                     loss = self.loss_fn(chosen_reward, reject_reward)
91                     loss_sum += loss.item()
92                 dist_mean = dist / self.eval_dataloader.__len__()
93                 loss_mean = loss_sum / self.eval_dataloader.__len__()
94             epoch_bar.update()
95             step_bar.set_postfix({'loss': loss_mean, 'dist_mean': dist_mean})
96             step_bar.close()
97
```

Path: `applications/ChatGPT/chatgpt/nn/reward_model.py`
Content:
```
1 from typing import Optional
2 
3 import torch
4 import torch.nn as nn
5 
6 from .lora import LoRAModule
7 
8 
9 class RewardModel(LoRAModule):
10     """
11     Reward model base class.
12 
13     Args:
14         model (nn.Module): Reward model.
15         value_head (nn.Module): Value head to get reward score.
16         lora_rank (int): LoRA rank.
17         lora_train_bias (str): LoRA bias training mode.
18     """
19 
20     def __init__(self,
21                  model: nn.Module,
22                  value_head: Optional[nn.Module] = None,
23                  lora_rank: int = 0,
24                  lora_train_bias: str = 'none') -> None:
25         super().__init__(lora_rank=lora_rank, lora_train_bias=lora_train_bias)
26         self.body = model
27         if value_head is not None:
28             if value_head.out_features != 1:
29                 raise ValueError("The value head of reward model's output dim should be 1!")
30             self.value_head = value_head
31 
32         else:
33             self.value_head = nn.Linear(model.config.n_embd, 1)
34         self.convert_to_lora()
35 
36     def forward(self, sequences: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
37         outputs = self.body(sequences, attention_mask=attention_mask)
38         last_hidden_states = outputs['last_hidden_state']
39         values = self.value_head(last_hidden_states)[:, :-1]
40         value = values.mean(dim=1).squeeze(1)    # ensure shape is (B)
41         return value
42
```

Path: `applications/ChatGPT/examples/train_reward_model.py`
Content:
```
1 import argparse
2 
3 import loralib as lora
4 import torch
5 from chatgpt.dataset import RewardDataset
6 from chatgpt.nn import BLOOMRM, GPTRM, OPTRM
7 from chatgpt.trainer import RewardModelTrainer
8 from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy
9 from datasets import load_dataset
10 from torch.optim import Adam
11 from transformers import AutoTokenizer, BloomTokenizerFast
12 from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer
13 
14 from colossalai.nn.optimizer import HybridAdam
15 
16 
17 def train(args):
18     # configure strategy
19     if args.strategy == 'naive':
20         strategy = NaiveStrategy()
21     elif args.strategy == 'ddp':
22         strategy = DDPStrategy()
23     elif args.strategy == 'colossalai_gemini':
24         strategy = ColossalAIStrategy(stage=3, placement_policy='cuda')
25     elif args.strategy == 'colossalai_zero2':
26         strategy = ColossalAIStrategy(stage=2, placement_policy='cuda')
27     else:
28         raise ValueError(f'Unsupported strategy "{args.strategy}"')
29 
30     # configure model
31     with strategy.model_init_context():
32         if args.model == 'bloom':
33             model = BLOOMRM(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()
34         elif args.model == 'opt':
35             model = OPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()
36         elif args.model == 'gpt2':
37             model = GPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()
38         else:
39             raise ValueError(f'Unsupported model "{args.model}"')
40 
41     # configure tokenizer
42     if args.model == 'gpt2':
43         tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
44         tokenizer.pad_token = tokenizer.eos_token
45     elif args.model == 'bloom':
46         tokenizer = BloomTokenizerFast.from_pretrained(args.pretrain)
47         tokenizer.pad_token = tokenizer.eos_token
48     elif args.model == 'opt':
49         tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
50     else:
51         raise ValueError(f'Unsupported model "{args.model}"')
52     tokenizer.pad_token = tokenizer.eos_token
53 
54     max_len = 512
55 
56     # configure optimizer
57     if args.strategy.startswith('colossalai'):
58         optim = HybridAdam(model.parameters(), lr=5e-5)
59     else:
60         optim = Adam(model.parameters(), lr=5e-5)
61 
62     # prepare for data and dataset
63     data = load_dataset(args.dataset)
64     train_data = data["train"].select(range(100))
65     eval_data = data['test'].select(range(5))
66     train_dataset = RewardDataset(train_data, tokenizer, max_len)
67     eval_dataset = RewardDataset(eval_data, tokenizer, max_len)
68 
69     # batch_size here is expected to be C(k,2), k means # response of each prompt
70     # be limited with the format of dataset 'Dahoas/rm-static', we'd better use batch_size as 1
71     trainer = RewardModelTrainer(model=model,
72                                  strategy=strategy,
73                                  optim=optim,
74                                  train_dataset=train_dataset,
75                                  eval_dataset=eval_dataset,
76                                  batch_size=args.batch_size,
77                                  max_epochs=args.max_epochs)
78 
79     trainer.fit(use_lora=args.lora_rank)
80 
81     # save model checkpoint after fitting on only rank0
82     strategy.save_model(model, 'rm_checkpoint.pt', only_rank0=True)
83     # save optimizer checkpoint on all ranks
84     strategy.save_optimizer(optim, 'rm_optim_checkpoint_%d.pt' % (torch.cuda.current_device()), only_rank0=False)
85 
86 
87 if __name__ == '__main__':
88     parser = argparse.ArgumentParser()
89     parser.add_argument('--strategy',
90                         choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'],
91                         default='naive')
92     parser.add_argument('--model', choices=['gpt2', 'bloom', 'opt'], default='bloom')
93     parser.add_argument('--pretrain', type=str, default=None)
94     parser.add_argument('--dataset', type=str, default='Dahoas/rm-static')
95     parser.add_argument('--save_path', type=str, default='rm_ckpt.pth')
96     parser.add_argument('--max_epochs', type=int, default=10)
97     parser.add_argument('--batch_size', type=int, default=4)
98     parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank")
99     args = parser.parse_args()
100     train(args)
101
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/applications/ChatGPT/chatgpt/nn/reward_model.py b/applications/ChatGPT/chatgpt/nn/reward_model.py
--- a/applications/ChatGPT/chatgpt/nn/reward_model.py
+++ b/applications/ChatGPT/chatgpt/nn/reward_model.py
@@ -23,7 +23,7 @@
                  lora_rank: int = 0,
                  lora_train_bias: str = 'none') -> None:
         super().__init__(lora_rank=lora_rank, lora_train_bias=lora_train_bias)
-        self.body = model
+        self.model = model
         if value_head is not None:
             if value_head.out_features != 1:
                 raise ValueError("The value head of reward model's output dim should be 1!")
@@ -34,7 +34,7 @@
         self.convert_to_lora()
 
     def forward(self, sequences: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
-        outputs = self.body(sequences, attention_mask=attention_mask)
+        outputs = self.model(sequences, attention_mask=attention_mask)
         last_hidden_states = outputs['last_hidden_state']
         values = self.value_head(last_hidden_states)[:, :-1]
         value = values.mean(dim=1).squeeze(1)    # ensure shape is (B)
diff --git a/applications/ChatGPT/chatgpt/trainer/rm.py b/applications/ChatGPT/chatgpt/trainer/rm.py
--- a/applications/ChatGPT/chatgpt/trainer/rm.py
+++ b/applications/ChatGPT/chatgpt/trainer/rm.py
@@ -44,6 +44,8 @@
         self.eval_dataloader = DataLoader(eval_dataset, batch_size=batch_size)
 
         self.model = strategy.setup_model(model)
+        if "DDP" in str(self.strategy):
+            self.model = self.model.module
         self.loss_fn = PairWiseLoss()
         self.optimizer = strategy.setup_optimizer(optim, self.model)
 
@@ -56,7 +58,7 @@
             # train
             if use_lora > 0:
                 print("Using Lora")
-                lora.mark_only_lora_as_trainable(self.model.body)
+                lora.mark_only_lora_as_trainable(self.model.model)
 
             else:
                 self.model.train()
diff --git a/applications/ChatGPT/examples/train_reward_model.py b/applications/ChatGPT/examples/train_reward_model.py
--- a/applications/ChatGPT/examples/train_reward_model.py
+++ b/applications/ChatGPT/examples/train_reward_model.py
@@ -61,8 +61,8 @@
 
     # prepare for data and dataset
     data = load_dataset(args.dataset)
-    train_data = data["train"].select(range(100))
-    eval_data = data['test'].select(range(5))
+    train_data = data["train"]
+    eval_data = data['test']
     train_dataset = RewardDataset(train_data, tokenizer, max_len)
     eval_dataset = RewardDataset(eval_data, tokenizer, max_len)
 
@@ -93,7 +93,7 @@
     parser.add_argument('--pretrain', type=str, default=None)
     parser.add_argument('--dataset', type=str, default='Dahoas/rm-static')
     parser.add_argument('--save_path', type=str, default='rm_ckpt.pth')
-    parser.add_argument('--max_epochs', type=int, default=10)
+    parser.add_argument('--max_epochs', type=int, default=1)
     parser.add_argument('--batch_size', type=int, default=4)
     parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank")
     args = parser.parse_args()

{"golden_diff": "diff --git a/applications/ChatGPT/chatgpt/nn/reward_model.py b/applications/ChatGPT/chatgpt/nn/reward_model.py\n--- a/applications/ChatGPT/chatgpt/nn/reward_model.py\n+++ b/applications/ChatGPT/chatgpt/nn/reward_model.py\n@@ -23,7 +23,7 @@\n                  lora_rank: int = 0,\n                  lora_train_bias: str = 'none') -> None:\n         super().__init__(lora_rank=lora_rank, lora_train_bias=lora_train_bias)\n-        self.body = model\n+        self.model = model\n         if value_head is not None:\n             if value_head.out_features != 1:\n                 raise ValueError(\"The value head of reward model's output dim should be 1!\")\n@@ -34,7 +34,7 @@\n         self.convert_to_lora()\n \n     def forward(self, sequences: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:\n-        outputs = self.body(sequences, attention_mask=attention_mask)\n+        outputs = self.model(sequences, attention_mask=attention_mask)\n         last_hidden_states = outputs['last_hidden_state']\n         values = self.value_head(last_hidden_states)[:, :-1]\n         value = values.mean(dim=1).squeeze(1)    # ensure shape is (B)\ndiff --git a/applications/ChatGPT/chatgpt/trainer/rm.py b/applications/ChatGPT/chatgpt/trainer/rm.py\n--- a/applications/ChatGPT/chatgpt/trainer/rm.py\n+++ b/applications/ChatGPT/chatgpt/trainer/rm.py\n@@ -44,6 +44,8 @@\n         self.eval_dataloader = DataLoader(eval_dataset, batch_size=batch_size)\n \n         self.model = strategy.setup_model(model)\n+        if \"DDP\" in str(self.strategy):\n+            self.model = self.model.module\n         self.loss_fn = PairWiseLoss()\n         self.optimizer = strategy.setup_optimizer(optim, self.model)\n \n@@ -56,7 +58,7 @@\n             # train\n             if use_lora > 0:\n                 print(\"Using Lora\")\n-                lora.mark_only_lora_as_trainable(self.model.body)\n+                lora.mark_only_lora_as_trainable(self.model.model)\n \n             else:\n                 self.model.train()\ndiff --git a/applications/ChatGPT/examples/train_reward_model.py b/applications/ChatGPT/examples/train_reward_model.py\n--- a/applications/ChatGPT/examples/train_reward_model.py\n+++ b/applications/ChatGPT/examples/train_reward_model.py\n@@ -61,8 +61,8 @@\n \n     # prepare for data and dataset\n     data = load_dataset(args.dataset)\n-    train_data = data[\"train\"].select(range(100))\n-    eval_data = data['test'].select(range(5))\n+    train_data = data[\"train\"]\n+    eval_data = data['test']\n     train_dataset = RewardDataset(train_data, tokenizer, max_len)\n     eval_dataset = RewardDataset(eval_data, tokenizer, max_len)\n \n@@ -93,7 +93,7 @@\n     parser.add_argument('--pretrain', type=str, default=None)\n     parser.add_argument('--dataset', type=str, default='Dahoas/rm-static')\n     parser.add_argument('--save_path', type=str, default='rm_ckpt.pth')\n-    parser.add_argument('--max_epochs', type=int, default=10)\n+    parser.add_argument('--max_epochs', type=int, default=1)\n     parser.add_argument('--batch_size', type=int, default=4)\n     parser.add_argument('--lora_rank', type=int, default=0, help=\"low-rank adaptation matrices rank\")\n     args = parser.parse_args()\n", "issue": "[BUG]: LoRA does not support the training of reward models\n### \ud83d\udc1b Describe the bug\n\n![image](https://user-images.githubusercontent.com/34190033/220584377-0467d526-d1c4-46a6-b4c8-c56a088bff96.png)\r\n\n\n### Environment\n\n### env\r\nOS\uff1aubuntu 20.04\r\nGPU\uff1a4 x A10\r\npython==3.9.0\r\ntorch==1.13.1-cu116\r\ncolossalai==0.2.5\r\n\r\n### command\r\n`python train_reward_model.py --pretrain \"bigscience/bloom-560m\" --lora_rank 16`\r\n\r\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "from abc import ABC\n\nimport loralib as lora\nimport torch\nfrom chatgpt.dataset import RewardDataset\nfrom chatgpt.nn import PairWiseLoss\nfrom torch.optim import Adam, Optimizer\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\n\nfrom .strategies import Strategy\nfrom .utils import is_rank_0\n\n\nclass RewardModelTrainer(ABC):\n    \"\"\"\n        Trainer to use while training reward model.\n\n    Args:\n        model (torch.nn.Module): the model to train\n        strategy (Strategy): the strategy to use for training\n        optim(Optimizer): the optimizer to use for training\n        train_dataset (RewardDataset): the dataset to use for training\n        eval_dataset (RewardDataset): the dataset to use for evaluation\n        batch_size (int, defaults to 1): the batch size while training\n        max_epochs (int, defaults to 2): the number of epochs to train\n        optim_kwargs (dict, defaults to {'lr':1e-4}): the kwargs to use while initializing optimizer\n    \"\"\"\n\n    def __init__(\n        self,\n        model,\n        strategy: Strategy,\n        optim: Optimizer,\n        train_dataset: RewardDataset,\n        eval_dataset: RewardDataset,\n        batch_size: int = 1,\n        max_epochs: int = 2,\n    ) -> None:\n        super().__init__()\n        self.strategy = strategy\n        self.epochs = max_epochs\n        self.train_dataloader = DataLoader(train_dataset, batch_size=batch_size)\n        self.eval_dataloader = DataLoader(eval_dataset, batch_size=batch_size)\n\n        self.model = strategy.setup_model(model)\n        self.loss_fn = PairWiseLoss()\n        self.optimizer = strategy.setup_optimizer(optim, self.model)\n\n    def fit(self, use_lora):\n        epoch_bar = tqdm(range(self.epochs), desc='Train epoch', disable=not is_rank_0())\n        for epoch in range(self.epochs):\n            step_bar = tqdm(range(self.train_dataloader.__len__()),\n                            desc='Train step of epoch %d' % epoch,\n                            disable=not is_rank_0())\n            # train\n            if use_lora > 0:\n                print(\"Using Lora\")\n                lora.mark_only_lora_as_trainable(self.model.body)\n\n            else:\n                self.model.train()\n            for chosen_ids, c_mask, reject_ids, r_mask in self.train_dataloader:\n                chosen_ids = chosen_ids.squeeze(1).cuda()\n                c_mask = c_mask.squeeze(1).cuda()\n                reject_ids = reject_ids.squeeze(1).cuda()\n                r_mask = r_mask.squeeze(1).cuda()\n                chosen_reward = self.model(chosen_ids, attention_mask=c_mask)\n                reject_reward = self.model(reject_ids, attention_mask=r_mask)\n                loss = self.loss_fn(chosen_reward, reject_reward)\n                self.strategy.backward(loss, self.model, self.optimizer)\n                self.strategy.optimizer_step(self.optimizer)\n                self.optimizer.zero_grad()\n                step_bar.update()\n                step_bar.set_postfix({'loss': loss.item()})\n\n            # eval\n            self.model.eval()\n            with torch.no_grad():\n                dist = 0\n                loss_sum = 0\n                for chosen_ids, c_mask, reject_ids, r_mask in self.eval_dataloader:\n                    chosen_ids = chosen_ids.squeeze(1).cuda()\n                    c_mask = c_mask.squeeze(1).cuda()\n                    reject_ids = reject_ids.squeeze(1).cuda()\n                    r_mask = r_mask.squeeze(1).cuda()\n                    chosen_reward = self.model(chosen_ids, attention_mask=c_mask)\n                    reject_reward = self.model(reject_ids, attention_mask=r_mask)\n                    dist += (chosen_reward - reject_reward).mean().item()\n                    loss = self.loss_fn(chosen_reward, reject_reward)\n                    loss_sum += loss.item()\n                dist_mean = dist / self.eval_dataloader.__len__()\n                loss_mean = loss_sum / self.eval_dataloader.__len__()\n            epoch_bar.update()\n            step_bar.set_postfix({'loss': loss_mean, 'dist_mean': dist_mean})\n            step_bar.close()\n", "path": "applications/ChatGPT/chatgpt/trainer/rm.py"}, {"content": "from typing import Optional\n\nimport torch\nimport torch.nn as nn\n\nfrom .lora import LoRAModule\n\n\nclass RewardModel(LoRAModule):\n    \"\"\"\n    Reward model base class.\n\n    Args:\n        model (nn.Module): Reward model.\n        value_head (nn.Module): Value head to get reward score.\n        lora_rank (int): LoRA rank.\n        lora_train_bias (str): LoRA bias training mode.\n    \"\"\"\n\n    def __init__(self,\n                 model: nn.Module,\n                 value_head: Optional[nn.Module] = None,\n                 lora_rank: int = 0,\n                 lora_train_bias: str = 'none') -> None:\n        super().__init__(lora_rank=lora_rank, lora_train_bias=lora_train_bias)\n        self.body = model\n        if value_head is not None:\n            if value_head.out_features != 1:\n                raise ValueError(\"The value head of reward model's output dim should be 1!\")\n            self.value_head = value_head\n\n        else:\n            self.value_head = nn.Linear(model.config.n_embd, 1)\n        self.convert_to_lora()\n\n    def forward(self, sequences: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:\n        outputs = self.body(sequences, attention_mask=attention_mask)\n        last_hidden_states = outputs['last_hidden_state']\n        values = self.value_head(last_hidden_states)[:, :-1]\n        value = values.mean(dim=1).squeeze(1)    # ensure shape is (B)\n        return value\n", "path": "applications/ChatGPT/chatgpt/nn/reward_model.py"}, {"content": "import argparse\n\nimport loralib as lora\nimport torch\nfrom chatgpt.dataset import RewardDataset\nfrom chatgpt.nn import BLOOMRM, GPTRM, OPTRM\nfrom chatgpt.trainer import RewardModelTrainer\nfrom chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy\nfrom datasets import load_dataset\nfrom torch.optim import Adam\nfrom transformers import AutoTokenizer, BloomTokenizerFast\nfrom transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer\n\nfrom colossalai.nn.optimizer import HybridAdam\n\n\ndef train(args):\n    # configure strategy\n    if args.strategy == 'naive':\n        strategy = NaiveStrategy()\n    elif args.strategy == 'ddp':\n        strategy = DDPStrategy()\n    elif args.strategy == 'colossalai_gemini':\n        strategy = ColossalAIStrategy(stage=3, placement_policy='cuda')\n    elif args.strategy == 'colossalai_zero2':\n        strategy = ColossalAIStrategy(stage=2, placement_policy='cuda')\n    else:\n        raise ValueError(f'Unsupported strategy \"{args.strategy}\"')\n\n    # configure model\n    with strategy.model_init_context():\n        if args.model == 'bloom':\n            model = BLOOMRM(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()\n        elif args.model == 'opt':\n            model = OPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()\n        elif args.model == 'gpt2':\n            model = GPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()\n        else:\n            raise ValueError(f'Unsupported model \"{args.model}\"')\n\n    # configure tokenizer\n    if args.model == 'gpt2':\n        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n        tokenizer.pad_token = tokenizer.eos_token\n    elif args.model == 'bloom':\n        tokenizer = BloomTokenizerFast.from_pretrained(args.pretrain)\n        tokenizer.pad_token = tokenizer.eos_token\n    elif args.model == 'opt':\n        tokenizer = AutoTokenizer.from_pretrained(\"facebook/opt-350m\")\n    else:\n        raise ValueError(f'Unsupported model \"{args.model}\"')\n    tokenizer.pad_token = tokenizer.eos_token\n\n    max_len = 512\n\n    # configure optimizer\n    if args.strategy.startswith('colossalai'):\n        optim = HybridAdam(model.parameters(), lr=5e-5)\n    else:\n        optim = Adam(model.parameters(), lr=5e-5)\n\n    # prepare for data and dataset\n    data = load_dataset(args.dataset)\n    train_data = data[\"train\"].select(range(100))\n    eval_data = data['test'].select(range(5))\n    train_dataset = RewardDataset(train_data, tokenizer, max_len)\n    eval_dataset = RewardDataset(eval_data, tokenizer, max_len)\n\n    # batch_size here is expected to be C(k,2), k means # response of each prompt\n    # be limited with the format of dataset 'Dahoas/rm-static', we'd better use batch_size as 1\n    trainer = RewardModelTrainer(model=model,\n                                 strategy=strategy,\n                                 optim=optim,\n                                 train_dataset=train_dataset,\n                                 eval_dataset=eval_dataset,\n                                 batch_size=args.batch_size,\n                                 max_epochs=args.max_epochs)\n\n    trainer.fit(use_lora=args.lora_rank)\n\n    # save model checkpoint after fitting on only rank0\n    strategy.save_model(model, 'rm_checkpoint.pt', only_rank0=True)\n    # save optimizer checkpoint on all ranks\n    strategy.save_optimizer(optim, 'rm_optim_checkpoint_%d.pt' % (torch.cuda.current_device()), only_rank0=False)\n\n\nif __name__ == '__main__':\n    parser = argparse.ArgumentParser()\n    parser.add_argument('--strategy',\n                        choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'],\n                        default='naive')\n    parser.add_argument('--model', choices=['gpt2', 'bloom', 'opt'], default='bloom')\n    parser.add_argument('--pretrain', type=str, default=None)\n    parser.add_argument('--dataset', type=str, default='Dahoas/rm-static')\n    parser.add_argument('--save_path', type=str, default='rm_ckpt.pth')\n    parser.add_argument('--max_epochs', type=int, default=10)\n    parser.add_argument('--batch_size', type=int, default=4)\n    parser.add_argument('--lora_rank', type=int, default=0, help=\"low-rank adaptation matrices rank\")\n    args = parser.parse_args()\n    train(args)\n", "path": "applications/ChatGPT/examples/train_reward_model.py"}], "after_files": [{"content": "from abc import ABC\n\nimport loralib as lora\nimport torch\nfrom chatgpt.dataset import RewardDataset\nfrom chatgpt.nn import PairWiseLoss\nfrom torch.optim import Adam, Optimizer\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\n\nfrom .strategies import Strategy\nfrom .utils import is_rank_0\n\n\nclass RewardModelTrainer(ABC):\n    \"\"\"\n        Trainer to use while training reward model.\n\n    Args:\n        model (torch.nn.Module): the model to train\n        strategy (Strategy): the strategy to use for training\n        optim(Optimizer): the optimizer to use for training\n        train_dataset (RewardDataset): the dataset to use for training\n        eval_dataset (RewardDataset): the dataset to use for evaluation\n        batch_size (int, defaults to 1): the batch size while training\n        max_epochs (int, defaults to 2): the number of epochs to train\n        optim_kwargs (dict, defaults to {'lr':1e-4}): the kwargs to use while initializing optimizer\n    \"\"\"\n\n    def __init__(\n        self,\n        model,\n        strategy: Strategy,\n        optim: Optimizer,\n        train_dataset: RewardDataset,\n        eval_dataset: RewardDataset,\n        batch_size: int = 1,\n        max_epochs: int = 2,\n    ) -> None:\n        super().__init__()\n        self.strategy = strategy\n        self.epochs = max_epochs\n        self.train_dataloader = DataLoader(train_dataset, batch_size=batch_size)\n        self.eval_dataloader = DataLoader(eval_dataset, batch_size=batch_size)\n\n        self.model = strategy.setup_model(model)\n        if \"DDP\" in str(self.strategy):\n            self.model = self.model.module\n        self.loss_fn = PairWiseLoss()\n        self.optimizer = strategy.setup_optimizer(optim, self.model)\n\n    def fit(self, use_lora):\n        epoch_bar = tqdm(range(self.epochs), desc='Train epoch', disable=not is_rank_0())\n        for epoch in range(self.epochs):\n            step_bar = tqdm(range(self.train_dataloader.__len__()),\n                            desc='Train step of epoch %d' % epoch,\n                            disable=not is_rank_0())\n            # train\n            if use_lora > 0:\n                print(\"Using Lora\")\n                lora.mark_only_lora_as_trainable(self.model.model)\n\n            else:\n                self.model.train()\n            for chosen_ids, c_mask, reject_ids, r_mask in self.train_dataloader:\n                chosen_ids = chosen_ids.squeeze(1).cuda()\n                c_mask = c_mask.squeeze(1).cuda()\n                reject_ids = reject_ids.squeeze(1).cuda()\n                r_mask = r_mask.squeeze(1).cuda()\n                chosen_reward = self.model(chosen_ids, attention_mask=c_mask)\n                reject_reward = self.model(reject_ids, attention_mask=r_mask)\n                loss = self.loss_fn(chosen_reward, reject_reward)\n                self.strategy.backward(loss, self.model, self.optimizer)\n                self.strategy.optimizer_step(self.optimizer)\n                self.optimizer.zero_grad()\n                step_bar.update()\n                step_bar.set_postfix({'loss': loss.item()})\n\n            # eval\n            self.model.eval()\n            with torch.no_grad():\n                dist = 0\n                loss_sum = 0\n                for chosen_ids, c_mask, reject_ids, r_mask in self.eval_dataloader:\n                    chosen_ids = chosen_ids.squeeze(1).cuda()\n                    c_mask = c_mask.squeeze(1).cuda()\n                    reject_ids = reject_ids.squeeze(1).cuda()\n                    r_mask = r_mask.squeeze(1).cuda()\n                    chosen_reward = self.model(chosen_ids, attention_mask=c_mask)\n                    reject_reward = self.model(reject_ids, attention_mask=r_mask)\n                    dist += (chosen_reward - reject_reward).mean().item()\n                    loss = self.loss_fn(chosen_reward, reject_reward)\n                    loss_sum += loss.item()\n                dist_mean = dist / self.eval_dataloader.__len__()\n                loss_mean = loss_sum / self.eval_dataloader.__len__()\n            epoch_bar.update()\n            step_bar.set_postfix({'loss': loss_mean, 'dist_mean': dist_mean})\n            step_bar.close()\n", "path": "applications/ChatGPT/chatgpt/trainer/rm.py"}, {"content": "from typing import Optional\n\nimport torch\nimport torch.nn as nn\n\nfrom .lora import LoRAModule\n\n\nclass RewardModel(LoRAModule):\n    \"\"\"\n    Reward model base class.\n\n    Args:\n        model (nn.Module): Reward model.\n        value_head (nn.Module): Value head to get reward score.\n        lora_rank (int): LoRA rank.\n        lora_train_bias (str): LoRA bias training mode.\n    \"\"\"\n\n    def __init__(self,\n                 model: nn.Module,\n                 value_head: Optional[nn.Module] = None,\n                 lora_rank: int = 0,\n                 lora_train_bias: str = 'none') -> None:\n        super().__init__(lora_rank=lora_rank, lora_train_bias=lora_train_bias)\n        self.model = model\n        if value_head is not None:\n            if value_head.out_features != 1:\n                raise ValueError(\"The value head of reward model's output dim should be 1!\")\n            self.value_head = value_head\n\n        else:\n            self.value_head = nn.Linear(model.config.n_embd, 1)\n        self.convert_to_lora()\n\n    def forward(self, sequences: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:\n        outputs = self.model(sequences, attention_mask=attention_mask)\n        last_hidden_states = outputs['last_hidden_state']\n        values = self.value_head(last_hidden_states)[:, :-1]\n        value = values.mean(dim=1).squeeze(1)    # ensure shape is (B)\n        return value\n", "path": "applications/ChatGPT/chatgpt/nn/reward_model.py"}, {"content": "import argparse\n\nimport loralib as lora\nimport torch\nfrom chatgpt.dataset import RewardDataset\nfrom chatgpt.nn import BLOOMRM, GPTRM, OPTRM\nfrom chatgpt.trainer import RewardModelTrainer\nfrom chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy\nfrom datasets import load_dataset\nfrom torch.optim import Adam\nfrom transformers import AutoTokenizer, BloomTokenizerFast\nfrom transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer\n\nfrom colossalai.nn.optimizer import HybridAdam\n\n\ndef train(args):\n    # configure strategy\n    if args.strategy == 'naive':\n        strategy = NaiveStrategy()\n    elif args.strategy == 'ddp':\n        strategy = DDPStrategy()\n    elif args.strategy == 'colossalai_gemini':\n        strategy = ColossalAIStrategy(stage=3, placement_policy='cuda')\n    elif args.strategy == 'colossalai_zero2':\n        strategy = ColossalAIStrategy(stage=2, placement_policy='cuda')\n    else:\n        raise ValueError(f'Unsupported strategy \"{args.strategy}\"')\n\n    # configure model\n    with strategy.model_init_context():\n        if args.model == 'bloom':\n            model = BLOOMRM(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()\n        elif args.model == 'opt':\n            model = OPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()\n        elif args.model == 'gpt2':\n            model = GPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()\n        else:\n            raise ValueError(f'Unsupported model \"{args.model}\"')\n\n    # configure tokenizer\n    if args.model == 'gpt2':\n        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n        tokenizer.pad_token = tokenizer.eos_token\n    elif args.model == 'bloom':\n        tokenizer = BloomTokenizerFast.from_pretrained(args.pretrain)\n        tokenizer.pad_token = tokenizer.eos_token\n    elif args.model == 'opt':\n        tokenizer = AutoTokenizer.from_pretrained(\"facebook/opt-350m\")\n    else:\n        raise ValueError(f'Unsupported model \"{args.model}\"')\n    tokenizer.pad_token = tokenizer.eos_token\n\n    max_len = 512\n\n    # configure optimizer\n    if args.strategy.startswith('colossalai'):\n        optim = HybridAdam(model.parameters(), lr=5e-5)\n    else:\n        optim = Adam(model.parameters(), lr=5e-5)\n\n    # prepare for data and dataset\n    data = load_dataset(args.dataset)\n    train_data = data[\"train\"]\n    eval_data = data['test']\n    train_dataset = RewardDataset(train_data, tokenizer, max_len)\n    eval_dataset = RewardDataset(eval_data, tokenizer, max_len)\n\n    # batch_size here is expected to be C(k,2), k means # response of each prompt\n    # be limited with the format of dataset 'Dahoas/rm-static', we'd better use batch_size as 1\n    trainer = RewardModelTrainer(model=model,\n                                 strategy=strategy,\n                                 optim=optim,\n                                 train_dataset=train_dataset,\n                                 eval_dataset=eval_dataset,\n                                 batch_size=args.batch_size,\n                                 max_epochs=args.max_epochs)\n\n    trainer.fit(use_lora=args.lora_rank)\n\n    # save model checkpoint after fitting on only rank0\n    strategy.save_model(model, 'rm_checkpoint.pt', only_rank0=True)\n    # save optimizer checkpoint on all ranks\n    strategy.save_optimizer(optim, 'rm_optim_checkpoint_%d.pt' % (torch.cuda.current_device()), only_rank0=False)\n\n\nif __name__ == '__main__':\n    parser = argparse.ArgumentParser()\n    parser.add_argument('--strategy',\n                        choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'],\n                        default='naive')\n    parser.add_argument('--model', choices=['gpt2', 'bloom', 'opt'], default='bloom')\n    parser.add_argument('--pretrain', type=str, default=None)\n    parser.add_argument('--dataset', type=str, default='Dahoas/rm-static')\n    parser.add_argument('--save_path', type=str, default='rm_ckpt.pth')\n    parser.add_argument('--max_epochs', type=int, default=1)\n    parser.add_argument('--batch_size', type=int, default=4)\n    parser.add_argument('--lora_rank', type=int, default=0, help=\"low-rank adaptation matrices rank\")\n    args = parser.parse_args()\n    train(args)\n", "path": "applications/ChatGPT/examples/train_reward_model.py"}]}