fix(training): honor max_steps when num_epochs is unset

[kabirvashisht4-glitch]

What does this PR do ?

Makes max_steps actually cap training when num_epochs is left unset.

StepSchedulerConfig.num_epochs defaulted to 10 instead of None, so the
derivation in StepScheduler.__init__ never ran on the config path the recipes
use:

if num_epochs is None:
    num_epochs = _calculate_num_epochs(max_steps, self.epoch_len)

asdict(self) in StepSchedulerConfig.build() always passed num_epochs=10
unless the user explicitly wrote num_epochs: null, leaving that branch dead.
Recipes go through this config — train_ft.py line 758 calls
self.cfg.step_scheduler.build(...).

The training loop is bounded by the epochs generator as well as by
max_steps:

@property
def epochs(self):
    for e in range(epoch, self.num_epochs):
        if self.step >= self.max_steps or self.sigterm_received:
            return
        yield e

so with num_epochs pinned at 10 the generator ran out before max_steps was
reached. A config setting only max_steps trained for 10 * epoch_len steps
and reported success — the symptom is an under-trained model, not an error.

Measured with a 100-step epoch:

config	num_epochs	steps requested	steps executed
max_steps: 5000 (before)	10	5000	1000
max_steps: 5000 (after)	50	5000	5000

The field's own docstring already described the intended behaviour ("When
None the builder derives it from max_steps"), so this makes the code match
the documented contract rather than changing it.

Changelog

• StepSchedulerConfig.num_epochs now defaults to None so max_steps drives
  the epoch count when num_epochs is not pinned.
• Attribute docstring updated to state the resolution order: derived from
  max_steps, falling back to 10 when max_steps is also unset.
• tests/unit_tests/training/test_training_config.py::test_defaults updated for
  the new default.
• Four CPU tests in
  tests/unit_tests/components/training/test_step_scheduler.py covering the
  resolution matrix: max-steps-only, num-epochs-only, both set, neither set.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

Behaviour now matches the existing docstring, so no separate doc change was
needed.

Compatibility. The both-unset case is deliberately unchanged:
_calculate_num_epochs returns its default_num_epochs=10 when max_steps is None, so a config pinning neither field still resolves to 10 epochs and the
same max_steps. test_config_neither_set_keeps_ten_epoch_fallback locks that
in. Configs that set num_epochs explicitly are unaffected.

Verification (CPU, no GPU needed):

• pytest tests/unit_tests/components/training/test_step_scheduler.py tests/unit_tests/training/test_training_config.py → 47 passed.
• Reverting only the one-line default makes
  test_config_max_steps_only_derives_num_epochs and test_defaults fail; the
  other three new tests pass either way, which is intended — they guard the
  paths this must not change.
• pytest tests/unit_tests/recipes/ tests/unit_tests/training/ tests/unit_tests/components/ → 1379 passed, 19 skipped, no regressions.
  (tests/unit_tests/recipes/dllm was excluded locally: it fails to import
  transformers.models.diffusion_gemma on main too, unrelated to this change.)
• ruff format / ruff check clean.

Additional Information

• Related to max_steps is silently ignored when num_epochs is left unset #3767

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[akoumpa]

/ok to test ed8f52a

[akoumpa]

Hi @kabirvashisht4-glitch , thanks for the PR, i'm not sure i follow 100%, i was under the impression that the max_steps' value would be derived from the num_epochs' default value (10).

[kabirvashisht4-glitch]

Thanks for looking @akoumpa — your understanding is right, and that direction is unchanged by this PR.

max_steps unset → derived from num_epochs via _calculate_max_steps, defaulting to 10 epochs. Still true after this change: _calculate_num_epochs returns its own default_num_epochs=10 when max_steps is None, so a config that pins neither field still resolves to num_epochs=10, max_steps=10*epoch_len. test_config_neither_set_keeps_ten_epoch_fallback locks that in.

The case this fixes is the opposite one: max_steps set explicitly, num_epochs left unset.

Because the dataclass default was 10 rather than None, asdict() in build() always passed num_epochs=10, so the num_epochs is None branch in __init__ was unreachable on the config path. And since the loop is bounded by the epochs generator as well as by max_steps, the generator ran out first:

for e in range(epoch, self.num_epochs):   # 10
    if self.step >= self.max_steps: return

With a 100-step epoch and max_steps: 5000, that is 10 epochs = 1000 steps. The run then reports success, so the symptom is a quietly under-trained model rather than an error:

config	num_epochs	requested	executed
max_steps: 5000 (before)	10	5000	1000
max_steps: 5000 (after)	50	5000	5000

This is also what StepScheduler.__init__ already documents:

num_epochs (int | None): Total number of epochs. Default: None or calculated from max_steps if num_epochs is None or 10 if max_steps and num_epochs are both None.
max_steps (int | None): Maximum number of steps to run. If None, calculated from num_epochs.

That describes a three-way resolution — explicit num_epochs wins; else derive from max_steps; else 10 — and the dataclass default of 10 made the middle branch dead. So this is the class's documented contract rather than new behaviour, and the resolution matrix is covered by the four tests added here.

Happy to go a different way if you'd rather keep the 10 default: the alternative is to leave it and instead raise (or warn) in __init__ when num_epochs * epoch_len < max_steps, so the conflict is loud instead of silent. That surfaces the same problem without changing any default — just say which you prefer.

For reference, full CI is green on ed8f52a (L0 CPU, both L0 GPU shards, and the L2 suites).