Make the WarmupCosineLR ratio flags reach the config

deepspeed/runtime/lr_schedules.py

@@ -195,6 +195,23 @@ def override_warmupLR_params(args, params):
         params[WARMUP_TYPE] = args.warmup_type
 
 
+def override_warmupCosineLR_params(args, params):
+    # WarmupCosineLR scales each param group's own lr by a ratio, so it takes
+    # warmup_min_ratio/cos_min_ratio and does not accept warmup_min_lr or
+    # warmup_max_lr.
+    if hasattr(args, WARMUP_NUM_STEPS) and args.warmup_num_steps is not None:
+        params[WARMUP_NUM_STEPS] = args.warmup_num_steps
+
+    if hasattr(args, WARMUP_TYPE) and args.warmup_type is not None:
+        params[WARMUP_TYPE] = args.warmup_type
+
+    if hasattr(args, WARMUP_MIN_RATIO) and args.warmup_min_ratio is not None:
+        params[WARMUP_MIN_RATIO] = args.warmup_min_ratio
+
+    if hasattr(args, COS_MIN_RATIO) and args.cos_min_ratio is not None:
+        params[COS_MIN_RATIO] = args.cos_min_ratio
+
+
 def override_params(args, params):
     # LR range test params
     override_lr_range_test_params(args, params)
@@ -205,6 +222,9 @@ def override_params(args, params):
     # WarmupLR params
     override_warmupLR_params(args, params)
 
+    # WarmupCosineLR params
+    override_warmupCosineLR_params(args, params)
+
 
 def get_config_from_args(args):
     if not hasattr(args, LR_SCHEDULE) or args.lr_schedule is None:
@@ -221,6 +241,8 @@ def get_config_from_args(args):
         override_lr_range_test_params(args, config['params'])
     elif args.lr_schedule == ONE_CYCLE:
         override_1cycle_params(args, config['params'])
+    elif args.lr_schedule == WARMUP_COSINE_LR:
+        override_warmupCosineLR_params(args, config['params'])
     else:
         override_warmupLR_params(args, config['params'])
 
@@ -244,6 +266,9 @@ def get_lr_from_config(config):
         return lr_params[LR_RANGE_TEST_MIN_LR], ''
     if lr_schedule == ONE_CYCLE:
         return lr_params[CYCLE_MAX_LR], ''
+    if lr_schedule == WARMUP_COSINE_LR:
+        return None, '{} scales the optimizer learning rate by a ratio, so its params define no lr'.format(
+            WARMUP_COSINE_LR)
     # Warmup LR
     return lr_params[WARMUP_MAX_LR], ''
 

tests/unit/runtime/test_lr_schedulers.py

@@ -3,6 +3,7 @@
 
 # DeepSpeed Team
 
+import argparse
 import math
 
 import torch
@@ -874,3 +875,52 @@ def test_one_cycle_rejects_wrong_length_per_group_lists(kwargs):
 
     with pytest.raises(ValueError):
         OneCycle(optimizer=optimizer, **{**defaults, **kwargs})
+
+
+def test_warmup_cosine_lr_config_from_args_carries_the_ratios():
+    # --warmup_min_ratio and --cos_min_ratio are declared by add_tuning_arguments and
+    # were then dropped: WarmupCosineLR fell into the WarmupLR branch, so the config
+    # came back with warmup_min_lr/warmup_max_lr instead, which WarmupCosineLR does
+    # not accept and which made building it from that config a TypeError.
+    parser = lrs.add_tuning_arguments(argparse.ArgumentParser())
+    args = parser.parse_args(
+        ["--lr_schedule", WARMUP_COSINE_LR, "--warmup_min_ratio", "0.1", "--cos_min_ratio", "0.05"])
+
+    config, err = lrs.get_config_from_args(args)
+    assert err is None
+    params = config["params"]
+
+    assert params[WARMUP_MIN_RATIO] == 0.1
+    assert params[COS_MIN_RATIO] == 0.05
+    assert WARMUP_MIN_LR not in params
+    assert WARMUP_MAX_LR not in params
+
+    param = torch.nn.Parameter(torch.zeros(1))
+    optimizer = torch.optim.Adam([param], lr=0.001)
+    scheduler = WarmupCosineLR(optimizer=optimizer, total_num_steps=2000, **params)
+    assert scheduler.warmup_min_ratio == 0.1
+    assert scheduler.cos_min_ratio == 0.05
+
+
+def test_warmup_cosine_lr_has_no_lr_in_its_config_params():
+    # get_lr_from_config read warmup_max_lr for every non-LRRangeTest, non-OneCycle
+    # schedule. WarmupCosineLR scales each param group's own lr by a ratio, so its
+    # params never hold one, and the lookup either returned an unrelated WarmupLR
+    # default or raised KeyError on a config written for this schedule.
+    config = {"type": WARMUP_COSINE_LR, "params": {WARMUP_MIN_RATIO: 0.1, COS_MIN_RATIO: 0.05}}
+
+    lr, err = lrs.get_lr_from_config(config)
+    assert lr is None
+    assert WARMUP_COSINE_LR in err
+
+
+def test_other_schedules_keep_their_config_params():
+    # The routing change must leave the three schedules that already worked alone.
+    parser = lrs.add_tuning_arguments(argparse.ArgumentParser())
+
+    for schedule, expected in ((LR_RANGE_TEST, LR_RANGE_TEST_MIN_LR), (ONE_CYCLE, CYCLE_MAX_LR),
+                               (WARMUP_LR, WARMUP_MAX_LR), (WARMUP_DECAY_LR, WARMUP_MAX_LR)):
+        config, err = lrs.get_config_from_args(parser.parse_args(["--lr_schedule", schedule]))
+        assert err is None
+        assert expected in config["params"]
+        assert lrs.get_lr_from_config(config)[0] == config["params"][expected]