Merge pull request #358 from leseb/opt-rm-checkpoints

feat: retain only last checkpoint directory

Author: mergify[bot]

README.md

@@ -105,6 +105,7 @@ for training jobs. There are a number of options you can specify, such as settin
 | fsdp_options | The settings for controlling FSDP when it's selected as the distributed backend. |
 | distributed_backend | Specifies which distributed training backend to use. Supported options are "fsdp" and "deepspeed". |
 | disable_flash_attn | Disables flash attention when set to true. This allows for training on older devices. |
+| keep_last_checkpoint_only | Determines whether we should only keep the last checkpoint directory - the previous checkpoint directory is always overwritten. The checkpoint directory is called `last_epoch`. |
 
 ### `DeepSpeedOptions`
 

src/instructlab/training/config.py

@@ -206,3 +206,8 @@ class TrainingArgs(BaseModel):
 
     # This field defines whether or not data processing will occur inside of `run_training()`
     process_data: Optional[bool] = True
+
+    # This field specifies whether only the last checkpoint should be retained. When set to true, it
+    # will overwrite the previous checkpoint directory, keeping only one directory called
+    # "last_epoch". This works alongside the '--checkpoint_at_epoch' flag.
+    keep_last_checkpoint_only: Optional[bool] = False

src/instructlab/training/main_ds.py

@@ -707,6 +707,7 @@ def run_training(torch_args: TorchrunArgs, train_args: TrainingArgs) -> None:
         f"--max_batch_len={train_args.max_batch_len}",
         f"--seed={train_args.random_seed}",
         f"--chat-tmpl-path={train_args.chat_tmpl_path}",
+        f"--keep_last_checkpoint_only={train_args.keep_last_checkpoint_only}",
     ]
 
     if train_args.checkpoint_at_epoch:
@@ -787,6 +788,9 @@ def run_training(torch_args: TorchrunArgs, train_args: TrainingArgs) -> None:
         f"--fsdp_sharding_strategy={train_args.fsdp_options.sharding_strategy.value}"
     )
 
+    if train_args.keep_last_checkpoint_only:
+        command.append("--keep_last_checkpoint_only")
+
     print(f"\033[92mRunning training command as subprocess: {' '.join(command)}\033[0m")
     process = None
     interrupt: KeyboardInterrupt | Exception | None = None
@@ -962,6 +966,14 @@ def run_training(torch_args: TorchrunArgs, train_args: TrainingArgs) -> None:
         ),
     )
     parser.add_argument("--disable_flash_attn", action="store_true")
+    parser.add_argument(
+        "--keep_last_checkpoint_only",
+        action="store_true",
+        help=(
+            "Keep only the last checkpoint directory - overwrite the previous ones. Useful for saving disk space."
+            "The last checkpoint will be saved as 'last_epoch'."
+        ),
+    )
     args = parser.parse_args()
     set_random_seed(args.seed)
     main(args)

src/instructlab/training/utils.py

@@ -925,8 +925,13 @@ def save_hf_format_accelerate(
     samples_seen,
     is_lora=False,
 ):
+    # Build the subdirectory name
+    subdir = (
+        "last_epoch" if args.keep_last_checkpoint_only else f"samples_{samples_seen}"
+    )
+
     log_rank_0(
-        f"\033[93mSaving model in huggingface format at samples_seen: {samples_seen}\033[0m",
+        f"\033[93mSaving model in huggingface format at: {subdir}\033[0m",
         to_print=True,
     )
     start = time.time()
@@ -936,7 +941,9 @@ def save_hf_format_accelerate(
     else:
         convert_dolomite = True
 
-    final_output_dir = Path(args.output_dir) / "hf_format" / f"samples_{samples_seen}"
+    # Build the final output directory path
+    final_output_dir = Path(args.output_dir) / "hf_format" / subdir
+
     if args.use_dolomite and convert_dolomite:
         tmpdir = TemporaryDirectory("w")  # pylint: disable=consider-using-with
         output_dir = Path(tmpdir.name)