NVIDIA-NeMo / NVIDIA-NeMo/Automodel

max_steps is silently ignored when num_epochs is left unset

Open Beginner friendly
#3,767 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

community-request waiting-on-maintainers
Dominant language
Python
Stars
960
Forks
316
Avg merge
3d 20h
Merged PRs (30d)
143

Description

Describe the bug

StepSchedulerConfig.num_epochs defaults to 10 instead of None, so the
max_steps-based derivation in StepScheduler.__init__ never runs on the
config path the recipes use. A config that sets only max_steps silently trains
for 10 * epoch_len steps instead of max_steps.

StepSchedulerConfig documents the intended contract
(nemo_automodel/components/training/step_scheduler.py):

num_epochs: Number of training epochs. When None the builder derives it
from max_steps. Default: 10.
max_steps: Hard cap on optimizer steps. None means derive from
num_epochs * epoch_len.

and StepScheduler.__init__ implements that derivation:

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

But the dataclass field is:

num_epochs: int | None = 10

asdict(self) in StepSchedulerConfig.build() therefore always passes
num_epochs=10 unless the user explicitly writes num_epochs: null, and the
num_epochs is None branch is dead on this path. Recipes build the scheduler
through this config — train_ft.py line 758 calls
self.cfg.step_scheduler.build(...).

The truncation happens because the training loop is bounded by epochs, not
only 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

With num_epochs stuck at 10, the generator is exhausted after 10 epochs and
max_steps is never reached. There is no warning.

Steps/Code to reproduce bug

step_scheduler:
  global_batch_size: 1
  local_batch_size: 1
  max_steps: 5000
  # num_epochs intentionally not set

Standalone repro (CPU, no GPU or tokenizer needed):

from nemo_automodel.components.training.step_scheduler import StepSchedulerConfig

class DL(list):
    pass

dl = DL(range(100))  # 100 micro-batches; grad_acc = 1 -> epoch_len = 100

cfg = StepSchedulerConfig(global_batch_size=1, max_steps=5000, preemption_signal=None)
ss = cfg.build(dataloader=dl, dp_group_size=1, local_batch_size=1)

total = 0
for epoch in ss.epochs:
    ss.set_epoch(epoch)
    for _ in ss:
        total += 1

print("num_epochs :", ss.num_epochs)   # 10   (expected 50 = ceil(5000/100))
print("max_steps  :", ss.max_steps)    # 5000
print("executed   :", total)           # 1000 -- 20% of the requested steps

Output:

num_epochs : 10
max_steps  : 5000
executed   : 1000

Passing num_epochs=None explicitly gives the documented behaviour:
num_epochs becomes 50 and all 5000 steps run.

Expected behavior

A config that sets max_steps and leaves num_epochs unset runs for
max_steps optimizer steps, as the docstring describes. Leaving both unset
should keep today's behaviour (10 epochs).

Environment overview

  • main at 3ddef9b1, CPU only — no GPU, dataset, or tokenizer required.

Additional context

Setting the dataclass default to None restores the documented behaviour and
leaves the both-unset case unchanged: _calculate_num_epochs already returns
its default_num_epochs=10 when max_steps is None, so a config that sets
neither field still resolves to 10 epochs.

Note this fails quietly — the run reports success after 1000 steps, so the
symptom is an under-trained model rather than an error. Happy to send a PR with
CPU unit tests covering max-steps-only, epochs-only, both, and neither.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in nemo_automodel/components/training/step_scheduler.py, reading StepSchedulerConfig and StepScheduler.init, then inspect the build entry point and the training-loop epochs property. Use the standalone CPU reproduction to verify max_steps-only, epochs-only, both-unset, and explicit-null cases. Done means leaving num_epochs unset derives enough epochs for max_steps while leaving both fields unset at 10 epochs.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, machine-learning
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
86/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.