[BUG]: Chat GPT inference. Error in load_state_dict for BloomForCasualLM

[sdvies]

🐛 Describe the bug

Hello!

Bug description

I've reproduced the ChatGPT Stages described in https://github.com/hpcaitech/ColossalAI/tree/main/applications/ChatGPT/examples#inference-exampleafter-stage3.

I've obtained the rm_checkpoint.pt from Stage 2 and I could obtain the actor_checkpoint_prompts.pt as well from Stage 3 (based on rm_checkpoint.pt).

The problem occurs when I try to run the inference process because the following error raises:

RuntimeError: Error(s) in loading state_dict for BloomForCausalLM:
	Unexpected key(s) in state_dict: "transformer.h.2.input_layernorm.weight", "transformer.h.2.input_layernorm.bias", "transformer.h.2.self_attention.query_key_value.weight", "transformer.h.2.self_attention.query_key_value.bias", "transformer.h.2.self_attention.query_key_value.lora_A", "transformer.h.2.self_attention.query_key_value.lora_B" ...

To reproduce

In order to obtain the rm_checkpoint.pt I executed the following:
python train_reward_model.py --pretrain bigscience/bloom-560m --model bloom --batch_size 1 --lora_rank 16 --strategy naive

The actor_checkpoint_prompts.pt was obtain by executing:

python train_prompts.py prompts.csv --pretrain bigscience/bloom-560m --strategy naive --model bloom --train_batch_size 1 --lora_rank 16 --experience_batch_size 1 --max_epochs

And finally I tried to do the inference by:
python inference.py --model_path actor_checkpoint_prompts.pt --pretrain bigscience/bloom-560m --model bloom

Thank you very much!

Environment

I include the environment information as follows:

• CUDA: 11.1
• cuDNN: 8
• Python: 3.8.8
• PyTorch: 1.13.1

[JinQiangWang2021]

@sdvies I meet the same problem with you. Have you solved it?

Traceback (most recent call last):
File "inference.py", line 59, in
eval(args)
File "inference.py", line 23, in eval
actor.model.load_state_dict(state_dict)
File "/home/user/miniconda3/envs/chatgpt/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1672, in load_state_dict
self.class.name, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for BloomForCausalLM:
Missing key(s) in state_dict: "transformer.word_embeddings.weight", "transformer.word_embeddings_layernorm.weight", "transformer.word_embeddings_layernorm.bias", "transformer.h.0.input_layernorm.weight", "transformer.h.0.input_layernorm.bias", "transformer.h.0.self_attention.query_key_value.weight", "transformer.h.0.self_attention.query_key_value.bias", "transformer.h.0.self_attention.dense.weight", "transformer.h.0.self_attention.dense.bias", "transformer.h.0.post_attention_layernorm.weight", "transformer.h.0.post_attention_layernorm.bias", "transformer.h.0.mlp.dense_h_to_4h.weight", "transformer.h.0.mlp.dense_h_to_4h.bias", "transformer.h.0.mlp.dense_4h_to_h.weight", "transformer.h.0.mlp.dense_4h_to_h.bias", "transformer.h.1.input_layernorm.weight", "transformer.h.1.input_layernorm.bias", "transformer.h.1.

• nvidia-cudnn-cu11 8.5.0.96
• python 3.7.16
• torch 1.13.1

[sdvies]

Hello @ht-zhou , thank you very much for your previous response!

I've obtained the actor_checkpoint_prompts.pt again with the new version of the train_prompts.py code and I'm still getting the same error:

RuntimeError: Error(s) in loading state_dict for BloomForCausalLM:
	Unexpected key(s) in state_dict: "transformer.h.2.input_layernorm.weight", "transformer.h.2.input_layernorm.bias", "transformer.h.2.self_attention.query_key_value.weight", "transformer.h.2.self_attention.query_key_value.bias", "transformer.h.2.self_attention.query_key_value.lora_A", "transformer.h.2.self_attention.query_key_value.lora_B" ...

Thank you!