How to save and load the Peft/LoRA Finetune

#21
by LazerJesus - opened

I am trying to further finetune Starchat-Beta, save my progress, load my progress, and continue training. But whatever I do, it doesn't come together. Whenever I load my progress and continue training, my loss starts back from zero (3.xxx in my case).
I'll run you through my code and then the problem.

tokenizer = AutoTokenizer.from_pretrained(BASEPATH)
model = AutoModelForCausalLM.from_pretrained(
    "/notebooks/starbaseplus"
    ...
)
# I get both the Tokenizer and the Foundation model from the starbaseplus repo (which I have locally). 

peftconfig = LoraConfig(
    "/notebooks/starchat-beta" 
    base_model_name_or_path = "/notebooks/starbaseplus",
    ...
)
model = get_peft_model(model, peftconfig)
# All Gucci so far, the model and the LoRA fine-tune are loaded from the starchat-beta repo (also local).

# important for later:
print_trainable_parameters(model)
# trainable params: 306 Million || all params: 15 Billion || trainable: 1.971%


trainer = Trainer(
    model=model,
    ...
)
trainer.train()
# I train, loss drops. from 3.xx to 1.xx.

# Now, either I follow the HugginFace docks:
model.save_pretrained("./huggingface_model") 
# -> saves /notebooks/huggingface_model/adapter_model.bin 16mb.

# or an alternative I found on SO:
trainer.save_model("./torch_model") 
# -> saves /notebooks/torch_model/pytorch_model.bin 60gb.

I have two alternatives saved to disk. Lets restart and try either of these approaches

First the huggingface docs approach:
I now have three sets of weights.

  1. the foundation model - starbase plus
  2. the chat finetune - starchat-beta
  3. the 16mb saved bin - adapter_model.bin

But I only have two opportunities to load weights.

  1. AutoModelForCausalLM.from_pretrained
  2. either get_peft_model or PeftModel.from_pretrained

Neither works. training restarts at a loss of 3.x.

Second approach:
Load the 60bg instead of the old starchat-beta repo model.
get_peft_model("/notebooks/torch_model/pytorch_model.bin", peftconfig)

Also doesn't work. The print_trainable_parameters(model) drops to trainable: 0.02% and training restarts at a loss of 3.x

There are 4 different ways to save a model.
model.save_pretrained(PATH)
torch.save({'model_state_dict': model.state_dict()})
trainer.save_model(PATH) and
TrainerArgs(save_strategy='steps').

Which one can I use to store the PeftModelForCausalLM(AutoModelForCausalLM()) and how to load it again?

Sign up or log in to comment