Make add_tuning_arguments' 1Cycle flags reach OneCycle

[alanhuangyoo]

add_tuning_arguments declares the 1Cycle flags and override_1cycle_params copies them into the scheduler config. Three of them do not survive the trip.

The plain CLI invocation does not build

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"])
>>> get_config_from_args(args)[0]["params"]
{'cycle_first_step_size': 4, 'cycle_first_stair_count': -1, 'cycle_second_step_size': -1,
 'cycle_second_stair_count': -1, 'decay_step_size': 1000, ...}
>>> OneCycle(optimizer, **_)
ValueError: cycle_second_step_size must be non-negative, got -1.0

-1 was survivable before #8166 added the step-size guards; it produced total_size = first + (-1) instead. It is now an error, and this path has no way around it short of passing the flag explicitly.

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 simply is not > 0, so --cycle_first_stair_count 5 gives the first half stairs and the second half none:

	via the CLI helper	built directly
first_stair_count	5	5
second_stair_count	-1	5

--cycle_momentum was never copied

It is declared, parsed, and dropped. OneCycle defaults cycle_momentum=True, so the flag could neither turn momentum cycling on (already on) nor off (what its own default=False asks for).

On an Adam-like optimizer the schedule rewrites param_groups["betas"][0] every step. Measured with --cycle_min_mom 0.80 --cycle_max_mom 0.99:

--cycle_momentum omitted   args=False  reached config: False  scheduler.cycle_momentum=True  beta1: [0.8, 0.9425, 0.895, 0.8475, 0.8]
--cycle_momentum passed    args=True   reached config: False  scheduler.cycle_momentum=True  beta1: [0.8, 0.9425, 0.895, 0.8475, 0.8]

Identical, and beta1 goes 0.9 -> 0.8 on the first step for a caller who never asked for it.

The fix

• the three sentinels default to None, so the existing is not None guard leaves them out and OneCycle's own defaults apply — which is what the help text already promises ("default first_step_size")
• CYCLE_MOMENTUM constant added next to the other 1Cycle momentum keys, and override_1cycle_params copies the flag

Behaviour change worth naming: a CLI-built OneCycle no longer cycles momentum unless --cycle_momentum is passed. That is what the flag documents; the old default was unreachable rather than chosen. store_true cannot distinguish "omitted" from "given as False", so the flag is copied unconditionally — there is no third state to preserve.

Scope: only this helper. engine._configure_lr_scheduler passes a JSON config straight to the scheduler, so "params": {"cycle_momentum": false} has always worked, and the other four schedules are untouched.

Testing

tests/unit/runtime/test_lr_schedulers.py, H20:

master   88 passed, 0 failed
branch   94 passed, 0 failed
only failing on branch (regressions) -> none
only failing on master               -> none

Six new cases. Verified they fail on master for the right reason, rather than only that they pass here — run against master's source with just the new constant added so the import resolves:

test_one_cycle_config_from_args_builds_a_scheduler        AssertionError: 'cycle_second_step_size' not in {... 'cycle_second_step_size': -1 ...}
test_one_cycle_second_stair_count_falls_back_to_the_first ValueError: cycle_second_step_size must be non-negative, got -1.0
test_cycle_momentum_reaches_scheduler_params[False]       KeyError: 'cycle_momentum'
test_cycle_momentum_reaches_scheduler_params[True]        KeyError: 'cycle_momentum'
test_cycle_momentum_flag_decides_whether_betas_move[False] ValueError: cycle_second_step_size must be non-negative, got -1.0
test_cycle_momentum_flag_decides_whether_betas_move[True]  ValueError: cycle_second_step_size must be non-negative, got -1.0

The last two report the sentinel rather than the momentum assertion because on master they cannot get far enough to reach it — the beta1 table above is from a probe that supplies --cycle_second_step_size explicitly to step over the first defect.

Related

This is the third of the same shape in this helper. #8268 fixed --warmup_min_ratio/--cos_min_ratio being dropped for WarmupCosineLR; #8337 fixed --lr_range_test_staircase being un-turn-off-able. The flag-parsing test #8337 added even cites --cycle_momentum as the correct store_true shape — but nothing checked that it reaches the scheduler, which is the half that was broken.

🤖 Generated with Claude Code