feat: Timer for the data sharding and job submission

[guyueh1]

What does this PR do ?

Wrap data sharding and job submission in timer to provide richer timing debugging info.

For example, for the newly added qwen30b 40K seqlen GRPO example, the timing info looks like this

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

Summary by CodeRabbit

• Refactor
  • Added optional timing instrumentation to training operations across multiple training algorithms
  • Extended policy method signatures to support performance metrics tracking during policy training and inference
  • Timing context is conditionally applied without affecting existing functionality

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

📝 Walkthrough

Walkthrough

This pull request adds optional timing instrumentation to policy training and inference methods across multiple algorithms by introducing timer and timer_tag_prefix parameters to method signatures and wrapping key operations with timer context managers.

Changes

Cohort / File(s)	Summary
Algorithm Updates nemo_rl/algorithms/distillation.py, nemo_rl/algorithms/dpo.py, nemo_rl/algorithms/grpo.py, nemo_rl/algorithms/sft.py	Added timer and timer_tag_prefix parameters to policy method calls (get_logprobs, get_reference_policy_logprobs, get_topk_logits, train) to enable granular timing measurements for inference and training operations. Changes are consistent across files with no modifications to control flow or error handling.
Interface Definitions nemo_rl/models/policy/interfaces.py	Extended PolicyInterface method signatures (get_logprobs, get_reference_policy_logprobs, train) to accept **kwargs for forwarding arbitrary keyword arguments including timer parameters.
Policy Implementation nemo_rl/models/policy/lm_policy.py	Implemented timer support in Policy class methods by adding optional Timer parameters, conditionally wrapping data sharding and futures submission operations in timer context managers using nullcontext guards. Added Timer and nullcontext imports.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	PR makes significant changes to core modules across multiple algorithms without documenting test results or validation data to demonstrate that timing instrumentation changes do not introduce regressions or affect algorithm behavior.	Add test results or testing information to PR description demonstrating existing unit tests pass, no performance regression occurs, convergence behavior remains unchanged, and new timing instrumentation works correctly.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main change: adding timer instrumentation to data sharding and job submission operations throughout the codebase.
Docstring Coverage	✅ Passed	Docstring coverage is 85.71% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

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: 0

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

nemo_rl/models/policy/interfaces.py (1)

54-112: Update all PolicyInterface implementers to accept **kwargs in get_logprobs, get_reference_policy_logprobs, and train methods.

The interface now accepts arbitrary keyword arguments in these three methods, but no concrete implementation does. This will cause TypeError when these methods are called through the interface with any unexpected keyword arguments.

Implementations to update:

• Policy (lm_policy.py): Currently accepts explicit timer and timer_tag_prefix parameters; add **kwargs or refactor to use kwargs
• MegatronPolicyWorker, DTensorPolicyWorker, DTensorPolicyWorkerV2: Add **kwargs to get_logprobs and train
• BasePolicyWorker: Add **kwargs to get_reference_policy_logprobs

[guyueh1]

@youngeunkwon0405 please review

[youngeunkwon0405]

LGTM