feat: add save_model_dir flag where final checkpoint saved

[anhuong]

Description of the change

• Add optional save_model_dir flag where final checkpoint can be saved to using trainer.save_model()
  • Note that this only saves the model, not the optimizer states
  • Adds minimal save() to sft_trainer
• output_dir is reserved for checkpoint saving and training logs. This param is still required to pass in even if no checkpoints are saved using save_strategy="no"
  • Note that with this update, the training logs will be streamed to output_dir
• Update accelerate_launch.py:
  • Remove tempdir which caused issues with ephemeral storage
  • Remove copy final checkpoint since this is now moved into sft_trainer
  • For lm_head removal, removes lm_head from save_model_dirif exists otherwise removes from final checkpoint

Related issue number

#217

How to verify the PR

Tested:

1. save_strategy=“no” and save_model_dir to diff path than output_dir --> verified saves final model and does not save any checkpoints in ´output_dir` only logs

2. save_total_limit=2 and output_dir set (akasave_model_dir not set) --> only checkpoints are saved with logs

3. save_strategy=“no” and output_dir==save_model_dir --> verified that logs and model saved to path

4. save_strategy="epoch” and save_total_limit=2 andoutput_dir==save_model_dir` --> checkpoint dirs, model, and training logs are all written to path

5. accelerate_launch.py

   • save_total_limit=3 and save_model_dir==output_dir --> same as 4, checkpoints, training logs, and model outputted to path
   • save_strategy="no" and save_model_dir==output_dir --> same as 3, only model and logs outputted to path
   • save_total_limit=1 and output_dir subdir of save_model_dir --> output_dir with checkpoints and logs inside of save_model_dir
   • save_total_limit=1 and save_model_dir subdir of output_dir --> output_dir has checkpoints, logs, and dir with model

6. accelerate_launch: Finally I also verified that the lm_head removal continued to work as expected:

• save_total_limit=1, save_model_dir==output_dir, granite-3b-code-base --> verified that with lora and ft the model was saved to given path with lm_head removed but the checkpoint didn't have lm_head removed
• save_strategy="no", save_model_dir==output_dir, granite-3b-code-base --> verified lm_head removed and no addt checkpoints saved
• save_total_limit=1, output_dir, granite-3b-code-base (aka no save_model_dir given) --> verified lm_head removed from final checkpoint

Ran vLLM inference on LoRA and fine tuned llama-13b-base model that was saved in separate and same dir. Fine tuning got good "no complaint" inference result, LoRA got poor results after tuning but marginal improvement from base model.

Was the PR tested

• I have added >=1 unit test(s) for every new method I have added.
• I have ensured all unit tests pass

[anhuong]

Wanted to call out a note from the description: For lm_head removal, removes lm_head from save_model_dir if exists, otherwise removes from final checkpoint. Is this the behavior we want or do we only want to remove lm_head if save_model_dir passed?

[anhuong]

This was the best way I could think of to test the save() without having to try to mock up the SFTTrainer. I had tried doing something like...

    training_args = SFTConfig(**transformer_kwargs)
    trainer = SFTTrainer(
        model=MODEL_NAME,
        args=TRAIN_ARGS,
        train_dataset=TWITTER_COMPLAINTS_DATA,
    )
    with tempfile.TemporaryDirectory() as tempdir:
        sft_trainer.save(tempdir, trainer)

but this requires needing to preprocess the train args so it can be parsed into SFTConfig and requires data preprocessing to run otherwise hits error AttributeError: 'str' object has no attribute 'column_names' when trying to tokenize the dataset

[ashokponkumar]

Thanks Anh! It looks good. Just a few questions.

[ashokponkumar]

Should we instead show an example of save_model_dir being a subfolder of output_dir? That way the the logs, checkpoints and final model area easily segregable.

[anhuong]

Yes can do! I agree this is recommended behavior, but wanted to point out the edge case that may cause the most confusion.

[ashokponkumar]

sure.

[ashokponkumar]

Do we know if removing the lm_head is required while resuming a training, if so we should think about removing the lm_head from each intermediate checkpoint as it is being written.

[anhuong]

Removing lm_head is required for loading the tuned model with vLLM, but I don't think it's required for resuming training. Is there a way I could check resuming training? If you have a command to run with sft_trainer, that would be super helpful!

[ashokponkumar]

I think we have to stop a run midway and then run the training again.

We have to ensure we use https://huggingface.co/docs/transformers/en/main_classes/trainer#transformers.Trainer.train.resume_from_checkpoint argument. Seems like we are not currently setting it. We should possibly set it. This will resume from the previously stopped last checkpoint.

[Ssukriti]

hi @ashokponkumar we have not heard of any requirement to remove lm_head to resume training. @Abhishek-TAMU will be working on resume training from checkpoint in follow up PR as per issue https://github.ibm.com/ai-foundation/watson-fm-stack-tracker/issues/1007 . Can we wait for him to add the change?

ypu can post any tips on the issue @ashokponkumar

[ashokponkumar]

Sure Sukriti.

[Ssukriti]

Suggested change

	"help": "Directory where final checkpoint will be saved to \
	"help": "Directory where final tuned model will be saved to \

[Ssukriti]

I think there has been some confusion here.

If the user sets a validation dataset and load_best_model_at_end, then the best checkpoint will be saved. If no additional flags are set, the final checkpoint will be saved.

This is for when save_total_limit = 1 , then you can. control best or last

save_model() is always for saving tuned model at end after training. By then any checkpoint information is lost. Can you confirm

[Ssukriti]

This README is confusing and can be simplified.

We can have. 2 subsections

1. Save checkpoints while training

• save_strategy = set to epoch , can also be set to steps or no (as described)
• checkpoints are saved to output_dir
• save_total_limit para

NOTE: load_best_model para applies when save_total_limit=1 , as it is a training argument . it does not apply to trainer.save_model as that is after saving. we can skip this para on load_best_model as users can look up TrainingArguments themselves. It si always described in URL you linked https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments.save_total_limit

2. Save model after training -
   save_model_dircan optionally be set to save the tuned model usingSFTTrainer.save_model(). This can be used in tandem with save_strategy="no"` to only save the final tuned and not any intermediate checkpoints, which can help to save space.

ways you can use save_model_dir and more tips (collapse this section)

examples on expanding:

[Ssukriti]

better to skip this section as described above, does not apply to save_model. It applies to save checkpoint only

[anhuong]

i ended up keeping it but moving it within the save_total_limits section and not in the save_model_dir section

[Ssukriti]

1. Have you verified saved model, using save() still infers on vLLM?

2. Comments on re-organizing README

3. merge conflicts to be addressed

[anhuong]

Have you verified saved model, using save() still infers on vLLM?

Yes this is described in the description of the PR above

Ran vLLM inference on LoRA and fine tuned llama-13b-base model that was saved in separate and same dir. Fine tuning got good "no complaint" inference result, LoRA got poor results after tuning but marginal improvement from base model.

[Ssukriti]

Nice work @anhuong !