github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[chat] fix train_prompts.py gemini strategy bug (#3666) 

* fix gemini strategy bug

* add comment

* add comment

* better solution
pull/3699/head
zhang-yi-chi 2023-05-06 16:46:38 +08:00 committed by GitHub
parent d556648885
commit 2da5d81dec
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 38 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
applications/Chat/examples/train_prompts.py
View File

@ -36,6 +36,7 @@ def main(args):
if args.rm_path is not None: if args.rm_path is not None:
state_dict = torch.load(args.rm_path, map_location='cpu') state_dict = torch.load(args.rm_path, map_location='cpu')
with strategy.model_init_context():
# configure model # configure model
if args.model == 'gpt2': if args.model == 'gpt2':
initial_model = GPTActor(pretrained=args.pretrain) initial_model = GPTActor(pretrained=args.pretrain)
@ -74,7 +75,6 @@ def main(args):
initial_model.to(torch.float16).to(torch.cuda.current_device()) initial_model.to(torch.float16).to(torch.cuda.current_device())
reward_model.to(torch.float16).to(torch.cuda.current_device()) reward_model.to(torch.float16).to(torch.cuda.current_device())
with strategy.model_init_context():
if args.model == 'gpt2': if args.model == 'gpt2':
actor = GPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank) actor = GPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
elif args.model == 'bloom': elif args.model == 'bloom':