add delete old checkpoints

remove --train_ln because we never changed it

peft_pretraining/args_utils.py

@@ -31,10 +31,6 @@ def check_args_torchrun_main(args):
     if args.batch_size is None:
         raise ValueError("batch_size must be specified")
 
-    if not args.train_ln:
-        logger.error("Are you sure? Not training LN is a bad idea.")
-        raise ValueError("Are you sure? Not training LN is a bad idea.")
-
     if args.tags is not None:
         args.tags = args.tags.split(",")
 

peft_pretraining/training_utils.py

@@ -402,3 +402,17 @@ def check_lr_and_alert(optimizer, max_lr):
             text=alert_message,
             level=wandb.AlertLevel.WARN,
         )
+
+def delete_old_checkpoints(save_dir, keep):
+    if keep is None:
+        return
+
+    checkpoints = [d for d in os.listdir(save_dir) if d.startswith(f"model_")]
+    if len(checkpoints) <= keep:
+        return
+
+    checkpoints = sorted(checkpoints, key=lambda x: int(x.split("_")[-1]))
+    for checkpoint in checkpoints[:-keep]:
+        checkpoint_path = os.path.join(save_dir, checkpoint)
+        logger.info(f"Deleting checkpoint {checkpoint_path}")
+        os.system(f"rm -rf {checkpoint_path}")

torchrun_main.py

@@ -88,7 +88,6 @@ def parse_args(args=None):
                         help=("Keep original model parameters even if relora is None. "
                               "Useful for making sure that full-LoRa model is equivalent to model+LoRa."))
 
-    parser.add_argument("--train_ln", default=True, action="store_true")
     parser.add_argument("--optimizer", default="Adam", help="Could be adam (for AdamW) or adam_zero for ZeroRedundancyOptimizer(AdamW)")
     parser.add_argument("--lr", type=float, default=1e-4)
     parser.add_argument("--scheduler", type=str, default="cosine", choices=["linear", "cosine", "cosine_restarts"])
@@ -114,6 +113,8 @@ def parse_args(args=None):
                              "You can use M and B suffixes, e.g. 100M or 1B.")
     parser.add_argument("--save_every", type=int, default=10_000)
     parser.add_argument("--save_dir", type=str, default=None)
+    parser.add_argument("--keep_checkpoints", type=int, default=None,
+                        help="Number of checkpoints to keep. By default, keep all checkpoints.")
     parser.add_argument("--tags", type=str, default=None)
     parser.add_argument("--dtype", type=str, default="bfloat16" if torch.cuda.is_bf16_supported() else "float32")
     parser.add_argument("--workers", type=int, default=8)
@@ -606,6 +607,13 @@ def main(args):
     trainable_params = [p for p in model.parameters() if p.requires_grad]
     lora_params = [p for n, p in model.named_parameters() if p.requires_grad and "lora_" in n]
     trainable_params_names = [name for name, p in model.named_parameters() if p.requires_grad]
+    non_trainable_params_names = [name for name, p in model.named_parameters() if not p.requires_grad]
+
+    logger.info("*" * 40)
+    logger.info("Non-trainable paramters:")
+    for name in non_trainable_params_names:
+        logger.info(f"{name:40}")
+    logger.info("*" * 40)
 
     if args.use_peft and len(lora_params) == 0:
         raise ValueError("No LoRA parameters found")
@@ -816,6 +824,8 @@ def main(args):
                 distributed_type=args.distributed_type,
                 save_dir=current_model_directory,
             )
+            if args.keep_checkpoints is not None:
+                training_utils.delete_old_checkpoints(args.save_dir, keep=args.keep_checkpoints)
 
         # ##############################
         # EVALUATION