Various updates

1 files changed, 8 insertions, 0 deletions

diff --git a/train_ti.py b/train_ti.py
index 6f116c3..1b60f64 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -260,6 +260,12 @@ def parse_args():
         help='If lr_annealing_func is "half_cos" or "cos", exponent to modify the function'
     )
     parser.add_argument(
+        "--lr_min_lr",
+        type=float,
+        default=None,
+        help="Minimum learning rate in the lr scheduler."
+    )
+    parser.add_argument(
         "--use_8bit_adam",
         action="store_true",
         help="Whether or not to use 8-bit Adam from bitsandbytes."
@@ -744,6 +750,7 @@ def main():
     if args.find_lr:
         lr_scheduler = None
     elif args.lr_scheduler == "one_cycle":
+        lr_min_lr = 0.04 if args.lr_min_lr is None else args.lr_min_lr / args.learning_rate
         lr_scheduler = get_one_cycle_schedule(
             optimizer=optimizer,
             num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
@@ -751,6 +758,7 @@ def main():
             annealing=args.lr_annealing_func,
             warmup_exp=args.lr_warmup_exp,
             annealing_exp=args.lr_annealing_exp,
+            min_lr=lr_min_lr,
         )
     elif args.lr_scheduler == "cosine_with_restarts":
         lr_scheduler = get_cosine_with_hard_restarts_schedule_with_warmup(