fix(lr_schedules): make --lr_range_test_staircase an opt-in flag

[Anai-Guo]

What

--lr_range_test_staircase is declared with type=bool, so argparse applies the builtin to the raw argv string. Every spelling of "off" is a non-empty string, and non-empty strings are truthy:

--lr_range_test_staircase False -> True
--lr_range_test_staircase false -> True
--lr_range_test_staircase 0     -> True
--lr_range_test_staircase no    -> True
--lr_range_test_staircase True  -> True

The flag can be turned on but never off. The only value that yields False is the empty string.

Why it matters

The parsed value is not cosmetic. override_lr_range_test_params copies it straight into the scheduler config:

https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/runtime/lr_schedules.py#L143-L144

and LRRangeTest turns it into the interval function that shapes the whole LR range test:

self.interval_fn = self._staircase_interval if lr_range_test_staircase else self._continuous_interval

So a user who writes --lr_range_test_staircase False silently gets the math.floor() staircase schedule instead of the continuous one. add_tuning_arguments is public API (re-exported from deepspeed/__init__.py), and docs/_tutorials/lrrt.md documents "lr_range_test_staircase": false as a supported setting — reachable through the JSON config, but not through the CLI flag that is supposed to mirror it.

The fix

Use action='store_true', matching --cycle_momentum — the other boolean in this very same parser, with the same default=False and the same override_* path:

group.add_argument('--cycle_momentum', default=False, action='store_true', help='Enable 1Cycle momentum schedule.')

--lr_range_test_staircase was the only type=bool argument in the package; every other boolean flag already uses store_true.

Behaviour change, stated plainly

This does change the CLI surface: --lr_range_test_staircase True previously parsed to True and now errors with unrecognized arguments: True. That seemed the better trade, because the alternative spelling --lr_range_test_staircase False is currently silently wrong — a loud error is easier to fix than a schedule that quietly differs from what was asked for. Omitting the flag still yields False, so the default is unchanged and existing runs that never passed the flag are unaffected.

If you would rather keep accepting an explicit value, the alternative is a str_to_bool-style converter — happy to switch, though there is no such helper in the repo today and it would make this flag the lone exception to the store_true convention instead of the lone exception to it in the other direction.

Tests

Two additions to tests/unit/runtime/test_lr_schedulers.py (plain functions, no distributed setup):

• test_lr_range_test_staircase_is_an_opt_in_flag — the absent/present contract, asserted alongside --cycle_momentum so the two booleans stay in agreement.
• test_lr_range_test_staircase_reaches_scheduler_params — both directions survive override_lr_range_test_params into the scheduler config.

Verified against this branch: 3/3 pass with the patch; 2/3 fail without it (the False-by-default case passes either way, since that direction was never broken). yapf==0.40.0 and flake8==5.0.4 per .pre-commit-config.yaml both report clean on the two changed files.

🤖 Generated with Claude Code

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 67370efe0d

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Add the required sign-off trailer

This is a non-merge commit, but its message contains no Signed-off-by trailer, violating the repository's commit/CI requirement and potentially failing DCO enforcement. Recreate the commit using git commit --signoff with the configured author identity.

AGENTS.md reference: AGENTS.md:L6-L8

Useful? React with 👍 / 👎.