diff --git a/deepspeed/runtime/lr_schedules.py b/deepspeed/runtime/lr_schedules.py index 96f25f4c05d0..edfc3a2b5ece 100644 --- a/deepspeed/runtime/lr_schedules.py +++ b/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], '' diff --git a/tests/unit/runtime/test_lr_schedulers.py b/tests/unit/runtime/test_lr_schedulers.py index ec43ee0d54ab..7889dbb8e25f 100644 --- a/tests/unit/runtime/test_lr_schedulers.py +++ b/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]