fix: moonlight flop calculation and DAPO config fix

[terrykong]

What does this PR do ?

The moonlight fix was on main, but the DAPO one was new (and needs to be fixed on main as well. CC @ashors1 )

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

• Chores

  • Updated model checkpointing configuration to use default save format instead of explicitly specifying DCP format.

• Bug Fixes

  • Improved parameter calculation accuracy for DeepSeek V3 models by making Q projection computation conditional based on model configuration, ensuring correct FLOPS estimation.

[coderabbitai]

📝 Walkthrough

Walkthrough

Configuration file updated to remove explicit model save format specification. FLOPS calculation formula for deepseekv3 Q projection now conditionally branches based on q_lora_rank presence, using MLA-based parameters when specified, standard attention parameters otherwise.

Changes

Cohort / File(s)	Summary
Configuration update examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml	Changed model_save_format from "dcp" to null, removing explicit save format specification in checkpointing configuration
FLOPS formula logic nemo_rl/utils/flops_formulas.py	Added conditional branching in deepseekv3 Q projection parameter calculation: uses MLA-based parameters when config.q_lora_rank is set, falls back to standard attention parameters otherwise

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

• Configuration change is straightforward (no logic verification needed)
• FLOPS formula logic change requires verification that conditional branching correctly handles both MLA and standard attention parameter computation paths
• Ensure the parameter accumulation formulas are accurate for both branches in deepseekv3

Suggested labels

r0.4.0

Suggested reviewers

• yfw
• chtruong814

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	Pull request modifies FLOP calculation logic without providing test results or validation documentation.	Provide test results validating FLOP calculations and confirming no regressions in the deepseekv3 function changes.

✅ Passed checks (2 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 two main changes: a moonlight flop calculation fix and a DAPO config fix, matching the file-level summaries provided.

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

---

📜 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 31722fe and 054ae7a.

📒 Files selected for processing (2)

• examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml (1 hunks)
• nemo_rl/utils/flops_formulas.py (1 hunks)

🧰 Additional context used

📓 Path-based instructions (6)

examples/configs/recipes/**/*.yaml

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

examples/configs/recipes/**/*.yaml: Recipe YAMLs under examples/configs/recipes/** are runnable snapshots and may omit documentation
When adding support for a new model, add a recipe YAML under examples/configs/recipes/ in the appropriate domain (llm/ or vlm/) with the correct name

Files:

• examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml

examples/configs/recipes/llm/*.yaml

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

LLM recipe YAML filenames must follow: --ng-[-modifiers][-long][.vN].yaml

Files:

• examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml

examples/configs/recipes/**/*.{yaml,sh}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Known exception: Deepscaler recipes may encode context length in place of the cluster tuple (e.g., grpo-deepscaler-1.5b-8K.*); allowed but document intended hardware in the script

Files:

• examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml

examples/configs/recipes/**

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Place recipe YAMLs under examples/configs/recipes//

Files:

• examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml

**/*.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/utils/flops_formulas.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/utils/flops_formulas.py

⏰ 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: build-container / main
• GitHub Check: sphinx-build / Build docs
• GitHub Check: Lint check
• GitHub Check: Post submodule check comment / Comment on PR
• GitHub Check: Post automodel integration comment / Comment on PR

🔇 Additional comments (2)

nemo_rl/utils/flops_formulas.py (1)

405-415: LGTM! Correct conditional logic for MLA vs. standard attention.

The implementation correctly differentiates between:

• MLA (DeepSeek-V3): Two-stage projection through q_lora_rank when the parameter is set
• Standard attention (Moonlight): Direct projection when q_lora_rank is None

The math for both paths is accurate, and the comments clearly explain each case.

examples/configs/recipes/llm/dapo-qwen2.5-7b.yaml (1)

32-32: Now I need to check the actual YAML file to understand the full context and what policy worker is being used:

Based on my verification, I have found comprehensive evidence in the codebase that explains the change. Here are my findings:

The system defaults model_save_format to "safetensors" when not provided. However, there's a critical constraint for v1 configurations: when using DTensorPolicyWorker without _v2=True, model_save_format must be None or omitted, and raises an error if set to any other value.

Additionally, the only valid formats supported are "torch_save" and "safetensors", which means "dcp" would be invalid anyway.

Setting model_save_format to null is the correct fix because:

1. It satisfies the v1 DTensorPolicyWorker requirement (None/null is allowed and expected)
2. It removes an unsupported format string
3. The system will apply appropriate defaults based on the configuration context

---

---

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