DeepSpeed

8d4074cb - fix(lr_schedules): make --lr_range_test_staircase an opt-in flag (#8337)

fix(lr_schedules): make --lr_range_test_staircase an opt-in flag (#8337) ## 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: ```python 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: ```python 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](https://claude.com/claude-code) Signed-off-by: Anai-Guo <antai12232931@outlook.com> Signed-off-by: Masahiro Tanaka <tanaka.masahiro@gmail.com> Co-authored-by: Masahiro Tanaka <tanaka.masahiro@gmail.com>