Make the OneCycle stair counts actually change the schedule

[vineethsaivs]

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.
3. 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:

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:

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:

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.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: aca90b02fc

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Preserve the peak when quantizing stairs

When stair counts are enabled with asymmetric cycle halves, this floor can quantize away the exact peak because the continuous scale at cycle_first_step_size is often just below 1 due to floating point math. For example, cycle_first_step_size=10, cycle_second_step_size=21, and cycle_first_stair_count=10 gives scale_factor=0.9999999999999998, which is floored to 0.9, so the scheduler never reaches cycle_max_lr at the cycle boundary. Please clamp values near 1 or quantize from integer batch offsets before flooring.

Useful? React with 👍 / 👎.

[tohtana]

Hi @vineethsaivs,
Thank you for submitting this PR!

The Codex's comment seems reasonable, but I think the required fix is broader: the same floating-point rounding can shift any mathematically exact interior stair boundary, not only values near 1.0.

My codex suggested this change to address it. Could you check it?

[vineethsaivs]

Thanks @tohtana, you are right that my version was too narrow, and your change is better than what I would have written. Pushed it in 22779f2.

To be concrete about why yours is the right shape: my math.floor(scale_factor * stair_count) / stair_count quantised the normalised scale, so it inherited whatever rounding error the normalisation introduced. Yours computes the stair position from the integer batch offset within the half cycle, which makes batch-aligned boundaries exact by construction instead of by luck of the rounding.

I re-ran the numbers against the patched source. With cycle_first_step_size=10, cycle_second_step_size=21:

	peak (step 10)	batch-aligned counts reproduce the continuous schedule
mine	0.0090	no
yours	0.0100	yes, exactly

So it is not only the peak. Any exact interior boundary could shift down a stair, which is the broader point you made.

I also confirmed the change does not weaken the original test: on the symmetric 8/8 cycle, counts 0 and 8 still reproduce the continuous schedule and counts 2 and 4 still hold the lr flat across each stair.

Your test_one_cycle_stair_count_handles_asymmetric_cycle is in as well, including the two halves-are-independent assertions, which my original test did not cover at all since it passed the same count to both halves. yapf and flake8 are clean.

One note for the record: your suggested diff is what landed here, so the credit for the approach is yours.

[tohtana]

Hi @vineethsaivs,
Thank you for the update and kind note! No worries, I really appreciate your contribution.