CarperAI / trlx

A repo for distributed training of language models with Reinforcement Learning via Human Feedback (RLHF)
MIT License
4.51k stars 471 forks source link

Load the checkpoint fails #600

Open AfraAmini opened 2 months ago

AfraAmini commented 2 months ago

🐛 Describe the bug

I am training a model using PPO with 8 unfrozen layers. I have the checkpoint saved under hf_model directory. When I attempt to load the checkpoint I get the error that layers are not initialized from the checkpoint and there are extra keys (base_model. and frozen_head.) that were not matched with any keys in the GPT model. Here is how I attempted to load the model:

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("ckpts/checkpoint_6400/hf_model")

Which trlX version are you using?

No response

Additional system and package information

No response