[Train] Adding more flexibility to the default checkpointing behavior (…

…ray-project#37897)

Signed-off-by: Kourosh Hakhamaneshi <[email protected]>

doc/source/templates/04_finetuning_llms_with_deepspeed/finetune_hf_llm.py

@@ -461,6 +461,15 @@ def parse_args():
  parser.add_argument(
  "--num-epochs", type=int, default=1, help="Number of epochs to train for."
  )
+ parser.add_argument(
+ "--num-checkpoints-to-keep",
+ type=int,
+ help=(
+ "Number of checkpoints to keep, if None, all checkpoints will be kept, "
+ "if set to n>=1, the top n checkpoint with min. evaluation perplexity "
+ "will be kept."
+ ),
+ )
  parser.add_argument("--lr", type=float, default=5e-6, help="Learning rate to use.")
 
  parser.add_argument(
@@ -550,7 +559,9 @@ def main():
  # sync_config=tune.SyncConfig(sync_artifacts=False),
  storage_path=storage_path,
  checkpoint_config=air.CheckpointConfig(
- num_to_keep=1,
+ num_to_keep=args.num_checkpoints_to_keep,
+ checkpoint_score_attribute="perplexity",
+ checkpoint_score_order="min",
  # Enable distributed checkpointing
  _checkpoint_keep_all_ranks=True,
  _checkpoint_upload_from_workers=True,

doc/source/templates/04_finetuning_llms_with_deepspeed/run_llama_ft.sh

@@ -53,6 +53,8 @@ fine_tune() {
  --train_path "${train_path}" \
  --test_path "${test_path}" \
  --special_token_path "${token_path}" \
+ --num-checkpoints-to-keep 1 \
+ --num-epochs 1 \
  "${params[@]}"; then
  echo "Failed to fine-tune the model. Exiting..."
  exit 1