We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[DOC]: Clarify More about Datasets for Chat Training Stage 3
### 📚 The doc issue

As raised by some users, the current guide to training Stages 3 regarding the required datasets is yet clear. Some user has composed a [guide](https://juejin.cn/post/7222250219570757690) which still confuses the datasets. 

Some clarification is needed in this [section](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat/examples#stage3---training-model-using-prompts-with-rl), regarding how to download and set `prompt_path` and `pretrain_dataset`.

Below is an example user query I received on Slack.

```
torchrun --standalone --nproc_per_node=4 train_prompts.py \
         --pretrain "bigscience/bloom-560m" \
         --model 'bloom' \
         --strategy colossalai_zero2 \
         --prompt_path /data/chenhao/train/ColossalAI/prompt_dataset/data.json \   -------Where is the data.json data
         --pretrain_dataset /data/chenhao/train/ColossalAI/pretrain_dataset/data.json \    -----------Where is the data.json data
         --rm_pretrain /data/chenhao/train/ColossalAI/Coati-7B \
         --rm_path /data/chenhao/train/ColossalAI/rmstatic.pt \
         --train_batch_size 4 \
         --experience_batch_size 4 \
         --max_epochs 1 \
         --num_episodes 1
```



[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/Chat/examples/train_prompts.py`
Content:
```
1 import argparse
2 
3 import pandas as pd
4 import torch
5 import torch.distributed as dist
6 from coati.dataset import DataCollatorForSupervisedDataset, PromptDataset, SupervisedDataset
7 from coati.models.bloom import BLOOMRM, BLOOMActor, BLOOMCritic
8 from coati.models.gpt import GPTRM, GPTActor, GPTCritic
9 from coati.models.llama import LlamaActor, LlamaCritic, LlamaRM
10 from coati.models.opt import OPTRM, OPTActor, OPTCritic
11 from coati.models.roberta import RoBERTaActor, RoBERTaCritic, RoBERTaRM
12 from coati.trainer import PPOTrainer
13 from coati.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy
14 from coati.utils import prepare_llama_tokenizer_and_embedding
15 from torch.optim import Adam
16 from torch.utils.data import DataLoader
17 from torch.utils.data.distributed import DistributedSampler
18 from transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer, RobertaTokenizer
19 
20 from colossalai.nn.optimizer import HybridAdam
21 
22 
23 def main(args):
24     # configure strategy
25     if args.strategy == 'naive':
26         strategy = NaiveStrategy()
27     elif args.strategy == 'ddp':
28         strategy = DDPStrategy()
29     elif args.strategy == 'colossalai_gemini':
30         strategy = ColossalAIStrategy(stage=3, placement_policy='cuda', initial_scale=2**5)
31     elif args.strategy == 'colossalai_zero2':
32         strategy = ColossalAIStrategy(stage=2, placement_policy='cuda')
33     else:
34         raise ValueError(f'Unsupported strategy "{args.strategy}"')
35 
36     if args.rm_path is not None:
37         state_dict = torch.load(args.rm_path, map_location='cpu')
38 
39     # configure model
40     if args.model == 'gpt2':
41         initial_model = GPTActor(pretrained=args.pretrain)
42     elif args.model == 'bloom':
43         initial_model = BLOOMActor(pretrained=args.pretrain)
44     elif args.model == 'opt':
45         initial_model = OPTActor(pretrained=args.pretrain)
46     elif args.model == 'llama':
47         initial_model = LlamaActor(pretrained=args.pretrain)
48     elif args.model == 'roberta':
49         initial_model = RoBERTaActor(pretrained=args.pretrain)
50     else:
51         raise ValueError(f'Unsupported actor model "{args.model}"')
52 
53     if args.rm_model == None:
54         rm_model_name = args.model
55     else:
56         rm_model_name = args.rm_model
57 
58     if rm_model_name == 'gpt2':
59         reward_model = GPTRM(pretrained=args.rm_pretrain)
60     elif rm_model_name == 'bloom':
61         reward_model = BLOOMRM(pretrained=args.rm_pretrain)
62     elif rm_model_name == 'opt':
63         reward_model = OPTRM(pretrained=args.rm_pretrain)
64     elif rm_model_name == 'llama':
65         reward_model = LlamaRM(pretrained=args.rm_pretrain)
66     elif rm_model_name == 'roberta':
67         reward_model = RoBERTaRM(pretrained=args.rm_pretrain)
68     else:
69         raise ValueError(f'Unsupported reward model "{rm_model_name}"')
70 
71     if args.rm_path is not None:
72         reward_model.load_state_dict(state_dict)
73 
74     initial_model.to(torch.float16).to(torch.cuda.current_device())
75     reward_model.to(torch.float16).to(torch.cuda.current_device())
76 
77     with strategy.model_init_context():
78         if args.model == 'gpt2':
79             actor = GPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
80         elif args.model == 'bloom':
81             actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
82         elif args.model == 'opt':
83             actor = OPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
84         elif args.model == 'llama':
85             actor = LlamaActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
86         elif args.model == 'roberta':
87             actor = RoBERTaActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
88         else:
89             raise ValueError(f'Unsupported actor model "{args.model}"')
90 
91         if rm_model_name == 'gpt2':
92             critic = GPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
93         elif rm_model_name == 'bloom':
94             critic = BLOOMCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
95         elif rm_model_name == 'opt':
96             critic = OPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
97         elif rm_model_name == 'llama':
98             critic = LlamaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
99         elif rm_model_name == 'roberta':
100             critic = RoBERTaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
101         else:
102             raise ValueError(f'Unsupported reward model "{rm_model_name}"')
103 
104         if args.rm_path is not None:
105             critic.load_state_dict(state_dict)
106             del state_dict
107 
108     if args.strategy != 'colossalai_gemini':
109         critic.to(torch.float16).to(torch.cuda.current_device())
110         actor.to(torch.float16).to(torch.cuda.current_device())
111 
112     # configure optimizer
113     if args.strategy.startswith('colossalai'):
114         actor_optim = HybridAdam(actor.parameters(), lr=1e-7)
115         critic_optim = HybridAdam(critic.parameters(), lr=1e-7)
116     else:
117         actor_optim = Adam(actor.parameters(), lr=1e-7)
118         critic_optim = Adam(critic.parameters(), lr=1e-7)
119 
120     # configure tokenizer
121     if args.model == 'gpt2':
122         tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
123     elif args.model == 'bloom':
124         tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m')
125     elif args.model == 'opt':
126         tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
127     elif args.model == 'llama':
128         tokenizer = LlamaTokenizer.from_pretrained(args.pretrain)
129         tokenizer.eos_token = '<\s>'
130     elif args.model == 'roberta':
131         tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
132     else:
133         raise ValueError(f'Unsupported model "{args.model}"')
134 
135     if args.model == 'llama':
136         tokenizer = prepare_llama_tokenizer_and_embedding(tokenizer, actor)
137     else:
138         tokenizer.pad_token = tokenizer.eos_token
139 
140     data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
141 
142     prompt_dataset = PromptDataset(tokenizer=tokenizer, data_path=args.prompt_path, max_datasets_size=16384)
143     if dist.is_initialized() and dist.get_world_size() > 1:
144         prompt_sampler = DistributedSampler(prompt_dataset, shuffle=True, seed=42, drop_last=True)
145     else:
146         prompt_sampler = None
147     prompt_dataloader = DataLoader(prompt_dataset,
148                                    shuffle=(prompt_sampler is None),
149                                    sampler=prompt_sampler,
150                                    batch_size=args.experience_batch_size)
151 
152     pretrain_dataset = SupervisedDataset(tokenizer=tokenizer,
153                                          data_path=args.pretrain_dataset,
154                                          max_datasets_size=16384,
155                                          max_length=args.max_input_len)
156     if dist.is_initialized() and dist.get_world_size() > 1:
157         pretrain_sampler = DistributedSampler(pretrain_dataset, shuffle=True, seed=42, drop_last=True)
158     else:
159         pretrain_sampler = None
160     pretrain_dataloader = DataLoader(pretrain_dataset,
161                                      shuffle=(pretrain_sampler is None),
162                                      sampler=pretrain_sampler,
163                                      batch_size=args.ptx_batch_size,
164                                      collate_fn=data_collator)
165 
166     (actor, actor_optim), (critic, critic_optim) = strategy.prepare((actor, actor_optim), (critic, critic_optim))
167 
168     # configure trainer
169     trainer = PPOTrainer(
170         strategy,
171         actor,
172         critic,
173         reward_model,
174         initial_model,
175         actor_optim,
176         critic_optim,
177         kl_coef=args.kl_coef,
178         ptx_coef=args.ptx_coef,
179         max_epochs=args.max_epochs,
180         train_batch_size=args.train_batch_size,
181         max_length=args.max_seq_len,
182         use_cache=True,
183         do_sample=True,
184         temperature=1.0,
185         top_k=50,
186         pad_token_id=tokenizer.pad_token_id,
187         eos_token_id=tokenizer.eos_token_id,
188     )
189 
190     trainer.fit(prompt_dataloader=prompt_dataloader,
191                 pretrain_dataloader=pretrain_dataloader,
192                 num_episodes=args.num_episodes,
193                 max_timesteps=args.max_timesteps,
194                 update_timesteps=args.update_timesteps)
195 
196     # save model checkpoint after fitting
197     trainer.save_model(args.save_path, only_rank0=True, tokenizer=tokenizer)
198     # save optimizer checkpoint on all ranks
199     if args.need_optim_ckpt:
200         strategy.save_optimizer(actor_optim,
201                                 'actor_optim_checkpoint_prompts_%d.pt' % (torch.cuda.current_device()),
202                                 only_rank0=False)
203 
204 
205 if __name__ == '__main__':
206     parser = argparse.ArgumentParser()
207     parser.add_argument('--prompt_path', type=str, default=None, help='path to the prompt dataset')
208     parser.add_argument('--pretrain_dataset', type=str, default=None, help='path to the pretrained dataset')
209     parser.add_argument('--strategy',
210                         choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'],
211                         default='naive',
212                         help='strategy to use')
213     parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama', 'roberta'])
214     parser.add_argument('--pretrain', type=str, default=None)
215     parser.add_argument('--rm_model', default=None, choices=['gpt2', 'bloom', 'opt', 'llama', 'roberta'])
216     parser.add_argument('--rm_path', type=str, default=None)
217     parser.add_argument('--rm_pretrain', type=str, default=None)
218     parser.add_argument('--save_path', type=str, default='actor_checkpoint_prompts')
219     parser.add_argument('--need_optim_ckpt', type=bool, default=False)
220     parser.add_argument('--num_episodes', type=int, default=10)
221     parser.add_argument('--max_timesteps', type=int, default=10)
222     parser.add_argument('--update_timesteps', type=int, default=10)
223     parser.add_argument('--max_epochs', type=int, default=5)
224     parser.add_argument('--train_batch_size', type=int, default=8)
225     parser.add_argument('--ptx_batch_size', type=int, default=1)
226     parser.add_argument('--experience_batch_size', type=int, default=8)
227     parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank")
228     parser.add_argument('--kl_coef', type=float, default=0.1)
229     parser.add_argument('--ptx_coef', type=float, default=0.9)
230     parser.add_argument('--max_input_len', type=int, default=96)
231     parser.add_argument('--max_seq_len', type=int, default=128)
232     args = parser.parse_args()
233     main(args)
234
```
--- 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/Chat/examples/train_prompts.py b/applications/Chat/examples/train_prompts.py
--- a/applications/Chat/examples/train_prompts.py
+++ b/applications/Chat/examples/train_prompts.py
@@ -139,7 +139,7 @@
 
     data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
 
-    prompt_dataset = PromptDataset(tokenizer=tokenizer, data_path=args.prompt_path, max_datasets_size=16384)
+    prompt_dataset = PromptDataset(tokenizer=tokenizer, data_path=args.prompt_dataset, max_datasets_size=16384)
     if dist.is_initialized() and dist.get_world_size() > 1:
         prompt_sampler = DistributedSampler(prompt_dataset, shuffle=True, seed=42, drop_last=True)
     else:
@@ -204,7 +204,7 @@
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('--prompt_path', type=str, default=None, help='path to the prompt dataset')
+    parser.add_argument('--prompt_dataset', type=str, default=None, help='path to the prompt dataset')
     parser.add_argument('--pretrain_dataset', type=str, default=None, help='path to the pretrained dataset')
     parser.add_argument('--strategy',
                         choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'],

{"golden_diff": "diff --git a/applications/Chat/examples/train_prompts.py b/applications/Chat/examples/train_prompts.py\n--- a/applications/Chat/examples/train_prompts.py\n+++ b/applications/Chat/examples/train_prompts.py\n@@ -139,7 +139,7 @@\n \n     data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)\n \n-    prompt_dataset = PromptDataset(tokenizer=tokenizer, data_path=args.prompt_path, max_datasets_size=16384)\n+    prompt_dataset = PromptDataset(tokenizer=tokenizer, data_path=args.prompt_dataset, max_datasets_size=16384)\n     if dist.is_initialized() and dist.get_world_size() > 1:\n         prompt_sampler = DistributedSampler(prompt_dataset, shuffle=True, seed=42, drop_last=True)\n     else:\n@@ -204,7 +204,7 @@\n \n if __name__ == '__main__':\n     parser = argparse.ArgumentParser()\n-    parser.add_argument('--prompt_path', type=str, default=None, help='path to the prompt dataset')\n+    parser.add_argument('--prompt_dataset', type=str, default=None, help='path to the prompt dataset')\n     parser.add_argument('--pretrain_dataset', type=str, default=None, help='path to the pretrained dataset')\n     parser.add_argument('--strategy',\n                         choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'],\n", "issue": "[DOC]: Clarify More about Datasets for Chat Training Stage 3\n### \ud83d\udcda The doc issue\n\nAs raised by some users, the current guide to training Stages 3 regarding the required datasets is yet clear. Some user has composed a [guide](https://juejin.cn/post/7222250219570757690) which still confuses the datasets. \r\n\r\nSome clarification is needed in this [section](https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat/examples#stage3---training-model-using-prompts-with-rl), regarding how to download and set `prompt_path` and `pretrain_dataset`.\r\n\r\nBelow is an example user query I received on Slack.\r\n\r\n```\r\ntorchrun --standalone --nproc_per_node=4 train_prompts.py \\\r\n         --pretrain \"bigscience/bloom-560m\" \\\r\n         --model 'bloom' \\\r\n         --strategy colossalai_zero2 \\\r\n         --prompt_path /data/chenhao/train/ColossalAI/prompt_dataset/data.json \\   -------Where is the data.json data\r\n         --pretrain_dataset /data/chenhao/train/ColossalAI/pretrain_dataset/data.json \\    -----------Where is the data.json data\r\n         --rm_pretrain /data/chenhao/train/ColossalAI/Coati-7B \\\r\n         --rm_path /data/chenhao/train/ColossalAI/rmstatic.pt \\\r\n         --train_batch_size 4 \\\r\n         --experience_batch_size 4 \\\r\n         --max_epochs 1 \\\r\n         --num_episodes 1\r\n```\r\n\r\n\r\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "import argparse\n\nimport pandas as pd\nimport torch\nimport torch.distributed as dist\nfrom coati.dataset import DataCollatorForSupervisedDataset, PromptDataset, SupervisedDataset\nfrom coati.models.bloom import BLOOMRM, BLOOMActor, BLOOMCritic\nfrom coati.models.gpt import GPTRM, GPTActor, GPTCritic\nfrom coati.models.llama import LlamaActor, LlamaCritic, LlamaRM\nfrom coati.models.opt import OPTRM, OPTActor, OPTCritic\nfrom coati.models.roberta import RoBERTaActor, RoBERTaCritic, RoBERTaRM\nfrom coati.trainer import PPOTrainer\nfrom coati.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy\nfrom coati.utils import prepare_llama_tokenizer_and_embedding\nfrom torch.optim import Adam\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data.distributed import DistributedSampler\nfrom transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer, RobertaTokenizer\n\nfrom colossalai.nn.optimizer import HybridAdam\n\n\ndef main(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', initial_scale=2**5)\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    if args.rm_path is not None:\n        state_dict = torch.load(args.rm_path, map_location='cpu')\n\n    # configure model\n    if args.model == 'gpt2':\n        initial_model = GPTActor(pretrained=args.pretrain)\n    elif args.model == 'bloom':\n        initial_model = BLOOMActor(pretrained=args.pretrain)\n    elif args.model == 'opt':\n        initial_model = OPTActor(pretrained=args.pretrain)\n    elif args.model == 'llama':\n        initial_model = LlamaActor(pretrained=args.pretrain)\n    elif args.model == 'roberta':\n        initial_model = RoBERTaActor(pretrained=args.pretrain)\n    else:\n        raise ValueError(f'Unsupported actor model \"{args.model}\"')\n\n    if args.rm_model == None:\n        rm_model_name = args.model\n    else:\n        rm_model_name = args.rm_model\n\n    if rm_model_name == 'gpt2':\n        reward_model = GPTRM(pretrained=args.rm_pretrain)\n    elif rm_model_name == 'bloom':\n        reward_model = BLOOMRM(pretrained=args.rm_pretrain)\n    elif rm_model_name == 'opt':\n        reward_model = OPTRM(pretrained=args.rm_pretrain)\n    elif rm_model_name == 'llama':\n        reward_model = LlamaRM(pretrained=args.rm_pretrain)\n    elif rm_model_name == 'roberta':\n        reward_model = RoBERTaRM(pretrained=args.rm_pretrain)\n    else:\n        raise ValueError(f'Unsupported reward model \"{rm_model_name}\"')\n\n    if args.rm_path is not None:\n        reward_model.load_state_dict(state_dict)\n\n    initial_model.to(torch.float16).to(torch.cuda.current_device())\n    reward_model.to(torch.float16).to(torch.cuda.current_device())\n\n    with strategy.model_init_context():\n        if args.model == 'gpt2':\n            actor = GPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank)\n        elif args.model == 'bloom':\n            actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank)\n        elif args.model == 'opt':\n            actor = OPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank)\n        elif args.model == 'llama':\n            actor = LlamaActor(pretrained=args.pretrain, lora_rank=args.lora_rank)\n        elif args.model == 'roberta':\n            actor = RoBERTaActor(pretrained=args.pretrain, lora_rank=args.lora_rank)\n        else:\n            raise ValueError(f'Unsupported actor model \"{args.model}\"')\n\n        if rm_model_name == 'gpt2':\n            critic = GPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)\n        elif rm_model_name == 'bloom':\n            critic = BLOOMCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)\n        elif rm_model_name == 'opt':\n            critic = OPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)\n        elif rm_model_name == 'llama':\n            critic = LlamaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)\n        elif rm_model_name == 'roberta':\n            critic = RoBERTaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)\n        else:\n            raise ValueError(f'Unsupported reward model \"{rm_model_name}\"')\n\n        if args.rm_path is not None:\n            critic.load_state_dict(state_dict)\n            del state_dict\n\n    if args.strategy != 'colossalai_gemini':\n        critic.to(torch.float16).to(torch.cuda.current_device())\n        actor.to(torch.float16).to(torch.cuda.current_device())\n\n    # configure optimizer\n    if args.strategy.startswith('colossalai'):\n        actor_optim = HybridAdam(actor.parameters(), lr=1e-7)\n        critic_optim = HybridAdam(critic.parameters(), lr=1e-7)\n    else:\n        actor_optim = Adam(actor.parameters(), lr=1e-7)\n        critic_optim = Adam(critic.parameters(), lr=1e-7)\n\n    # configure tokenizer\n    if args.model == 'gpt2':\n        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n    elif args.model == 'bloom':\n        tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m')\n    elif args.model == 'opt':\n        tokenizer = AutoTokenizer.from_pretrained(\"facebook/opt-350m\")\n    elif args.model == 'llama':\n        tokenizer = LlamaTokenizer.from_pretrained(args.pretrain)\n        tokenizer.eos_token = '<\\s>'\n    elif args.model == 'roberta':\n        tokenizer = RobertaTokenizer.from_pretrained(\"roberta-base\")\n    else:\n        raise ValueError(f'Unsupported model \"{args.model}\"')\n\n    if args.model == 'llama':\n        tokenizer = prepare_llama_tokenizer_and_embedding(tokenizer, actor)\n    else:\n        tokenizer.pad_token = tokenizer.eos_token\n\n    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)\n\n    prompt_dataset = PromptDataset(tokenizer=tokenizer, data_path=args.prompt_path, max_datasets_size=16384)\n    if dist.is_initialized() and dist.get_world_size() > 1:\n        prompt_sampler = DistributedSampler(prompt_dataset, shuffle=True, seed=42, drop_last=True)\n    else:\n        prompt_sampler = None\n    prompt_dataloader = DataLoader(prompt_dataset,\n                                   shuffle=(prompt_sampler is None),\n                                   sampler=prompt_sampler,\n                                   batch_size=args.experience_batch_size)\n\n    pretrain_dataset = SupervisedDataset(tokenizer=tokenizer,\n                                         data_path=args.pretrain_dataset,\n                                         max_datasets_size=16384,\n                                         max_length=args.max_input_len)\n    if dist.is_initialized() and dist.get_world_size() > 1:\n        pretrain_sampler = DistributedSampler(pretrain_dataset, shuffle=True, seed=42, drop_last=True)\n    else:\n        pretrain_sampler = None\n    pretrain_dataloader = DataLoader(pretrain_dataset,\n                                     shuffle=(pretrain_sampler is None),\n                                     sampler=pretrain_sampler,\n                                     batch_size=args.ptx_batch_size,\n                                     collate_fn=data_collator)\n\n    (actor, actor_optim), (critic, critic_optim) = strategy.prepare((actor, actor_optim), (critic, critic_optim))\n\n    # configure trainer\n    trainer = PPOTrainer(\n        strategy,\n        actor,\n        critic,\n        reward_model,\n        initial_model,\n        actor_optim,\n        critic_optim,\n        kl_coef=args.kl_coef,\n        ptx_coef=args.ptx_coef,\n        max_epochs=args.max_epochs,\n        train_batch_size=args.train_batch_size,\n        max_length=args.max_seq_len,\n        use_cache=True,\n        do_sample=True,\n        temperature=1.0,\n        top_k=50,\n        pad_token_id=tokenizer.pad_token_id,\n        eos_token_id=tokenizer.eos_token_id,\n    )\n\n    trainer.fit(prompt_dataloader=prompt_dataloader,\n                pretrain_dataloader=pretrain_dataloader,\n                num_episodes=args.num_episodes,\n                max_timesteps=args.max_timesteps,\n                update_timesteps=args.update_timesteps)\n\n    # save model checkpoint after fitting\n    trainer.save_model(args.save_path, only_rank0=True, tokenizer=tokenizer)\n    # save optimizer checkpoint on all ranks\n    if args.need_optim_ckpt:\n        strategy.save_optimizer(actor_optim,\n                                'actor_optim_checkpoint_prompts_%d.pt' % (torch.cuda.current_device()),\n                                only_rank0=False)\n\n\nif __name__ == '__main__':\n    parser = argparse.ArgumentParser()\n    parser.add_argument('--prompt_path', type=str, default=None, help='path to the prompt dataset')\n    parser.add_argument('--pretrain_dataset', type=str, default=None, help='path to the pretrained dataset')\n    parser.add_argument('--strategy',\n                        choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'],\n                        default='naive',\n                        help='strategy to use')\n    parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama', 'roberta'])\n    parser.add_argument('--pretrain', type=str, default=None)\n    parser.add_argument('--rm_model', default=None, choices=['gpt2', 'bloom', 'opt', 'llama', 'roberta'])\n    parser.add_argument('--rm_path', type=str, default=None)\n    parser.add_argument('--rm_pretrain', type=str, default=None)\n    parser.add_argument('--save_path', type=str, default='actor_checkpoint_prompts')\n    parser.add_argument('--need_optim_ckpt', type=bool, default=False)\n    parser.add_argument('--num_episodes', type=int, default=10)\n    parser.add_argument('--max_timesteps', type=int, default=10)\n    parser.add_argument('--update_timesteps', type=int, default=10)\n    parser.add_argument('--max_epochs', type=int, default=5)\n    parser.add_argument('--train_batch_size', type=int, default=8)\n    parser.add_argument('--ptx_batch_size', type=int, default=1)\n    parser.add_argument('--experience_batch_size', type=int, default=8)\n    parser.add_argument('--lora_rank', type=int, default=0, help=\"low-rank adaptation matrices rank\")\n    parser.add_argument('--kl_coef', type=float, default=0.1)\n    parser.add_argument('--ptx_coef', type=float, default=0.9)\n    parser.add_argument('--max_input_len', type=int, default=96)\n    parser.add_argument('--max_seq_len', type=int, default=128)\n    args = parser.parse_args()\n    main(args)\n", "path": "applications/Chat/examples/train_prompts.py"}], "after_files": [{"content": "import argparse\n\nimport pandas as pd\nimport torch\nimport torch.distributed as dist\nfrom coati.dataset import DataCollatorForSupervisedDataset, PromptDataset, SupervisedDataset\nfrom coati.models.bloom import BLOOMRM, BLOOMActor, BLOOMCritic\nfrom coati.models.gpt import GPTRM, GPTActor, GPTCritic\nfrom coati.models.llama import LlamaActor, LlamaCritic, LlamaRM\nfrom coati.models.opt import OPTRM, OPTActor, OPTCritic\nfrom coati.models.roberta import RoBERTaActor, RoBERTaCritic, RoBERTaRM\nfrom coati.trainer import PPOTrainer\nfrom coati.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy\nfrom coati.utils import prepare_llama_tokenizer_and_embedding\nfrom torch.optim import Adam\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data.distributed import DistributedSampler\nfrom transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer, RobertaTokenizer\n\nfrom colossalai.nn.optimizer import HybridAdam\n\n\ndef main(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', initial_scale=2**5)\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    if args.rm_path is not None:\n        state_dict = torch.load(args.rm_path, map_location='cpu')\n\n    # configure model\n    if args.model == 'gpt2':\n        initial_model = GPTActor(pretrained=args.pretrain)\n    elif args.model == 'bloom':\n        initial_model = BLOOMActor(pretrained=args.pretrain)\n    elif args.model == 'opt':\n        initial_model = OPTActor(pretrained=args.pretrain)\n    elif args.model == 'llama':\n        initial_model = LlamaActor(pretrained=args.pretrain)\n    elif args.model == 'roberta':\n        initial_model = RoBERTaActor(pretrained=args.pretrain)\n    else:\n        raise ValueError(f'Unsupported actor model \"{args.model}\"')\n\n    if args.rm_model == None:\n        rm_model_name = args.model\n    else:\n        rm_model_name = args.rm_model\n\n    if rm_model_name == 'gpt2':\n        reward_model = GPTRM(pretrained=args.rm_pretrain)\n    elif rm_model_name == 'bloom':\n        reward_model = BLOOMRM(pretrained=args.rm_pretrain)\n    elif rm_model_name == 'opt':\n        reward_model = OPTRM(pretrained=args.rm_pretrain)\n    elif rm_model_name == 'llama':\n        reward_model = LlamaRM(pretrained=args.rm_pretrain)\n    elif rm_model_name == 'roberta':\n        reward_model = RoBERTaRM(pretrained=args.rm_pretrain)\n    else:\n        raise ValueError(f'Unsupported reward model \"{rm_model_name}\"')\n\n    if args.rm_path is not None:\n        reward_model.load_state_dict(state_dict)\n\n    initial_model.to(torch.float16).to(torch.cuda.current_device())\n    reward_model.to(torch.float16).to(torch.cuda.current_device())\n\n    with strategy.model_init_context():\n        if args.model == 'gpt2':\n            actor = GPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank)\n        elif args.model == 'bloom':\n            actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank)\n        elif args.model == 'opt':\n            actor = OPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank)\n        elif args.model == 'llama':\n            actor = LlamaActor(pretrained=args.pretrain, lora_rank=args.lora_rank)\n        elif args.model == 'roberta':\n            actor = RoBERTaActor(pretrained=args.pretrain, lora_rank=args.lora_rank)\n        else:\n            raise ValueError(f'Unsupported actor model \"{args.model}\"')\n\n        if rm_model_name == 'gpt2':\n            critic = GPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)\n        elif rm_model_name == 'bloom':\n            critic = BLOOMCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)\n        elif rm_model_name == 'opt':\n            critic = OPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)\n        elif rm_model_name == 'llama':\n            critic = LlamaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)\n        elif rm_model_name == 'roberta':\n            critic = RoBERTaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)\n        else:\n            raise ValueError(f'Unsupported reward model \"{rm_model_name}\"')\n\n        if args.rm_path is not None:\n            critic.load_state_dict(state_dict)\n            del state_dict\n\n    if args.strategy != 'colossalai_gemini':\n        critic.to(torch.float16).to(torch.cuda.current_device())\n        actor.to(torch.float16).to(torch.cuda.current_device())\n\n    # configure optimizer\n    if args.strategy.startswith('colossalai'):\n        actor_optim = HybridAdam(actor.parameters(), lr=1e-7)\n        critic_optim = HybridAdam(critic.parameters(), lr=1e-7)\n    else:\n        actor_optim = Adam(actor.parameters(), lr=1e-7)\n        critic_optim = Adam(critic.parameters(), lr=1e-7)\n\n    # configure tokenizer\n    if args.model == 'gpt2':\n        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n    elif args.model == 'bloom':\n        tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m')\n    elif args.model == 'opt':\n        tokenizer = AutoTokenizer.from_pretrained(\"facebook/opt-350m\")\n    elif args.model == 'llama':\n        tokenizer = LlamaTokenizer.from_pretrained(args.pretrain)\n        tokenizer.eos_token = '<\\s>'\n    elif args.model == 'roberta':\n        tokenizer = RobertaTokenizer.from_pretrained(\"roberta-base\")\n    else:\n        raise ValueError(f'Unsupported model \"{args.model}\"')\n\n    if args.model == 'llama':\n        tokenizer = prepare_llama_tokenizer_and_embedding(tokenizer, actor)\n    else:\n        tokenizer.pad_token = tokenizer.eos_token\n\n    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)\n\n    prompt_dataset = PromptDataset(tokenizer=tokenizer, data_path=args.prompt_dataset, max_datasets_size=16384)\n    if dist.is_initialized() and dist.get_world_size() > 1:\n        prompt_sampler = DistributedSampler(prompt_dataset, shuffle=True, seed=42, drop_last=True)\n    else:\n        prompt_sampler = None\n    prompt_dataloader = DataLoader(prompt_dataset,\n                                   shuffle=(prompt_sampler is None),\n                                   sampler=prompt_sampler,\n                                   batch_size=args.experience_batch_size)\n\n    pretrain_dataset = SupervisedDataset(tokenizer=tokenizer,\n                                         data_path=args.pretrain_dataset,\n                                         max_datasets_size=16384,\n                                         max_length=args.max_input_len)\n    if dist.is_initialized() and dist.get_world_size() > 1:\n        pretrain_sampler = DistributedSampler(pretrain_dataset, shuffle=True, seed=42, drop_last=True)\n    else:\n        pretrain_sampler = None\n    pretrain_dataloader = DataLoader(pretrain_dataset,\n                                     shuffle=(pretrain_sampler is None),\n                                     sampler=pretrain_sampler,\n                                     batch_size=args.ptx_batch_size,\n                                     collate_fn=data_collator)\n\n    (actor, actor_optim), (critic, critic_optim) = strategy.prepare((actor, actor_optim), (critic, critic_optim))\n\n    # configure trainer\n    trainer = PPOTrainer(\n        strategy,\n        actor,\n        critic,\n        reward_model,\n        initial_model,\n        actor_optim,\n        critic_optim,\n        kl_coef=args.kl_coef,\n        ptx_coef=args.ptx_coef,\n        max_epochs=args.max_epochs,\n        train_batch_size=args.train_batch_size,\n        max_length=args.max_seq_len,\n        use_cache=True,\n        do_sample=True,\n        temperature=1.0,\n        top_k=50,\n        pad_token_id=tokenizer.pad_token_id,\n        eos_token_id=tokenizer.eos_token_id,\n    )\n\n    trainer.fit(prompt_dataloader=prompt_dataloader,\n                pretrain_dataloader=pretrain_dataloader,\n                num_episodes=args.num_episodes,\n                max_timesteps=args.max_timesteps,\n                update_timesteps=args.update_timesteps)\n\n    # save model checkpoint after fitting\n    trainer.save_model(args.save_path, only_rank0=True, tokenizer=tokenizer)\n    # save optimizer checkpoint on all ranks\n    if args.need_optim_ckpt:\n        strategy.save_optimizer(actor_optim,\n                                'actor_optim_checkpoint_prompts_%d.pt' % (torch.cuda.current_device()),\n                                only_rank0=False)\n\n\nif __name__ == '__main__':\n    parser = argparse.ArgumentParser()\n    parser.add_argument('--prompt_dataset', type=str, default=None, help='path to the prompt dataset')\n    parser.add_argument('--pretrain_dataset', type=str, default=None, help='path to the pretrained dataset')\n    parser.add_argument('--strategy',\n                        choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'],\n                        default='naive',\n                        help='strategy to use')\n    parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama', 'roberta'])\n    parser.add_argument('--pretrain', type=str, default=None)\n    parser.add_argument('--rm_model', default=None, choices=['gpt2', 'bloom', 'opt', 'llama', 'roberta'])\n    parser.add_argument('--rm_path', type=str, default=None)\n    parser.add_argument('--rm_pretrain', type=str, default=None)\n    parser.add_argument('--save_path', type=str, default='actor_checkpoint_prompts')\n    parser.add_argument('--need_optim_ckpt', type=bool, default=False)\n    parser.add_argument('--num_episodes', type=int, default=10)\n    parser.add_argument('--max_timesteps', type=int, default=10)\n    parser.add_argument('--update_timesteps', type=int, default=10)\n    parser.add_argument('--max_epochs', type=int, default=5)\n    parser.add_argument('--train_batch_size', type=int, default=8)\n    parser.add_argument('--ptx_batch_size', type=int, default=1)\n    parser.add_argument('--experience_batch_size', type=int, default=8)\n    parser.add_argument('--lora_rank', type=int, default=0, help=\"low-rank adaptation matrices rank\")\n    parser.add_argument('--kl_coef', type=float, default=0.1)\n    parser.add_argument('--ptx_coef', type=float, default=0.9)\n    parser.add_argument('--max_input_len', type=int, default=96)\n    parser.add_argument('--max_seq_len', type=int, default=128)\n    args = parser.parse_args()\n    main(args)\n", "path": "applications/Chat/examples/train_prompts.py"}]}