DeepSpeed

3169c2dc - Make the OneCycle stair counts actually change the schedule (#8179)

Make the OneCycle stair counts actually change the schedule (#8179) ## Problem `cycle_first_stair_count` and `cycle_second_stair_count` do nothing. Every value produces the same learning rate schedule. They are not obscure knobs. They are documented in the 1Cycle tutorial: > 2. `cycle_first_stair_count`: The count of updates (or stairs) in first step of cycle phase. > 4. `cycle_second_stair_count`: The count of updates (or stairs) in the second step of cycle phase. exposed as CLI flags (`--cycle_first_stair_count`, `--cycle_second_stair_count`), plumbed through `get_config_from_args` into the scheduler params, documented on `OneCycle` itself ("Number of stairs in first half of cycle phase. This means lr/mom are changed in staircase fashion. Default 0, means staircase disabled."), and set in four of the repo's own test configs (`tests/unit/runtime/half_precision/test_fp16.py`, `tests/unit/v1/half_precision/test_bf16.py`, `tests/unit/checkpoint/test_pipeline.py`, `tests/unit/checkpoint/test_other_optimizer.py`). `_initialize_cycle` stores them: ```python self.first_stair_count = cycle_first_stair_count self.second_stair_count = cycle_first_stair_count if cycle_second_stair_count is None else cycle_second_stair_count ``` and `self.first_stair_count` / `self.second_stair_count` are read nowhere in the repository. `_get_scale_factor` is unconditionally continuous: ```python if x <= self.step_ratio: scale_factor = x / self.step_ratio else: scale_factor = (x - 1) / (self.step_ratio - 1) return scale_factor ``` ## Reproduction `cycle_min_lr=0.001`, `cycle_max_lr=0.01`, `cycle_first_step_size=8`, `cycle_second_step_size=8`: ``` stair_count=0 [0.002125, 0.00325, 0.004375, 0.0055, 0.006625, 0.00775, 0.008875, 0.01, ...] stair_count=2 identical stair_count=4 identical stair_count=8 identical ``` A user who sets a stair count silently gets the continuous schedule, with no warning. ## Fix Quantise the half-cycle scale factor into `stair_count` steps, using the same `floor()` that `LRRangeTest._staircase_interval` already uses for the LR range test staircase: ```python if stair_count > 0: scale_factor = math.floor(scale_factor * stair_count) / stair_count ``` Because `_get_scale_factor` backs both `_get_cycle_lr` and `_get_cycle_mom`, this makes lr and momentum staircase together, which is what the docstring describes. After the fix, on the same config: ``` stair_count=0 [0.002125, 0.00325, 0.004375, 0.0055, 0.006625, 0.00775, 0.008875, 0.01, ...] 8 distinct values stair_count=2 [0.001, 0.001, 0.001, 0.0055, 0.0055, 0.0055, 0.0055, 0.01, ...] 3 distinct values stair_count=4 [0.001, 0.00325, 0.00325, 0.0055, 0.0055, 0.00775, 0.00775, 0.01, ...] 5 distinct values ``` `stair_count=0` is byte-identical to before, so the default path is untouched. A stair count equal to the half-cycle length is also the continuous schedule, which is the correct degenerate case: quantising 8 batches into 8 steps changes nothing. Documentation needs no update. The tutorial and the docstring already describe this behaviour; only the code was missing. The exact quantisation is the one design choice here. I matched `LRRangeTest`, which is the only other staircase in this file, and kept it to a single expression so it is easy to change if you would rather the stairs be derived from the batch index than from the scale factor. ## Testing `test_one_cycle_stair_count_holds_lr_flat` in `tests/unit/runtime/test_lr_schedulers.py`, added next to the existing non-distributed scheduler tests. It asserts the continuous schedule still varies every batch, that a stair count of 2 differs from it, changes the lr less often, and is flat within a stair, and that a stair count equal to the half-cycle length reproduces the continuous schedule. It fails on master with `AssertionError: stair count still has no effect` and passes with this change. `yapf --style .style.yapf` and `flake8 --config .flake8` clean on both files. Commit is signed off. --------- Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com> Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com> Co-authored-by: Masahiro Tanaka <mtanaka@anyscale.com>