Change lr_min to min_lr_factor (#1471)

[fixed] #1457

---------

Co-authored-by: shuaipengli <[email protected]>

Author: unlimblue

tests/unit_tests/test_lr_scheduler.py

@@ -37,7 +37,7 @@ def create_job_config(
         warmup_steps=None,
         decay_ratio=None,
         decay_type=None,
-        lr_min=None,
+        min_lr_factor=None,
     ):
         # Create a job config with the specified parameters
         args = [
@@ -58,7 +58,11 @@ def create_job_config(
         args += (
             ["--lr_scheduler.decay_type", decay_type] if decay_type is not None else []
         )
-        args += ["--lr_scheduler.lr_min", str(lr_min)] if lr_min is not None else []
+        args += (
+            ["--lr_scheduler.min_lr_factor", str(min_lr_factor)]
+            if min_lr_factor is not None
+            else []
+        )
 
         config_manager = ConfigManager()
         # Create base config with parameters passed directly
@@ -74,7 +78,7 @@ def test_linear_warmup_decay(self):
             warmup_steps=2,
             decay_ratio=None,  # Use default decay: start decay immediately
             decay_type=None,
-            lr_min=None,
+            min_lr_factor=None,
         )
 
         # Build the lr scheduler
@@ -116,7 +120,7 @@ def test_warmup_stable_decay(self):
             warmup_steps=2,
             decay_ratio=0.5,  # 50% of steps for decay
             decay_type="linear",
-            lr_min=0.0,
+            min_lr_factor=0.0,
         )
 
         # Build the lr scheduler
@@ -157,7 +161,7 @@ def test_min_lr(self):
             warmup_steps=2,
             decay_ratio=None,
             decay_type="linear",
-            lr_min=0.2,  # 20% of base LR as minimum
+            min_lr_factor=0.2,  # 20% of base LR as minimum
         )
 
         # Build the lr scheduler
@@ -180,7 +184,7 @@ def test_warmup_exceeds_training(self):
             warmup_steps=10,  # More than training steps
             decay_ratio=None,
             decay_type="linear",
-            lr_min=0.0,
+            min_lr_factor=0.0,
         )
 
         # Build the lr scheduler - should adjust warmup steps
@@ -216,7 +220,7 @@ def test_warmup_stable_only(self):
             warmup_steps=2,
             decay_ratio=0.0,  # 0% of steps for decay (no decay)
             decay_type="linear",
-            lr_min=0.0,
+            min_lr_factor=0.0,
         )
 
         # Build the lr scheduler
@@ -258,7 +262,7 @@ def test_warmup_plus_decay_exceeds_training(self):
             warmup_steps=5,
             decay_ratio=0.8,  # 80% of steps for decay (8 steps)
             decay_type="linear",
-            lr_min=0.0,
+            min_lr_factor=0.0,
         )
 
         # Build the lr scheduler - should adjust warmup steps

torchtitan/components/lr_scheduler.py

@@ -127,15 +127,15 @@ def build_lr_schedulers(
     # Add a vitual last step to prevent the learning rate from dropping to 0
     stable_steps = training_steps + 1 - warmup_steps - decay_steps
     lr_decay_type = lr_scheduler_config.decay_type
-    lr_min = lr_scheduler_config.lr_min
+    min_lr_factor = lr_scheduler_config.min_lr_factor
 
     def linear_warmup_stable_decay(
         current_step: int,
         warmup_steps: int,
         stable_steps: int,
         decay_steps: int,
         lr_decay_type: str,
-        lr_min: float,
+        min_lr_factor: float,
     ):
         """
         Computes linear warmup followed by stable learning rate for a while,
@@ -150,7 +150,7 @@ def linear_warmup_stable_decay(
         2. `sqrt`: decays as 1 minus the square root of the decay progress.
         3. `cosine`: follows a cosine curve, decaying according to the values of the half-period of the cosine function.
 
-        If `lr_min` is specified, the decay range is scaled from 1 to `lr_min`
+        If `min_lr_factor` is specified, the decay range is scaled from 1 to `min_lr_factor`
         to ensure the learning rate does not drop below this minimum value.
         """
         warmup_stable_steps = warmup_steps + stable_steps
@@ -176,7 +176,7 @@ def linear_warmup_stable_decay(
                 curr_adjustment = 1 - math.sqrt(progress)
             elif lr_decay_type == "cosine":
                 curr_adjustment = 0.5 * (1.0 + math.cos(math.pi * progress))
-            curr_adjustment = lr_min + (1 - lr_min) * curr_adjustment
+            curr_adjustment = min_lr_factor + (1 - min_lr_factor) * curr_adjustment
         return curr_adjustment
 
     lr_lambda = functools.partial(
@@ -185,6 +185,6 @@ def linear_warmup_stable_decay(
         stable_steps=stable_steps,
         decay_steps=decay_steps,
         lr_decay_type=lr_decay_type,
-        lr_min=lr_min,
+        min_lr_factor=min_lr_factor,
     )
     return LRSchedulersContainer(optimizers, lr_lambda)

torchtitan/config/job_config.py

@@ -155,11 +155,11 @@ class LRScheduler:
     - 'cosine': smoothly decays learning rate following a cosine curve
     """
 
-    lr_min: float = 0.0
+    min_lr_factor: float = 0.0
     """
     Min lr ratio for lr scheduler.
-    If provided, the range of decay factor is scaled from 1 to `lr_min`
-    to ensure the learning rate does not drop below `optimizer.lr * lr_scheduler.lr_min`.
+    If provided, the range of decay factor is scaled from 1 to `min_lr_factor`
+    to ensure the learning rate does not drop below `optimizer.lr * lr_scheduler.min_lr_factor`.
     """
 
 

torchtitan/experiments/deepseek_v3/train_configs/deepseek_v2.toml

@@ -35,7 +35,7 @@ implementation = "foreach"
 warmup_steps = 100  # lr scheduler warm up, normally 20% of the train steps
 decay_ratio = 0.8  # lr scheduler decay ratio, 80% of the train steps
 decay_type = "linear"
-lr_min = 0.1
+min_lr_factor = 0.1
 
 [training]
 local_batch_size = 2 # 8

torchtitan/experiments/llama4/train_configs/debug_model.toml

@@ -34,7 +34,7 @@ eps = 1e-15
 warmup_steps = 2  # lr scheduler warm up, normally 20% of the train steps
 decay_ratio = 0.8  # lr scheduler decay ratio, 80% of the train steps
 decay_type = "linear"
-lr_min = 0.1
+min_lr_factor = 0.1
 
 [training]
 local_batch_size = 8

torchtitan/experiments/llama4/train_configs/llama4_17bx128e.toml

@@ -27,7 +27,7 @@ eps = 1e-15
 
 [lr_scheduler]
 warmup_steps = 600
-lr_min = 0.1
+min_lr_factor = 0.1
 
 [training]
 local_batch_size = 1

torchtitan/experiments/llama4/train_configs/llama4_17bx16e.toml

@@ -27,7 +27,7 @@ eps = 1e-15
 
 [lr_scheduler]
 warmup_steps = 600
-lr_min = 0.1
+min_lr_factor = 0.1
 
 [training]
 local_batch_size = 8

torchtitan/models/deepseek_v3/train_configs/debug_model.toml

@@ -36,7 +36,7 @@ eps = 1e-8
 warmup_steps = 2  # lr scheduler warm up, normally 20% of the train steps
 decay_ratio = 0.8  # lr scheduler decay ratio, 80% of the train steps
 decay_type = "linear"
-lr_min = 0.0
+min_lr_factor = 0.0
 
 [training]
 local_batch_size = 8

torchtitan/models/deepseek_v3/train_configs/deepseek_v3_16b.toml

@@ -33,8 +33,8 @@ eps = 1e-8
 [lr_scheduler]
 warmup_steps = 200  # lr scheduler warm up, normally 20% of the train steps
 decay_ratio = 0.8  # lr scheduler decay ratio, 80% of the train steps
-decay_type = "linear"
-lr_min = 2.2e-5
+decay_type = "cosine"
+min_lr_factor = 0.1
 
 [training]
 local_batch_size = 8

torchtitan/models/deepseek_v3/train_configs/deepseek_v3_671b.toml

@@ -33,8 +33,8 @@ eps = 1e-8
 [lr_scheduler]
 warmup_steps = 2_000  # lr scheduler warm up, normally 20% of the train steps
 decay_ratio = 0.8  # lr scheduler decay ratio, 80% of the train steps
-decay_type = "linear"
-lr_min = 2.2e-5
+decay_type = "cosine"
+min_lr_factor = 0.1
 
 [training]
 local_batch_size = 4