From 01e73b59344214ed69bfc9dc9cb0981a5f239c20 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Thu, 10 Sep 2026 18:07:07 +0800 Subject: [PATCH] Make add_tuning_arguments' 1Cycle flags reach OneCycle `add_tuning_arguments` declares the 1Cycle flags and `override_1cycle_params` copies them into the scheduler config, but three of them never arrive intact. The plain CLI invocation for this schedule does not build at all. Three flags use -1 as an "unset" sentinel, and the copy guard is `is not None`, so -1 is forwarded as a real value: $ parser = add_tuning_arguments(argparse.ArgumentParser()) $ args = parser.parse_args(["--lr_schedule", "OneCycle", "--cycle_min_lr", "1e-4", "--cycle_max_lr", "1e-3", "--cycle_first_step_size", "4"]) $ OneCycle(optimizer, **get_config_from_args(args)[0]["params"]) ValueError: cycle_second_step_size must be non-negative, got -1.0 `cycle_second_stair_count` has the same shape without the error: `None` means "same as the first half", and -1 is a real value that is simply not > 0, so `--cycle_first_stair_count 5` gives the first half stairs and the second none. `--cycle_momentum` was declared but never copied. OneCycle defaults it to True, so the flag could neither switch momentum cycling on (already on) nor off (what its own default asks for). On an Adam-like optimizer the schedule rewrites `param_groups["betas"][0]` every step, and the trajectory was identical either way -- beta1 0.9 -> 0.8 on the first step for a caller who never asked: --cycle_momentum omitted scheduler.cycle_momentum=True [0.8, 0.9425, 0.895, ...] --cycle_momentum passed scheduler.cycle_momentum=True [0.8, 0.9425, 0.895, ...] Behaviour change worth naming: a CLI-built OneCycle no longer cycles momentum unless `--cycle_momentum` is passed. That is what the flag documents, and the old default was unreachable rather than chosen. Only this helper is affected -- `engine._configure_lr_scheduler` passes a JSON config straight through, so `"params": {"cycle_momentum": false}` has always worked. The JSON path and the other four schedules are untouched. This is the sibling of #8268 and #8337 in the same helper; the flag-parsing test added by #8337 even cites --cycle_momentum as the correct shape, but nothing checked that it reaches the scheduler. Co-Authored-By: Claude Opus 5 Signed-off-by: alanhuangyoo --- deepspeed/runtime/lr_schedules.py | 12 ++-- tests/unit/runtime/test_lr_schedulers.py | 89 ++++++++++++++++++++++++ 2 files changed, 97 insertions(+), 4 deletions(-) diff --git a/deepspeed/runtime/lr_schedules.py b/deepspeed/runtime/lr_schedules.py index 7b8b27157e1d..b6f169f34209 100644 --- a/deepspeed/runtime/lr_schedules.py +++ b/deepspeed/runtime/lr_schedules.py @@ -41,6 +41,7 @@ CYCLE_MAX_LR = 'cycle_max_lr' DECAY_LR_RATE = 'decay_lr_rate' +CYCLE_MOMENTUM = 'cycle_momentum' CYCLE_MIN_MOM = 'cycle_min_mom' CYCLE_MAX_MOM = 'cycle_max_mom' DECAY_MOM_RATE = 'decay_mom_rate' @@ -80,16 +81,16 @@ def add_tuning_arguments(parser): help='size of first step of 1Cycle schedule (training steps).') group.add_argument("--cycle_first_stair_count", type=int, - default=-1, + default=None, help='first stair count for 1Cycle schedule.') group.add_argument("--cycle_second_step_size", type=int, - default=-1, + default=None, help='size of second step of 1Cycle schedule (default first_step_size).') group.add_argument("--cycle_second_stair_count", type=int, - default=-1, - help='second stair count for 1Cycle schedule.') + default=None, + help='second stair count for 1Cycle schedule (default first stair count).') group.add_argument("--decay_step_size", type=int, default=1000, @@ -171,6 +172,9 @@ def override_1cycle_params(args, params): params[DECAY_LR_RATE] = args.decay_lr_rate # 1Cycle MOM params + if hasattr(args, CYCLE_MOMENTUM): + params[CYCLE_MOMENTUM] = args.cycle_momentum + if hasattr(args, CYCLE_MIN_MOM) and args.cycle_min_mom is not None: params[CYCLE_MIN_MOM] = args.cycle_min_mom diff --git a/tests/unit/runtime/test_lr_schedulers.py b/tests/unit/runtime/test_lr_schedulers.py index e3d095e9c693..0ad9b962b9ce 100644 --- a/tests/unit/runtime/test_lr_schedulers.py +++ b/tests/unit/runtime/test_lr_schedulers.py @@ -15,6 +15,7 @@ from deepspeed.runtime.lr_schedules import WARMUP_LR, WARMUP_MIN_LR, WARMUP_MAX_LR, WARMUP_NUM_STEPS, WARMUP_TYPE, WARMUP_LOG_RATE, WARMUP_LINEAR_RATE from deepspeed.runtime.lr_schedules import ONE_CYCLE, CYCLE_MIN_LR, CYCLE_MAX_LR, CYCLE_FIRST_STEP_SIZE, DECAY_LR_RATE, DECAY_STEP_SIZE from deepspeed.runtime.lr_schedules import CYCLE_MIN_MOM, CYCLE_MAX_MOM, DECAY_MOM_RATE +from deepspeed.runtime.lr_schedules import CYCLE_MOMENTUM, CYCLE_FIRST_STAIR_COUNT, CYCLE_SECOND_STEP_SIZE, CYCLE_SECOND_STAIR_COUNT from deepspeed.runtime.lr_schedules import WARMUP_DECAY_LR, TOTAL_NUM_STEPS from deepspeed.runtime.lr_schedules import WARMUP_COSINE_LR, WARMUP_MIN_RATIO, COS_MIN_RATIO, WarmupCosineLR from deepspeed.runtime.lr_schedules import WarmupLR, WarmupDecayLR, LRRangeTest, OneCycle @@ -949,3 +950,91 @@ def test_other_schedules_keep_their_config_params(): assert err is None assert expected in config["params"] assert lrs.get_lr_from_config(config)[0] == config["params"][expected] + + +def test_one_cycle_config_from_args_builds_a_scheduler(): + # The three "unset" flags defaulted to -1 and override_1cycle_params copied them + # through, since -1 is not None. OneCycle rejects a negative second step size, so + # the plain CLI invocation for this schedule raised before it ever ran: + # ValueError: cycle_second_step_size must be non-negative, got -1.0 + parser = lrs.add_tuning_arguments(argparse.ArgumentParser()) + args = parser.parse_args([ + "--lr_schedule", ONE_CYCLE, "--cycle_min_lr", "1e-4", "--cycle_max_lr", "1e-3", + "--cycle_first_step_size", "4" + ]) + + config, err = lrs.get_config_from_args(args) + assert err is None + params = config["params"] + # Left out entirely, so OneCycle's own defaults apply rather than a sentinel. + assert CYCLE_SECOND_STEP_SIZE not in params + assert CYCLE_FIRST_STAIR_COUNT not in params + assert CYCLE_SECOND_STAIR_COUNT not in params + + optimizer = torch.optim.Adam([torch.nn.Parameter(torch.zeros(1))], lr=1e-3) + scheduler = OneCycle(optimizer, **params) + # The help text says the second step defaults to the first. + assert scheduler.second_step_size == scheduler.first_step_size == 4 + + +def test_one_cycle_second_stair_count_falls_back_to_the_first(): + # cycle_second_stair_count=None means "same as the first"; the -1 default reached + # OneCycle as a real value, so the second half of the cycle lost its stairs while + # the first half kept them. + parser = lrs.add_tuning_arguments(argparse.ArgumentParser()) + args = parser.parse_args([ + "--lr_schedule", ONE_CYCLE, "--cycle_min_lr", "1e-4", "--cycle_max_lr", "1e-3", + "--cycle_first_step_size", "4", "--cycle_first_stair_count", "5" + ]) + + config, _ = lrs.get_config_from_args(args) + optimizer = torch.optim.Adam([torch.nn.Parameter(torch.zeros(1))], lr=1e-3) + scheduler = OneCycle(optimizer, **config["params"]) + + assert scheduler.first_stair_count == 5 + assert scheduler.second_stair_count == 5 + + direct = OneCycle(torch.optim.Adam([torch.nn.Parameter(torch.zeros(1))], lr=1e-3), + cycle_min_lr=1e-4, + cycle_max_lr=1e-3, + cycle_first_step_size=4, + cycle_first_stair_count=5) + assert scheduler.second_stair_count == direct.second_stair_count + + +@pytest.mark.parametrize("argv, expected", [([], False), (["--cycle_momentum"], True)]) +def test_cycle_momentum_reaches_scheduler_params(argv, expected): + # --cycle_momentum was declared but never copied into the config, so OneCycle fell + # back to its own cycle_momentum=True either way: the flag could not turn momentum + # cycling on (it was already on) or off (its documented default). + args = lrs.add_tuning_arguments(argparse.ArgumentParser()).parse_args(argv) + params = {} + lrs.override_1cycle_params(args, params) + + assert params[CYCLE_MOMENTUM] is expected + + +@pytest.mark.parametrize("argv, cycles", [([], False), (["--cycle_momentum"], True)]) +def test_cycle_momentum_flag_decides_whether_betas_move(argv, cycles): + # The observable half: OneCycle rewrites param_groups["betas"][0] on every step when + # momentum cycling is on. Without the flag the optimizer's own beta1 must survive. + parser = lrs.add_tuning_arguments(argparse.ArgumentParser()) + args = parser.parse_args([ + "--lr_schedule", ONE_CYCLE, "--cycle_min_lr", "1e-4", "--cycle_max_lr", "1e-3", + "--cycle_first_step_size", "4", "--cycle_min_mom", "0.80", "--cycle_max_mom", "0.99" + ] + argv) + + config, _ = lrs.get_config_from_args(args) + optimizer = torch.optim.Adam([torch.nn.Parameter(torch.zeros(1))], lr=1e-3, betas=(0.9, 0.999)) + scheduler = OneCycle(optimizer, **config["params"]) + assert scheduler.cycle_momentum is cycles + + seen = {optimizer.param_groups[0]["betas"][0]} + for _ in range(4): + scheduler.step() + seen.add(optimizer.param_groups[0]["betas"][0]) + + if cycles: + assert len(seen) > 1 + else: + assert seen == {0.9}