cp: fix: fix mcore train_iters in grpo (1383) into r0.4.0

[chtruong814]

beep boop [🤖]: Hi @yuki-97 👋,

we've cherry picked #1383 into  for you! 🚀

Please review and approve this cherry pick by your convenience!

Summary by CodeRabbit

• Bug Fixes
  • Updated training iteration calculation to incorporate epoch count, enabling more training steps per Megatron run while respecting maximum step limits.

[coderabbitai]

📝 Walkthrough

Walkthrough

Modifies total_train_iters calculation in GRPO algorithm to incorporate epoch count: changes from len(dataloader) to min(grpo_config["max_num_steps"], grpo_config["max_num_epochs"] * len(dataloader)), allowing more training steps based on epoch configuration.

Changes

Cohort / File(s)	Summary
GRPO training iteration calculation nemo_rl/algorithms/grpo.py	Updated total_train_iters computation to use grpo_config["max_num_epochs"] * len(dataloader) as upper bound instead of len(dataloader), capped by grpo_config["max_num_steps"]

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Possibly related PRs

• fix: fix mcore train_iters in grpo #1383: Modifies the same total_train_iters calculation in nemo_rl/algorithms/grpo.py using the identical formula combining max_num_steps and max_num_epochs * len(dataloader).

Suggested labels

CI:L1, r0.4.0

Suggested reviewers

• yuki-97
• terrykong
• ashors1

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	While this is a focused bug fix to a single calculation in the Megatron backend's training iteration setup, the change affects how training iterations are computed, which directly influences convergence behavior. The original PR #1383 was merged into main but the commit message contains no test results or validation details. The cherry-pick PR #1385 similarly provides no reference to testing evidence, only stating it is a cherry-pick into r0.4.0. Per the custom check requirements, changes affecting numerics or convergence should include evidence of no regression, and the PR description should document test results or reference validation from the original PR.	The PR description for this cherry-pick should be enhanced to reference the testing and validation performed on the original PR #1383, or include test results demonstrating the fix resolves the training iteration calculation bug without regression. At minimum, add a note that references PR #1383's approval/merge, or provide validation showing training with the Megatron backend now correctly incorporates max_num_epochs in the iteration count.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title Check	✅ Passed	The substantive part of the title, "fix mcore train_iters in grpo," directly corresponds to the main change in the pull request. The changeset modifies the calculation of total_train_iters for Megatron in nemo_rl/algorithms/grpo.py, which is precisely what the title indicates. While the title includes cherry-pick metadata (the "cp:" prefix and "into r0.4.0" notation) and the reference to issue #1383, these do not obscure the core message. The title is sufficiently specific to convey the primary fix to a teammate scanning the commit history.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch cherry-pick-1383-r0.4.0

---

📜 Recent review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a3b700a and f4e5564.

📒 Files selected for processing (1)

• nemo_rl/algorithms/grpo.py (1 hunks)

🧰 Additional context used

📓 Path-based instructions (2)

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: Follow the Google Python Style Guide for all Python code
Target Python 3.12+ for all Python code in NeMo-RL
Indent Python code with 4 spaces; do not use tabs
Python filenames should be snake_case (e.g., some_file.py)
Class names should be PascalCase
Function and method names should be snake_case
Local variable names should be snake_case; if starting with a number, prefix with k (e.g., k_99th_percentile)
Global variables should be UPPER_SNAKE_CASE and prefixed with G_ (e.g., G_MY_GLOBAL)
Constants should be UPPER_SNAKE_CASE
Avoid shadowing variables declared in an outer scope
Initialize all externally visible members of a class in the constructor
For public interfaces used outside a file, prefer docstrings over comments
Use comments mainly for code within a function or interfaces local to a file
Commented-out code must include a nearby comment explaining usage and why it is commented out; otherwise remove before merging
Use Google-style docstrings for classes and functions (Sphinx-parseable)
Avoid using reflection when functionality can be easily achieved without it
Limit except clauses to the smallest specific set of exceptions possible
For duck-typing via try/except, keep the try body minimal and use else for main logic
Add the NVIDIA copyright header (with current year) at the top of all Python files, excluding tests/ and test-only scripts

Files:

• nemo_rl/algorithms/grpo.py

nemo_rl/**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

nemo_rl/**/*.py: Do not set non-None configuration defaults in code; YAML is the single source of truth for defaults
Access required config attributes directly (e.g., policy_cfg["precision"]) and assume presence; do not introduce hidden defaults
Express configuration optionality via TypedDict using typing.NotRequired
When adding a new config key to a TypedDict subclass, document the key’s purpose, valid values/types, and recommended default in code
For any class or function decorated with @ray.remote, add '# pragma: no cover' on the class/def line (and on remote functions)

Files:

• nemo_rl/algorithms/grpo.py

🧬 Code graph analysis (1)

nemo_rl/algorithms/grpo.py (1)

tests/check_metrics.py (1)

• min (24-26)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)

• GitHub Check: Post automodel integration comment / Comment on PR
• GitHub Check: Post submodule check comment / Comment on PR

🔇 Additional comments (1)

nemo_rl/algorithms/grpo.py (1)

423-426: LGTM! Correct fix for multi-epoch training.

The calculation now properly accounts for the total number of training iterations across all epochs by multiplying max_num_epochs * len(dataloader) and capping at max_num_steps. This aligns with the training loop's dual control mechanism (line 650) that respects both epoch and step limits.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}