NVIDIA-NeMo / NVIDIA-NeMo/Automodel
max_steps is silently ignored when num_epochs is left unset
Nobody has claimed this yet.
- 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. WhenNonethe builder derives it
frommax_steps. Default: 10.
max_steps: Hard cap on optimizer steps.Nonemeans 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
mainat 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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