Skip to content

fix: leave num_epochs unset if max_steps is specified - #1107

Merged
akoumpa merged 3 commits into
NVIDIA-NeMo:mainfrom
Sparlitu:fix_scheduler_num_epochs
Jan 30, 2026
Merged

akoumpa merged 3 commits into
NVIDIA-NeMo:mainfrom
Sparlitu:fix_scheduler_num_epochs

Conversation

@Sparlitu

Copy link
Copy Markdown
Contributor

What does this PR do ?

This PR aims to fix the case where max_steps is specified in the scheduler and num_epochs is not, effectively defaulting to 10 and affecting the number of steps in training.

Changelog

  • Changed StepScheduler logic to allow num_epochs to be None if max_steps is provided.
  • Added backward compatibility check to default num_epochs to 10 if nothing is set.
  • Added a unit test test_scheduler_max_steps_priority in tests/unit_tests/components/training/test_step_scheduler.py.
  • Applied ruff formatting to modified files.

Fixes #1095

Signed-off-by: Andrei Doca <docaandrei62@yahoo.com>
@copy-pr-bot

copy-pr-bot Bot commented Jan 25, 2026

Copy link
Copy Markdown

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.

Comment thread nemo_automodel/components/training/step_scheduler.py
@chtruong814 chtruong814 added the needs-follow-up Issue needs follow-up label Jan 27, 2026
@akoumpa akoumpa removed the needs-follow-up Issue needs follow-up label Jan 27, 2026
@chtruong814 chtruong814 added the needs-follow-up Issue needs follow-up label Jan 27, 2026
Signed-off-by: Andrei Doca <docaandrei62@yahoo.com>
@Sparlitu
Sparlitu requested a review from akoumpa January 28, 2026 22:13
@akoumpa

akoumpa commented Jan 29, 2026

Copy link
Copy Markdown
Contributor

/ok to test becd773

@akoumpa

akoumpa commented Jan 29, 2026

Copy link
Copy Markdown
Contributor

/ok to test d00298e

Comment thread nemo_automodel/components/training/step_scheduler.py

@akoumpa akoumpa left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks a lot for contributing @Sparlitu!

@akoumpa
akoumpa merged commit c1b5705 into NVIDIA-NeMo:main Jan 30, 2026
49 checks passed
@Sparlitu

Copy link
Copy Markdown
Contributor Author

You're welcome @akoumpa. I hope I can contribute more for this project

linnanwang pushed a commit that referenced this pull request Apr 24, 2026
Signed-off-by: Andrei Doca <docaandrei62@yahoo.com>
Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>
kashif pushed a commit to kashif/Automodel that referenced this pull request May 25, 2026
Signed-off-by: Andrei Doca <docaandrei62@yahoo.com>
Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>

This branch was previously deployed

2 inactive deployments
nemo-ci — d00298ec Deployed Jan 29, 2026 by copy-pr-bot[bot] via L2_HF_Transformer_LLM #3401
test — d00298ec Deployed Jan 29, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #3401
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

num_epochs should not default to 10 if max_steps is specified in step scheduler

3 participants