Repository: hpcaitech/ColossalAI

Issue:
[BUG]: inference ERROR again
### 🐛 Describe the bug

Traceback (most recent call last):
  File "inference.py", line 59, in <module>
    eval(args)
  File "inference.py", line 23, in eval
    actor.model.load_state_dict(state_dict)
  File "/home/rst/ColossalAI/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1671, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for OPTForCausalLM:
	Missing key(s) in state_dict: "model.decoder.project_out.weight", "model.decoder.project_in.weight", "model.decoder.layers.12.self_attn.k_proj.weight", "model.decoder.layers.12.self_attn.k_proj.bias", "model.decoder.layers.12.self_attn.v_proj.weight", "model.decoder.layers.12.self_attn.v_proj.bias", "model.decoder.layers.12.self_attn.q_proj.weight", "model.decoder.layers.12.self_attn.q_proj.bias", "model.decoder.layers.12.self_attn.out_proj.weight", "model.decoder.layers.12.self_attn.out_proj.bias", "model.decoder.layers.12.self_attn_layer_norm.weight", "model.decoder.layers.12.self_attn_layer_norm.bias", "model.decoder.layers.12.fc1.weight", "model.decoder.layers.12.fc1.bias", "model.decoder.layers.12.fc2.weight", "model.decoder.layers.12.fc2.bias", "model.decoder.layers.12.final_layer_norm.weight", "model.decoder.layers.12.final_layer_norm.bias", "model.decoder.layers.13.self_attn.k_proj.weight", "model.decoder.layers.13.self_attn.k_proj.bias", "model.decoder.layers.13.self_attn.v_proj.weight", "model.decoder.layers.13.self_attn.v_proj.bias", "model.decoder.layers.13.self_attn.q_proj.weight", "model.decoder.layers.13.self_attn.q_proj.bias", "model.decoder.layers.13.self_attn.out_proj.weight", "model.decoder.layers.13.self_attn.out_proj.bias", "model.decoder.layers.13.self_attn_layer_norm.weight", "model.decoder.layers.13.self_attn_layer_norm.bias", "model.decoder.layers.13.fc1.weight", "model.decoder.layers.13.fc1.bias", "model.decoder.layers.13.fc2.weight", "model.decoder.layers.13.fc2.bias", "model.decoder.layers.13.final_layer_norm.weight", "model.decoder.layers.13.final_layer_norm.bias", "model.decoder.layers.14.self_attn.k_proj.weight", "model.decoder.layers.14.self_attn.k_proj.bias", "model.decoder.layers.14.self_attn.v_proj.weight", "model.decoder.layers.14.self_attn.v_proj.bias", "model.decoder.layers.14.self_attn.q_proj.weight", "model.decoder.layers.14.self_attn.q_proj.bias", "model.decoder.layers.14.self_attn.out_proj.weight", "model.decoder.layers.14.self_attn.out_proj.bias", "model.decoder.layers.14.self_attn_layer_norm.weight", "model.decoder.layers.14.self_attn_layer_norm.bias", "model.decoder.layers.14.fc1.weight", "model.decoder.layers.14.fc1.bias", "model.decoder.layers.14.fc2.weight", "model.decoder.layers.14.fc2.bias", "model.decoder.layers.14.final_layer_norm.weight", "model.decoder.layers.14.final_layer_norm.bias", "model.decoder.layers.15.self_attn.k_proj.weight", "model.decoder.layers.15.self_attn.k_proj.bias", "model.decoder.layers.15.self_attn.v_proj.weight", "model.decoder.layers.15.self_attn.v_proj.bias", "model.decoder.layers.15.self_attn.q_proj.weight", "model.decoder.layers.15.self_attn.q_proj.bias", "model.decoder.layers.15.self_attn.out_proj.weight", "model.decoder.layers.15.self_attn.out_proj.bias", "model.decoder.layers.15.self_attn_layer_norm.weight", "model.decoder.layers.15.self_attn_layer_norm.bias", "model.decoder.layers.15.fc1.weight", "model.decoder.layers.15.fc1.bias", "model.decoder.layers.15.fc2.weight", "model.decoder.layers.15.fc2.bias", "model.decoder.layers.15.final_layer_norm.weight", "model.decoder.layers.15.final_layer_norm.bias", "model.decoder.layers.16.self_attn.k_proj.weight", "model.decoder.layers.16.self_attn.k_proj.bias", "model.decoder.layers.16.self_attn.v_proj.weight", "model.decoder.layers.16.self_attn.v_proj.bias", "model.decoder.layers.16.self_attn.q_proj.weight", "model.decoder.layers.16.self_attn.q_proj.bias", "model.decoder.layers.16.self_attn.out_proj.weight", "model.decoder.layers.16.self_attn.out_proj.bias", "model.decoder.layers.16.self_attn_layer_norm.weight", "model.decoder.layers.16.self_attn_layer_norm.bias", "model.decoder.layers.16.fc1.weight", "model.decoder.layers.16.fc1.bias", "model.decoder.layers.16.fc2.weight", "model.decoder.layers.16.fc2.bias", "model.decoder.layers.16.final_layer_norm.weight", "model.decoder.layers.16.final_layer_norm.bias", "model.decoder.layers.17.self_attn.k_proj.weight", "model.decoder.layers.17.self_attn.k_proj.bias", "model.decoder.layers.17.self_attn.v_proj.weight", "model.decoder.layers.17.self_attn.v_proj.bias", "model.decoder.layers.17.self_attn.q_proj.weight", "model.decoder.layers.17.self_attn.q_proj.bias", "model.decoder.layers.17.self_attn.out_proj.weight", "model.decoder.layers.17.self_attn.out_proj.bias", "model.decoder.layers.17.self_attn_layer_norm.weight", "model.decoder.layers.17.self_attn_layer_norm.bias", "model.decoder.layers.17.fc1.weight", "model.decoder.layers.17.fc1.bias", "model.decoder.layers.17.fc2.weight", "model.decoder.layers.17.fc2.bias", "model.decoder.layers.17.final_layer_norm.weight", "model.decoder.layers.17.final_layer_norm.bias", "model.decoder.layers.18.self_attn.k_proj.weight", "model.decoder.layers.18.self_attn.k_proj.bias", "model.decoder.layers.18.self_attn.v_proj.weight", "model.decoder.layers.18.self_attn.v_proj.bias", "model.decoder.layers.18.self_attn.q_proj.weight", "model.decoder.layers.18.self_attn.q_proj.bias", "model.decoder.layers.18.self_attn.out_proj.weight", "model.decoder.layers.18.self_attn.out_proj.bias", "model.decoder.layers.18.self_attn_layer_norm.weight", "model.decoder.layers.18.self_attn_layer_norm.bias", "model.decoder.layers.18.fc1.weight", "model.decoder.layers.18.fc1.bias", "model.decoder.layers.18.fc2.weight", "model.decoder.layers.18.fc2.bias", "model.decoder.layers.18.final_layer_norm.weight", "model.decoder.layers.18.final_layer_norm.bias", "model.decoder.layers.19.self_attn.k_proj.weight", "model.decoder.layers.19.self_attn.k_proj.bias", "model.decoder.layers.19.self_attn.v_proj.weight", "model.decoder.layers.19.self_attn.v_proj.bias", "model.decoder.layers.19.self_attn.q_proj.weight", "model.decoder.layers.19.self_attn.q_proj.bias", "model.decoder.layers.19.self_attn.out_proj.weight", "model.decoder.layers.19.self_attn.out_proj.bias", "model.decoder.layers.19.self_attn_layer_norm.weight", "model.decoder.layers.19.self_attn_layer_norm.bias", "model.decoder.layers.19.fc1.weight", "model.decoder.layers.19.fc1.bias", "model.decoder.layers.19.fc2.weight", "model.decoder.layers.19.fc2.bias", "model.decoder.layers.19.final_layer_norm.weight", "model.decoder.layers.19.final_layer_norm.bias", "model.decoder.layers.20.self_attn.k_proj.weight", "model.decoder.layers.20.self_attn.k_proj.bias", "model.decoder.layers.20.self_attn.v_proj.weight", "model.decoder.layers.20.self_attn.v_proj.bias", "model.decoder.layers.20.self_attn.q_proj.weight", "model.decoder.layers.20.self_attn.q_proj.bias", "model.decoder.layers.20.self_attn.out_proj.weight", "model.decoder.layers.20.self_attn.out_proj.bias", "model.decoder.layers.20.self_attn_layer_norm.weight", "model.decoder.layers.20.self_attn_layer_norm.bias", "model.decoder.layers.20.fc1.weight", "model.decoder.layers.20.fc1.bias", "model.decoder.layers.20.fc2.weight", "model.decoder.layers.20.fc2.bias", "model.decoder.layers.20.final_layer_norm.weight", "model.decoder.layers.20.final_layer_norm.bias", "model.decoder.layers.21.self_attn.k_proj.weight", "model.decoder.layers.21.self_attn.k_proj.bias", "model.decoder.layers.21.self_attn.v_proj.weight", "model.decoder.layers.21.self_attn.v_proj.bias", "model.decoder.layers.21.self_attn.q_proj.weight", "model.decoder.layers.21.self_attn.q_proj.bias", "model.decoder.layers.21.self_attn.out_proj.weight", "model.decoder.layers.21.self_attn.out_proj.bias", "model.decoder.layers.21.self_attn_layer_norm.weight", "model.decoder.layers.21.self_attn_layer_norm.bias", "model.decoder.layers.21.fc1.weight", "model.decoder.layers.21.fc1.bias", "model.decoder.layers.21.fc2.weight", "model.decoder.layers.21.fc2.bias", "model.decoder.layers.21.final_layer_norm.weight", "model.decoder.layers.21.final_layer_norm.bias", "model.decoder.layers.22.self_attn.k_proj.weight", "model.decoder.layers.22.self_attn.k_proj.bias", "model.decoder.layers.22.self_attn.v_proj.weight", "model.decoder.layers.22.self_attn.v_proj.bias", "model.decoder.layers.22.self_attn.q_proj.weight", "model.decoder.layers.22.self_attn.q_proj.bias", "model.decoder.layers.22.self_attn.out_proj.weight", "model.decoder.layers.22.self_attn.out_proj.bias", "model.decoder.layers.22.self_attn_layer_norm.weight", "model.decoder.layers.22.self_attn_layer_norm.bias", "model.decoder.layers.22.fc1.weight", "model.decoder.layers.22.fc1.bias", "model.decoder.layers.22.fc2.weight", "model.decoder.layers.22.fc2.bias", "model.decoder.layers.22.final_layer_norm.weight", "model.decoder.layers.22.final_layer_norm.bias", "model.decoder.layers.23.self_attn.k_proj.weight", "model.decoder.layers.23.self_attn.k_proj.bias", "model.decoder.layers.23.self_attn.v_proj.weight", "model.decoder.layers.23.self_attn.v_proj.bias", "model.decoder.layers.23.self_attn.q_proj.weight", "model.decoder.layers.23.self_attn.q_proj.bias", "model.decoder.layers.23.self_attn.out_proj.weight", "model.decoder.layers.23.self_attn.out_proj.bias", "model.decoder.layers.23.self_attn_layer_norm.weight", "model.decoder.layers.23.self_attn_layer_norm.bias", "model.decoder.layers.23.fc1.weight", "model.decoder.layers.23.fc1.bias", "model.decoder.layers.23.fc2.weight", "model.decoder.layers.23.fc2.bias", "model.decoder.layers.23.final_layer_norm.weight", "model.decoder.layers.23.final_layer_norm.bias". 
	Unexpected key(s) in state_dict: "model.decoder.final_layer_norm.weight", "model.decoder.final_layer_norm.bias". 


### Environment

_No response_
[tensor] fix some unittests

[tensor] fix some unittests



Discussion:
python inference.py --model_path ./actor_checkpoint_prompts.pt --pretrain facebook/opt-350m --model opt
Thank you for your feedback, we will fix it soon.
Still met the same error



Resolving patch:
diff --git a/applications/ChatGPT/examples/train_dummy.py b/applications/ChatGPT/examples/train_dummy.py
--- a/applications/ChatGPT/examples/train_dummy.py
+++ b/applications/ChatGPT/examples/train_dummy.py
@@ -38,19 +38,19 @@ def main(args):
     # configure model
     with strategy.model_init_context():
         if args.model == 'gpt2':
-            actor = GPTActor().cuda()
-            critic = GPTCritic().cuda()
+            actor = GPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
+            critic = GPTCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
         elif args.model == 'bloom':
-            actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()
-            critic = BLOOMCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()
+            actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
+            critic = BLOOMCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
         elif args.model == 'opt':
-            actor = OPTActor().cuda()
-            critic = OPTCritic().cuda()
+            actor = OPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
+            critic = OPTCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
         else:
             raise ValueError(f'Unsupported model "{args.model}"')
 
-        initial_model = deepcopy(actor).cuda()
-        reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).cuda()
+        initial_model = deepcopy(actor).to(torch.cuda.current_device())
+        reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).to(torch.cuda.current_device())
 
     # configure optimizer
     if args.strategy.startswith('colossalai'):
@@ -114,12 +114,13 @@ def main(args):
                 max_timesteps=args.max_timesteps,
                 update_timesteps=args.update_timesteps)
 
-    # save model checkpoint after fitting on only rank0
+    # save model checkpoint after fitting
     strategy.save_model(actor, 'actor_checkpoint_dummy.pt', only_rank0=True)
     # save optimizer checkpoint on all ranks
-    strategy.save_optimizer(actor_optim,
-                            'actor_optim_checkpoint_dummy_%d.pt' % (torch.cuda.current_device()),
-                            only_rank0=False)
+    if args.need_optim_ckpt:
+        strategy.save_optimizer(actor_optim,
+                                'actor_optim_checkpoint_dummy_%d.pt' % (torch.cuda.current_device()),
+                                only_rank0=False)
 
 
 if __name__ == '__main__':
@@ -129,6 +130,7 @@ def main(args):
                         default='naive')
     parser.add_argument('--model', type=str, default='gpt2', choices=['gpt2', 'bloom', 'opt'])
     parser.add_argument('--pretrain', type=str, default=None)
+    parser.add_argument('--need_optim_ckpt', type=bool, default=False)
     parser.add_argument('--num_episodes', type=int, default=50)
     parser.add_argument('--max_timesteps', type=int, default=10)
     parser.add_argument('--update_timesteps', type=int, default=10)
diff --git a/applications/ChatGPT/examples/train_prompts.py b/applications/ChatGPT/examples/train_prompts.py
--- a/applications/ChatGPT/examples/train_prompts.py
+++ b/applications/ChatGPT/examples/train_prompts.py
@@ -32,19 +32,20 @@ def main(args):
     # configure model
     with strategy.model_init_context():
         if args.model == 'gpt2':
-            actor = GPTActor().cuda()
-            critic = GPTCritic().cuda()
+            actor = GPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
+            critic = GPTCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
         elif args.model == 'bloom':
-            actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()
-            critic = BLOOMCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda()
+            actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
+            critic = BLOOMCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
         elif args.model == 'opt':
-            actor = OPTActor(lora_rank=args.lora_rank).cuda()
-            critic = OPTCritic(lora_rank=args.lora_rank).cuda()
+            actor = OPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
+            critic = OPTCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
         else:
             raise ValueError(f'Unsupported model "{args.model}"')
 
         initial_model = deepcopy(actor)
-        reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).cuda()
+        reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).to(torch.cuda.current_device())
+
 
     # configure optimizer
     if args.strategy.startswith('colossalai'):
@@ -100,12 +101,13 @@ def tokenize_fn(texts):
                 num_episodes=args.num_episodes,
                 max_timesteps=args.max_timesteps,
                 update_timesteps=args.update_timesteps)
-    # save model checkpoint after fitting on only rank0
+    # save model checkpoint after fitting 
     strategy.save_model(actor, 'actor_checkpoint_prompts.pt', only_rank0=True)
     # save optimizer checkpoint on all ranks
-    strategy.save_optimizer(actor_optim,
-                            'actor_optim_checkpoint_prompts_%d.pt' % (torch.cuda.current_device()),
-                            only_rank0=False)
+    if args.need_optim_ckpt:
+        strategy.save_optimizer(actor_optim,
+                                'actor_optim_checkpoint_prompts_%d.pt' % (torch.cuda.current_device()),
+                                only_rank0=False)
 
 
 if __name__ == '__main__':
@@ -116,6 +118,7 @@ def tokenize_fn(texts):
                         default='naive')
     parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt'])
     parser.add_argument('--pretrain', type=str, default=None)
+    parser.add_argument('--need_optim_ckpt', type=bool, default=False)
     parser.add_argument('--num_episodes', type=int, default=10)
     parser.add_argument('--max_timesteps', type=int, default=10)
     parser.add_argument('--update_timesteps', type=int, default=10)


Test patch: