[Feature] Option not to save optimizer states to save disk space (#1333)

Author: yzlnew

slime/backends/fsdp_utils/checkpoint.py

@@ -214,14 +214,15 @@ def save(actor: Any, iteration: int) -> None:
     state_dict = {"model_state": model_state}
     dcp.save(state_dict, checkpoint_id=str(model_dir))
 
-    # Save optimizer state
-    if hasattr(actor, "optimizer") and actor.optimizer is not None:
+    # Save optimizer state (skip if --no-save-optim is set)
+    save_optimizer_state = not getattr(actor.args, "no_save_optim", False)
+    if save_optimizer_state and hasattr(actor, "optimizer") and actor.optimizer is not None:
         optimizer_state = OptimizerState(actor.model, actor.optimizer)
         optim_state_dict = {"optim_state": optimizer_state}
         dcp.save(optim_state_dict, checkpoint_id=str(optimizer_dir))
 
-    # Save LR scheduler state
-    if hasattr(actor, "lr_scheduler") and actor.lr_scheduler is not None:
+    # Save LR scheduler state (skip if --no-save-optim is set)
+    if save_optimizer_state and hasattr(actor, "lr_scheduler") and actor.lr_scheduler is not None:
         lr_scheduler_state = LRSchedulerState(actor.lr_scheduler)
         lr_scheduler_state_dict = {"lr_scheduler_state": lr_scheduler_state}
         dcp.save(lr_scheduler_state_dict, checkpoint_id=str(lr_scheduler_dir))

slime/utils/arguments.py

@@ -695,6 +695,16 @@ def add_algo_arguments(parser):
             reset_arg(parser, "--save", type=str, default=None)
             reset_arg(parser, "--save-interval", type=int, default=None)
             reset_arg(parser, "--async-save", action="store_true")
+            reset_arg(
+                parser,
+                "--no-save-optim",
+                action="store_true",
+                default=False,
+                help=(
+                    "If set, do not save the optimizer state when saving checkpoints. "
+                    "This reduces checkpoint size but disables training resumption from the saved checkpoint."
+                ),
+            )
             parser.add_argument(
                 "--save-hf",
                 type=str,