cp: perf: Add num_workers in DPO, GRPO and RM (1314) into r0.4.0

[chtruong814]

beep boop [🤖]: Hi @katec846 👋,

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

Please review and approve this cherry pick by your convenience!

Summary by CodeRabbit

• New Features

  • Added a configurable data loading parallelism option (num_workers) across training configurations (e.g., DPO, GRPO, RM, SFT, VLM variants).
  • Training and validation data loaders now honor num_workers for improved throughput; example configs default to 1 but can be increased as needed.

• Tests

  • Updated unit tests to validate handling of the new num_workers configuration.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds a new data.num_workers configuration to multiple example YAMLs and wires this value into StatefulDataLoader initializations in DPO, GRPO, and RM algorithms. Updates a unit test to include the new field in the data config.

Changes

Cohort / File(s)	Change Summary
Configs: add data.num_workers examples/configs/dpo.yaml, examples/configs/grpo_math_1B.yaml, examples/configs/rm.yaml, examples/configs/sft_openmathinstruct2_megatron.yaml, examples/configs/vlm_grpo_3B.yaml, examples/configs/vlm_grpo_3B_megatron.yaml	Introduced data.num_workers: 1 in each config’s data section; no other keys changed.
Algorithms: pass num_workers to dataloaders nemo_rl/algorithms/dpo.py, nemo_rl/algorithms/grpo.py, nemo_rl/algorithms/rm.py	Added num_workers=data_config["num_workers"] to train/val StatefulDataLoader instantiations; signatures unchanged.
Tests: include num_workers in data config tests/unit/algorithms/test_grpo.py	Updated test data configs to include "num_workers": 1 alongside existing fields.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  participant U as Runner
  participant CFG as YAML Config
  participant ALG as Algorithm Setup (DPO/GRPO/RM)
  participant DLtrain as StatefulDataLoader (train)
  participant DLval as StatefulDataLoader (val)

  U->>CFG: Load config
  CFG-->>ALG: data{..., num_workers}
  ALG->>DLtrain: init(..., num_workers=data.num_workers)
  ALG->>DLval: init(..., num_workers=data.num_workers)
  note over DLtrain,DLval: Data loading uses specified worker count
  U->>ALG: start training/validation
  ALG->>DLtrain: iterate batches
  ALG->>DLval: iterate batches

Loading

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• perf: Add a field in SFT data config to modify num_workers for loading data #1143: Adds data.num_workers in SFT config and wires into StatefulDataLoader, similar config-to-loader wiring.
• perf: Add num_workers in DPO, GRPO and SFT for loading data #1314: Mirrors this PR’s pattern adding data.num_workers to configs and algorithms (DPO/GRPO/RM) with test updates.

Suggested labels

CI:L1, r0.4.0

Suggested reviewers

• terrykong

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Title Check	⚠️ Warning	The pull request title includes cherry-pick metadata, backticks, the original PR number, and the target branch name, which adds noise and detracts from a clear description of the actual change. While it mentions adding num_workers to DPO, GRPO, and RM, it omits other impacted modules and uses internal notation rather than a concise summary of the functional update. Therefore the title does not meet guidelines for clarity and specificity.	Rename the title to a concise summary of the change without cherry-pick metadata or branch details, for example: “perf: Add num_workers to data loaders across DPO, GRPO, RM, SFT, and VLM configurations.” This removes noise and accurately describes the functional update.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Test Results For Major Changes	✅ Passed	The changes in this PR are limited to adding a configurable num_workers parameter in example YAMLs and wiring it through existing data loaders without altering numeric algorithms or model convergence, and the default behavior remains sequential loading; this is a minor feature enhancement that does not require additional test results or performance benchmarks.

✨ 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-1314-r0.4.0

---

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.}

[coderabbitai]

Actionable comments posted: 6

📜 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 8ce81a0 and 1cd6060.

📒 Files selected for processing (10)

• examples/configs/dpo.yaml (1 hunks)
• examples/configs/grpo_math_1B.yaml (1 hunks)
• examples/configs/rm.yaml (1 hunks)
• examples/configs/sft_openmathinstruct2_megatron.yaml (1 hunks)
• examples/configs/vlm_grpo_3B.yaml (1 hunks)
• examples/configs/vlm_grpo_3B_megatron.yaml (1 hunks)
• nemo_rl/algorithms/dpo.py (2 hunks)
• nemo_rl/algorithms/grpo.py (2 hunks)
• nemo_rl/algorithms/rm.py (2 hunks)
• tests/unit/algorithms/test_grpo.py (2 hunks)

🧰 Additional context used

📓 Path-based instructions (3)

**/*.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/dpo.py
• nemo_rl/algorithms/rm.py
• nemo_rl/algorithms/grpo.py
• tests/unit/algorithms/test_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/dpo.py
• nemo_rl/algorithms/rm.py
• nemo_rl/algorithms/grpo.py

examples/configs/*.yaml

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

examples/configs/*.yaml: Exemplar configs under examples/configs/.yaml must include documented defaults
When adding a new config key, reflect its recommended default in exemplar YAMLs under examples/configs/.yaml

Files:

• examples/configs/vlm_grpo_3B.yaml
• examples/configs/vlm_grpo_3B_megatron.yaml
• examples/configs/dpo.yaml
• examples/configs/sft_openmathinstruct2_megatron.yaml
• examples/configs/rm.yaml
• examples/configs/grpo_math_1B.yaml

⏰ 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). (5)

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

🔇 Additional comments (5)

tests/unit/algorithms/test_grpo.py (1)

243-243: LGTM!

The test configurations correctly include the new num_workers field, ensuring test coverage for the updated data configuration schema. This prevents test failures due to missing required config attributes.

Also applies to: 299-299

nemo_rl/algorithms/grpo.py (1)

210-210: LGTM!

The implementation correctly wires num_workers from the data config to both training and validation dataloaders, following the guideline to access required config attributes directly without introducing hidden defaults. The changes are consistent with the similar updates in nemo_rl/algorithms/rm.py and nemo_rl/algorithms/dpo.py.

As per coding guidelines: "Access required config attributes directly and assume presence; do not introduce hidden defaults"

Also applies to: 232-232

nemo_rl/algorithms/dpo.py (2)

202-202: LGTM!

Consistent implementation for validation DataLoader, correctly using the same configuration parameter.

---

179-179: LGTM – DataConfig defines and documents num_workers
DataConfig includes num_workers: NotRequired[int] with accompanying comments per coding guidelines.

examples/configs/vlm_grpo_3B_megatron.yaml (1)

159-159: Verified num_workers: 1 in all example configs The six targeted YAMLs (dpo.yaml, grpo_math_1B.yaml, rm.yaml, sft_openmathinstruct2_megatron.yaml, vlm_grpo_3B.yaml, vlm_grpo_3B_megatron.yaml) all include the default.