feat: Add deepseek flops tracker

[guyueh1]

What does this PR do ?

Add deepseek v3 flops tracker

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

• New Features

  • FLOPS tracker now supports DeepSeek V3 model configurations, enabling accurate FLOPS estimation for models with MoE and LoRA-related settings.

• Chores

  • Enabled more detailed NCCL debug logging during initialization to improve runtime diagnostics.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds NCCL debug environment variables during MegatronPolicyWorker initialization. Extends FLOPS tracking to support DeepSeekV3Config by importing the config and formula symbol, and adding a new branch in convert_config_to_flops_config to construct a DeepSeekV3-specific FLOPSConfig.

Changes

Cohort / File(s)	Summary of Changes
NCCL debug setup in policy worker nemo_rl/models/policy/megatron_policy_worker.py	Sets environment variables NCCL_DEBUG="INFO" and NCCL_DEBUG_SUBSYS="TUNING" in MegatronPolicyWorker.init. No API/signature changes.
FLOPS tracker: DeepSeek V3 support nemo_rl/utils/flops_tracker.py	Imports DeepSeekV3Config and deepseekv3 formula. Extends convert_config_to_flops_config to handle DeepSeekV3Config, constructing FLOPSConfig with DeepSeekV3-specific fields; retains existing handling for Qwen2, Qwen3/Qwen3Moe, Llama; raises ValueError otherwise.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  actor Caller
  participant FT as flops_tracker.convert_config_to_flops_config
  participant HF as transformers.Config (Qwen2/Qwen3/Llama/DeepSeekV3)
  participant FF as flops_formulas

  Caller->>FT: convert_config_to_flops_config(config)
  FT->>HF: inspect type(config)
  alt DeepSeekV3Config (new)
    FT->>FT: build FLOPSConfig (DeepSeekV3 fields)
    FT->>FF: select formula: deepseekv3
  else Qwen2/Qwen3/Qwen3Moe/Llama
    FT->>FT: build corresponding FLOPSConfig
    FT->>FF: select respective formula
  else Unsupported
    FT-->>Caller: raise ValueError
  end
  FT-->>Caller: (FLOPSConfig, formula)

Loading

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.
Test Results For Major Changes	⚠️ Warning	The PR introduces a new DeepSeek V3 FLOPS tracking path and modifies NCCL debugging behavior, which qualifies as a new feature-level change under this check, yet the PR description contains no testing or performance evidence demonstrating correctness or regressions. Because major changes must document test results, the absence of this information means the check cannot pass.	Please update the PR description with the relevant test results or performance validation (including configurations and outcomes) showing that the new FLOPS tracker and NCCL debug initialization behave as intended without regressions.

✅ 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 succinctly describes the principal feature introduced—DeepSeek V3 FLOPS tracker support—which aligns with the primary change in this pull request. It is concise and specific enough for teammates to understand the core addition without extraneous detail. Since the minor NCCL debug environment variable updates are secondary, focusing on the flops tracker in the title is appropriate.

✨ Finishing touches

• 📝 Generate Docstrings

🧪 Generate unit tests

• Create PR with unit tests
• Post copyable unit tests in a comment

---

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

📜 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 d82ca75 and b36c803.

📒 Files selected for processing (2)

• nemo_rl/models/policy/megatron_policy_worker.py (1 hunks)
• nemo_rl/utils/flops_tracker.py (2 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/models/policy/megatron_policy_worker.py
• nemo_rl/utils/flops_tracker.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/models/policy/megatron_policy_worker.py
• nemo_rl/utils/flops_tracker.py

🧬 Code graph analysis (1)

nemo_rl/utils/flops_tracker.py (2)

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

• sliding_window_overwrite (169-195)

nemo_rl/utils/flops_formulas.py (2)

• FLOPSConfig (21-59)
• deepseekv3 (386-458)

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

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

[terrykong]

awesome. thanks @guyueh1 !

could you add a unit test here https://github.com/NVIDIA-NeMo/RL/blob/main/tests/unit/utils/test_flops_counter.py#L23-L30 ?

[guyueh1]

@terrykong added test for this

[terrykong]

closing in favor of #1305