add lr_scheduler_kwargs and set default to min_lr_rate=0.1 (#23)

fixie-ai · Dec 17, 2024 · 9e4239e · 9e4239e
1 parent c357138
commit 9e4239e
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 1 deletion.
diff --git a/ultravox/training/config_base.py b/ultravox/training/config_base.py
@@ -102,6 +102,7 @@ def get_val_sets(self) -> List[DatasetOptions]:
     batch_size: int = 2
     lr: float = 1e-5
     lr_scheduler: str = "cosine"
+    lr_scheduler_kwargs: Dict[str, Any] = simple_parsing.field(default_factory=dict)
     lr_warmup_steps: int = 0
     weight_decay: float = 0.0
     seed: int = 42

diff --git a/ultravox/training/configs/meta_config.yaml b/ultravox/training/configs/meta_config.yaml
@@ -19,7 +19,8 @@ eval_max_new_tokens: 32
 eval_num_procs: 16
 
 optimizer: "adamw_torch" # options: adamw_torch, adamw_bnb_8bit
-lr_scheduler: "cosine" # options: linear, cosine, cosine_with_restarts, etc.
+lr_scheduler: "cosine_with_min_lr" # options: linear, cosine, cosine_with_restarts, etc.
+lr_scheduler_kwargs: {"min_lr_rate": 0.1}
 lr: 2.e-3
 grad_accum_steps: 1
 lr_warmup_steps: 1000

diff --git a/ultravox/training/train.py b/ultravox/training/train.py
@@ -228,6 +228,7 @@ def train(args: config_base.TrainConfig):
             ddp_find_unused_parameters=False,
             learning_rate=args.lr,
             lr_scheduler_type=args.lr_scheduler,
+            lr_scheduler_kwargs=args.lr_scheduler_kwargs,
             warmup_steps=args.lr_warmup_steps,
             weight_decay=args.weight_decay,
             # fp16=dtype == torch.float16,