We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG]: [chat] Unable to run train_prompts.sh by single card
### 🐛 Describe the bug

Missing definition for prompt_sampler and pretrain_sampler in `examples/train_prompts.py` when dist.get_world_size() == 1.

### Environment

_No response_

--- 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 RoBERTaRM, RoBERTaActor, RoBERTaCritic
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     if args.strategy != 'colossalai_gemini':
75         initial_model.to(torch.float16).to(torch.cuda.current_device())
76         reward_model.to(torch.float16).to(torch.cuda.current_device())
77 
78     with strategy.model_init_context():
79         if args.model == 'gpt2':
80             actor = GPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
81         elif args.model == 'bloom':
82             actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
83         elif args.model == 'opt':
84             actor = OPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
85         elif args.model == 'llama':
86             actor = LlamaActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
87         elif args.model == 'roberta':
88             actor = RoBERTaActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
89         else:
90             raise ValueError(f'Unsupported actor model "{args.model}"')
91 
92         if rm_model_name == 'gpt2':
93             critic = GPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
94         elif rm_model_name == 'bloom':
95             critic = BLOOMCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
96         elif rm_model_name == 'opt':
97             critic = OPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
98         elif rm_model_name == 'llama':
99             critic = LlamaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
100         elif rm_model_name == 'roberta':
101             critic = RoBERTaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
102         else:
103             raise ValueError(f'Unsupported reward model "{rm_model_name}"')
104 
105         if args.rm_path is not None:
106             critic.load_state_dict(state_dict)
107             del state_dict
108 
109     if args.strategy != 'colossalai_gemini':
110         critic.to(torch.float16).to(torch.cuda.current_device())
111         actor.to(torch.float16).to(torch.cuda.current_device())
112 
113     # configure optimizer
114     if args.strategy.startswith('colossalai'):
115         actor_optim = HybridAdam(actor.parameters(), lr=1e-7)
116         critic_optim = HybridAdam(critic.parameters(), lr=1e-7)
117     else:
118         actor_optim = Adam(actor.parameters(), lr=1e-7)
119         critic_optim = Adam(critic.parameters(), lr=1e-7)
120 
121     # configure tokenizer
122     if args.model == 'gpt2':
123         tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
124     elif args.model == 'bloom':
125         tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m')
126     elif args.model == 'opt':
127         tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
128     elif args.model == 'llama':
129         tokenizer = LlamaTokenizer.from_pretrained(args.pretrain)
130         tokenizer.eos_token = '<\s>'
131     elif args.model == 'roberta':
132         tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
133     else:
134         raise ValueError(f'Unsupported model "{args.model}"')
135 
136     if args.model == 'llama':
137         tokenizer = prepare_llama_tokenizer_and_embedding(tokenizer, actor)
138     else:
139         tokenizer.pad_token = tokenizer.eos_token
140 
141     data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
142 
143     prompt_dataset = PromptDataset(tokenizer=tokenizer, data_path=args.prompt_path, max_datasets_size=16384)
144     if dist.is_initialized() and dist.get_world_size() > 1:
145         prompt_sampler = DistributedSampler(prompt_dataset, shuffle=True, seed=42, drop_last=True)
146     prompt_dataloader = DataLoader(prompt_dataset,
147                                    shuffle=(prompt_sampler is None),
148                                    sampler=prompt_sampler,
149                                    batch_size=args.train_batch_size)
150 
151     pretrain_dataset = SupervisedDataset(tokenizer=tokenizer, data_path=args.pretrain_dataset, max_datasets_size=16384)
152     if dist.is_initialized() and dist.get_world_size() > 1:
153         pretrain_sampler = DistributedSampler(pretrain_dataset, shuffle=True, seed=42, drop_last=True)
154     pretrain_dataloader = DataLoader(pretrain_dataset,
155                                      shuffle=(pretrain_sampler is None),
156                                      sampler=pretrain_sampler,
157                                      batch_size=args.ptx_batch_size,
158                                      collate_fn=data_collator)
159 
160     def tokenize_fn(texts):
161         # MUST padding to max length to ensure inputs of all ranks have the same length
162         # Different length may lead to hang when using gemini, as different generation steps
163         batch = tokenizer(texts, return_tensors='pt', max_length=96, padding='max_length', truncation=True)
164         return {k: v.to(torch.cuda.current_device()) for k, v in batch.items()}
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         experience_batch_size=args.experience_batch_size,
182         tokenizer=tokenize_fn,
183         max_length=128,
184         do_sample=True,
185         temperature=1.0,
186         top_k=50,
187         pad_token_id=tokenizer.pad_token_id,
188         eos_token_id=tokenizer.eos_token_id,
189     )
190 
191     trainer.fit(prompt_dataloader=prompt_dataloader,
192                 pretrain_dataloader=pretrain_dataloader,
193                 num_episodes=args.num_episodes,
194                 max_timesteps=args.max_timesteps,
195                 update_timesteps=args.update_timesteps)
196 
197     # save model checkpoint after fitting
198     trainer.save_model(args.save_path, only_rank0=True, tokenizer=tokenizer)
199     # save optimizer checkpoint on all ranks
200     if args.need_optim_ckpt:
201         strategy.save_optimizer(actor_optim,
202                                 'actor_optim_checkpoint_prompts_%d.pt' % (torch.cuda.current_device()),
203                                 only_rank0=False)
204 
205 
206 if __name__ == '__main__':
207     parser = argparse.ArgumentParser()
208     parser.add_argument('--prompt_path', type=str, default=None, help='path to the prompt dataset')
209     parser.add_argument('--pretrain_dataset', type=str, default=None, help='path to the pretrained dataset')
210     parser.add_argument('--strategy',
211                         choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'],
212                         default='naive',
213                         help='strategy to use')
214     parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama', 'roberta'])
215     parser.add_argument('--pretrain', type=str, default=None)
216     parser.add_argument('--rm_model', default=None, choices=['gpt2', 'bloom', 'opt', 'llama', 'roberta'])
217     parser.add_argument('--rm_path', type=str, default=None)
218     parser.add_argument('--rm_pretrain', type=str, default=None)
219     parser.add_argument('--save_path', type=str, default='actor_checkpoint_prompts')
220     parser.add_argument('--need_optim_ckpt', type=bool, default=False)
221     parser.add_argument('--num_episodes', type=int, default=10)
222     parser.add_argument('--max_timesteps', type=int, default=10)
223     parser.add_argument('--update_timesteps', type=int, default=10)
224     parser.add_argument('--max_epochs', type=int, default=5)
225     parser.add_argument('--train_batch_size', type=int, default=8)
226     parser.add_argument('--ptx_batch_size', type=int, default=1)
227     parser.add_argument('--experience_batch_size', type=int, default=8)
228     parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank")
229     parser.add_argument('--kl_coef', type=float, default=0.1)
230     parser.add_argument('--ptx_coef', type=float, default=0.9)
231     args = parser.parse_args()
232     main(args)
233
```
--- 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
@@ -8,7 +8,7 @@
 from coati.models.gpt import GPTRM, GPTActor, GPTCritic
 from coati.models.llama import LlamaActor, LlamaCritic, LlamaRM
 from coati.models.opt import OPTRM, OPTActor, OPTCritic
-from coati.models.roberta import RoBERTaRM, RoBERTaActor, RoBERTaCritic
+from coati.models.roberta import RoBERTaActor, RoBERTaCritic, RoBERTaRM
 from coati.trainer import PPOTrainer
 from coati.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy
 from coati.utils import prepare_llama_tokenizer_and_embedding
@@ -143,6 +143,8 @@
     prompt_dataset = PromptDataset(tokenizer=tokenizer, data_path=args.prompt_path, 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:
+        prompt_sampler = None
     prompt_dataloader = DataLoader(prompt_dataset,
                                    shuffle=(prompt_sampler is None),
                                    sampler=prompt_sampler,
@@ -151,6 +153,8 @@
     pretrain_dataset = SupervisedDataset(tokenizer=tokenizer, data_path=args.pretrain_dataset, max_datasets_size=16384)
     if dist.is_initialized() and dist.get_world_size() > 1:
         pretrain_sampler = DistributedSampler(pretrain_dataset, shuffle=True, seed=42, drop_last=True)
+    else:
+        pretrain_sampler = None
     pretrain_dataloader = DataLoader(pretrain_dataset,
                                      shuffle=(pretrain_sampler is None),
                                      sampler=pretrain_sampler,

{"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@@ -8,7 +8,7 @@\n from coati.models.gpt import GPTRM, GPTActor, GPTCritic\n from coati.models.llama import LlamaActor, LlamaCritic, LlamaRM\n from coati.models.opt import OPTRM, OPTActor, OPTCritic\n-from coati.models.roberta import RoBERTaRM, RoBERTaActor, RoBERTaCritic\n+from coati.models.roberta import RoBERTaActor, RoBERTaCritic, RoBERTaRM\n from coati.trainer import PPOTrainer\n from coati.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy\n from coati.utils import prepare_llama_tokenizer_and_embedding\n@@ -143,6 +143,8 @@\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@@ -151,6 +153,8 @@\n     pretrain_dataset = SupervisedDataset(tokenizer=tokenizer, data_path=args.pretrain_dataset, max_datasets_size=16384)\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", "issue": "[BUG]: [chat] Unable to run train_prompts.sh by single card\n### \ud83d\udc1b Describe the bug\n\nMissing definition for prompt_sampler and pretrain_sampler in `examples/train_prompts.py` when dist.get_world_size() == 1.\n\n### Environment\n\n_No response_\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 RoBERTaRM, RoBERTaActor, RoBERTaCritic\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    if args.strategy != 'colossalai_gemini':\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    prompt_dataloader = DataLoader(prompt_dataset,\n                                   shuffle=(prompt_sampler is None),\n                                   sampler=prompt_sampler,\n                                   batch_size=args.train_batch_size)\n\n    pretrain_dataset = SupervisedDataset(tokenizer=tokenizer, data_path=args.pretrain_dataset, max_datasets_size=16384)\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    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    def tokenize_fn(texts):\n        # MUST padding to max length to ensure inputs of all ranks have the same length\n        # Different length may lead to hang when using gemini, as different generation steps\n        batch = tokenizer(texts, return_tensors='pt', max_length=96, padding='max_length', truncation=True)\n        return {k: v.to(torch.cuda.current_device()) for k, v in batch.items()}\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        experience_batch_size=args.experience_batch_size,\n        tokenizer=tokenize_fn,\n        max_length=128,\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    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    if args.strategy != 'colossalai_gemini':\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.train_batch_size)\n\n    pretrain_dataset = SupervisedDataset(tokenizer=tokenizer, data_path=args.pretrain_dataset, max_datasets_size=16384)\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    def tokenize_fn(texts):\n        # MUST padding to max length to ensure inputs of all ranks have the same length\n        # Different length may lead to hang when using gemini, as different generation steps\n        batch = tokenizer(texts, return_tensors='pt', max_length=96, padding='max_length', truncation=True)\n        return {k: v.to(torch.cuda.current_device()) for k, v in batch.items()}\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        experience_batch_size=args.experience_batch_size,\n        tokenizer=tokenize_fn,\n        max_length=128,\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    args = parser.parse_args()\n    main(args)\n", "path": "applications/Chat/examples/train_prompts.py"}]}