We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

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

{"golden_diff": "diff --git a/applications/Chat/examples/community/peft/train_peft_prompts.py b/applications/Chat/examples/community/peft/train_peft_prompts.py\n--- a/applications/Chat/examples/community/peft/train_peft_prompts.py\n+++ b/applications/Chat/examples/community/peft/train_peft_prompts.py\n@@ -24,7 +24,9 @@\n     if args.strategy == \"ddp\":\n         strategy = DDPStrategy()\n     elif args.strategy == \"colossalai_gemini\":\n-        strategy = GeminiStrategy(placement_policy=\"static\", offload_optim_frac=1.0, offload_param_frac=1.0, initial_scale=2**5)\n+        strategy = GeminiStrategy(\n+            placement_policy=\"static\", offload_optim_frac=1.0, offload_param_frac=1.0, initial_scale=2**5\n+        )\n     elif args.strategy == \"colossalai_zero2\":\n         strategy = LowLevelZeroStrategy(stage=2, placement_policy=\"cpu\")\n     else:\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "import argparse\n\nimport torch\nimport torch.distributed as dist\nfrom coati.dataset import DataCollatorForSupervisedDataset\nfrom coati.models.bloom import BLOOMRM, BLOOMCritic\nfrom coati.models.gpt import GPTRM, GPTCritic\nfrom coati.models.llama import LlamaCritic, LlamaRM\nfrom coati.models.opt import OPTRM, OPTCritic\nfrom coati.trainer import PPOTrainer\nfrom coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy\nfrom easy_dataset import EasyPromptsDataset, EasySupervisedDataset\nfrom easy_models import BLOOMActor\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\n\nfrom colossalai.nn.optimizer import HybridAdam\n\n\ndef main(args):\n    # configure strategy\n    if args.strategy == \"ddp\":\n        strategy = DDPStrategy()\n    elif args.strategy == \"colossalai_gemini\":\n        strategy = GeminiStrategy(placement_policy=\"static\", offload_optim_frac=1.0, offload_param_frac=1.0, initial_scale=2**5)\n    elif args.strategy == \"colossalai_zero2\":\n        strategy = LowLevelZeroStrategy(stage=2, placement_policy=\"cpu\")\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 == \"bloom\":\n        # initial_model = BLOOMActor(pretrained=args.pretrain)\n        print(\"Using peft lora to load Bloom model as initial_model\")\n        initial_model = BLOOMActor(pretrained=args.pretrain, lora_path=args.sft_lora_path)\n        print(\"Using peft lora to load Bloom model as initial_model (Done)\")\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        print(\"load bloom reward model \", args.rm_pretrain)\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    else:\n        raise ValueError(f'Unsupported reward model \"{rm_model_name}\"')\n\n    if args.rm_path is not None:\n        print(\"Loading reward model from\", args.rm_path)\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 == \"bloom\":\n            # actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank)\n            print(\"Using peft lora to load Bloom model as Actor\")\n            actor = BLOOMActor(pretrained=args.pretrain, lora_path=args.sft_lora_path)\n            print(\"Using peft lora to load Bloom model as Actor (Done)\")\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            print(\"load bloom critic \", args.rm_pretrain, \" lora_rank \", args.lora_rank, \" use_action_mask \", True)\n            critic = BLOOMCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)\n            print(\"load bloom critic (Done) \")\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        else:\n            raise ValueError(f'Unsupported reward model \"{rm_model_name}\"')\n\n        if args.rm_path is not None:\n            print(\"Loading reward model from\", args.rm_path)\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(args.rm_pretrain)\n        tokenizer.pad_token = tokenizer.eos_token\n    elif args.model == \"bloom\":\n        tokenizer = BloomTokenizerFast.from_pretrained(args.rm_pretrain)\n        tokenizer.pad_token = tokenizer.eos_token\n    elif args.model == \"opt\":\n        tokenizer = AutoTokenizer.from_pretrained(args.rm_pretrain)\n        tokenizer.pad_token = tokenizer.eos_token\n    elif args.model == \"llama\":\n        tokenizer = LlamaTokenizer.from_pretrained(args.pretrain)\n        tokenizer.eos_token = \"</s>\"\n        tokenizer.pad_token = tokenizer.unk_token\n    else:\n        raise ValueError(f'Unsupported model \"{args.model}\"')\n\n    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)\n\n    prompt_dataset = EasyPromptsDataset(args.prompt_path, tokenizer)\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(\n        prompt_dataset, shuffle=(prompt_sampler is None), sampler=prompt_sampler, batch_size=args.train_batch_size\n    )\n\n    pretrain_dataset = EasySupervisedDataset(args.pretrain_dataset, tokenizer)\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(\n        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\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        train_batch_size=args.train_batch_size,\n        experience_batch_size=args.experience_batch_size,\n        tokenizer=tokenize_fn,\n        max_length=512,\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(\n        prompt_dataloader=prompt_dataloader,\n        pretrain_dataloader=pretrain_dataloader,\n        num_episodes=args.num_episodes,\n        num_update_steps=args.num_update_steps,\n        num_collect_steps=args.num_collect_steps,\n    )\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(\n            actor_optim, \"actor_optim_checkpoint_prompts_%d.pt\" % (torch.cuda.current_device()), only_rank0=False\n        )\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(\n        \"--strategy\", choices=[\"ddp\", \"colossalai_gemini\", \"colossalai_zero2\"], default=\"ddp\", help=\"strategy to use\"\n    )\n    parser.add_argument(\"--model\", default=\"gpt2\", choices=[\"gpt2\", \"bloom\", \"opt\", \"llama\"])\n    parser.add_argument(\"--pretrain\", type=str, default=None)\n    parser.add_argument(\"--sft_lora_path\", type=str, default=None)\n    parser.add_argument(\"--rm_model\", default=None, choices=[\"gpt2\", \"bloom\", \"opt\", \"llama\"])\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(\"--num_collect_steps\", type=int, default=10)\n    parser.add_argument(\"--num_update_steps\", type=int, default=5)\n    parser.add_argument(\"--train_batch_size\", type=int, default=2)\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/community/peft/train_peft_prompts.py"}], "after_files": [{"content": "import argparse\n\nimport torch\nimport torch.distributed as dist\nfrom coati.dataset import DataCollatorForSupervisedDataset\nfrom coati.models.bloom import BLOOMRM, BLOOMCritic\nfrom coati.models.gpt import GPTRM, GPTCritic\nfrom coati.models.llama import LlamaCritic, LlamaRM\nfrom coati.models.opt import OPTRM, OPTCritic\nfrom coati.trainer import PPOTrainer\nfrom coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy\nfrom easy_dataset import EasyPromptsDataset, EasySupervisedDataset\nfrom easy_models import BLOOMActor\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\n\nfrom colossalai.nn.optimizer import HybridAdam\n\n\ndef main(args):\n    # configure strategy\n    if args.strategy == \"ddp\":\n        strategy = DDPStrategy()\n    elif args.strategy == \"colossalai_gemini\":\n        strategy = GeminiStrategy(\n            placement_policy=\"static\", offload_optim_frac=1.0, offload_param_frac=1.0, initial_scale=2**5\n        )\n    elif args.strategy == \"colossalai_zero2\":\n        strategy = LowLevelZeroStrategy(stage=2, placement_policy=\"cpu\")\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 == \"bloom\":\n        # initial_model = BLOOMActor(pretrained=args.pretrain)\n        print(\"Using peft lora to load Bloom model as initial_model\")\n        initial_model = BLOOMActor(pretrained=args.pretrain, lora_path=args.sft_lora_path)\n        print(\"Using peft lora to load Bloom model as initial_model (Done)\")\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        print(\"load bloom reward model \", args.rm_pretrain)\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    else:\n        raise ValueError(f'Unsupported reward model \"{rm_model_name}\"')\n\n    if args.rm_path is not None:\n        print(\"Loading reward model from\", args.rm_path)\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 == \"bloom\":\n            # actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank)\n            print(\"Using peft lora to load Bloom model as Actor\")\n            actor = BLOOMActor(pretrained=args.pretrain, lora_path=args.sft_lora_path)\n            print(\"Using peft lora to load Bloom model as Actor (Done)\")\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            print(\"load bloom critic \", args.rm_pretrain, \" lora_rank \", args.lora_rank, \" use_action_mask \", True)\n            critic = BLOOMCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)\n            print(\"load bloom critic (Done) \")\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        else:\n            raise ValueError(f'Unsupported reward model \"{rm_model_name}\"')\n\n        if args.rm_path is not None:\n            print(\"Loading reward model from\", args.rm_path)\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(args.rm_pretrain)\n        tokenizer.pad_token = tokenizer.eos_token\n    elif args.model == \"bloom\":\n        tokenizer = BloomTokenizerFast.from_pretrained(args.rm_pretrain)\n        tokenizer.pad_token = tokenizer.eos_token\n    elif args.model == \"opt\":\n        tokenizer = AutoTokenizer.from_pretrained(args.rm_pretrain)\n        tokenizer.pad_token = tokenizer.eos_token\n    elif args.model == \"llama\":\n        tokenizer = LlamaTokenizer.from_pretrained(args.pretrain)\n        tokenizer.eos_token = \"</s>\"\n        tokenizer.pad_token = tokenizer.unk_token\n    else:\n        raise ValueError(f'Unsupported model \"{args.model}\"')\n\n    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)\n\n    prompt_dataset = EasyPromptsDataset(args.prompt_path, tokenizer)\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(\n        prompt_dataset, shuffle=(prompt_sampler is None), sampler=prompt_sampler, batch_size=args.train_batch_size\n    )\n\n    pretrain_dataset = EasySupervisedDataset(args.pretrain_dataset, tokenizer)\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(\n        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\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        train_batch_size=args.train_batch_size,\n        experience_batch_size=args.experience_batch_size,\n        tokenizer=tokenize_fn,\n        max_length=512,\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(\n        prompt_dataloader=prompt_dataloader,\n        pretrain_dataloader=pretrain_dataloader,\n        num_episodes=args.num_episodes,\n        num_update_steps=args.num_update_steps,\n        num_collect_steps=args.num_collect_steps,\n    )\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(\n            actor_optim, \"actor_optim_checkpoint_prompts_%d.pt\" % (torch.cuda.current_device()), only_rank0=False\n        )\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(\n        \"--strategy\", choices=[\"ddp\", \"colossalai_gemini\", \"colossalai_zero2\"], default=\"ddp\", help=\"strategy to use\"\n    )\n    parser.add_argument(\"--model\", default=\"gpt2\", choices=[\"gpt2\", \"bloom\", \"opt\", \"llama\"])\n    parser.add_argument(\"--pretrain\", type=str, default=None)\n    parser.add_argument(\"--sft_lora_path\", type=str, default=None)\n    parser.add_argument(\"--rm_model\", default=None, choices=[\"gpt2\", \"bloom\", \"opt\", \"llama\"])\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(\"--num_collect_steps\", type=int, default=10)\n    parser.add_argument(\"--num_update_steps\", type=int, default=5)\n    parser.add_argument(\"--train_batch_size\", type=int, default=2)\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/community/peft/train_peft_prompts.py"}]}