feat: Add deepseek flops tracker (original #1250)

[guyueh1]

What does this PR do ?

This is a replica of #1250 in the main repo for CI testing.

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

  • Added FLOPS estimation support for the DeepSeek V3 model family, enabling accurate compute metrics alongside existing models.
  • Improved configuration handling to automatically recognize DeepSeek V3 parameters.

• Tests

  • Added a unit test case validating FLOPS calculations for a DeepSeek V3 checkpoint with representative settings.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds DeepSeek-V3 support to FLOPS config conversion by introducing a new model_type branch in nemo_rl/utils/flops_tracker.py and importing the deepseekv3 FLOPS formula. Updates unit tests to include a DeepSeek-V3 test case with expected FLOPS.

Changes

Cohort / File(s)	Summary
FLOPS tracker: DeepSeek-V3 support nemo_rl/utils/flops_tracker.py	Adds import for deepseekv3 formula and a new branch for model_type == "deepseek_v3" that builds a FLOPSConfig with DeepSeek-V3-specific fields and returns the deepseekv3 formula.
Unit tests: DeepSeek-V3 case tests/unit/utils/test_flops_counter.py	Adds a test entry for deepseek-ai/DeepSeek-V3 with gbs=1, seqlen=4096, and expected_flops=1.023e15.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  actor User
  participant Trainer as Trainer/Caller
  participant Tracker as FLOPSTracker
  participant Conv as convert_config_to_flops_config
  participant Formula as FLOPS Formula

  User->>Trainer: Configure model (DeepSeek-V3) and run
  Trainer->>Tracker: init/start tracking
  Tracker->>Conv: convert(config)
  alt model_type == "deepseek_v3"
    Conv->>Conv: Build FLOPSConfig (DeepSeek-V3 fields)
    Conv-->>Tracker: (FLOPSConfig, deepseekv3)
    Tracker->>Formula: compute(FLOPSConfig) [deepseekv3]
  else other model types
    Conv-->>Tracker: (FLOPSConfig, corresponding formula)
    Tracker->>Formula: compute(FLOPSConfig)
  end
  Formula-->>Tracker: FLOPS value
  Tracker-->>Trainer: Report FLOPS

Loading

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• fix: report the correct number of workers during FLOPs calculation #1034 — Refactors convert_config_to_flops_config and related imports in flops_tracker.py, overlapping directly with the new DeepSeek-V3 branch addition.

Suggested reviewers

• parthchadha
• terrykong
• wangshangsam

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% 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 conversion path, which constitutes a functional feature addition rather than a minor change, yet the PR description contains only placeholder text and does not document any tests, numerical validation, or performance measurements, so the requirement for reporting test results on major changes is unmet.	Please update the PR description to include the relevant test results or performance/numerical validation details demonstrating that the newly added DeepSeek V3 FLOPS tracker functionality has been verified.

✅ 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 clearly summarizes the primary change—adding deepseek FLOPS tracking—and accurately reflects the implementation in this PR without introducing unrelated details.

✨ 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 guyueh/deepseek_flops_tracker

---

📜 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 806e285 and 6132d60.

📒 Files selected for processing (2)

• nemo_rl/utils/flops_tracker.py (2 hunks)
• tests/unit/utils/test_flops_counter.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/utils/flops_tracker.py
• tests/unit/utils/test_flops_counter.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_tracker.py

🧬 Code graph analysis (1)

nemo_rl/utils/flops_tracker.py (1)

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

🔇 Additional comments (3)

nemo_rl/utils/flops_tracker.py (2)

28-28: LGTM!

The import of deepseekv3 formula is correctly added alongside existing formula imports.

---

80-100: Consistent type checking and config-driven defaults in DeepSeek-V3 branch

• In nemo_rl/utils/flops_tracker.py (lines 80–100), replace
  config.__class__.model_type == "deepseek_v3"
  with an isinstance(config, DeepSeekV3Config) check (import from transformers.models.deepseek_v3.configuration_deepseek_v3, falling back to model_type only if the class isn’t available).
• Avoid hard-coding moe_layer_freq=1, mtp_num_layers=0, and causal_self_attn=True; derive these values from the config or document why those defaults are always correct.

tests/unit/utils/test_flops_counter.py (1)

31-31: Confirm DeepSeek-V3 test parameters and FLOPS value.

• Verify that using gbs=1 (vs. gbs=128 in other tests) is intentional.
• Manually confirm that 1.023e15 FLOPS is correct for DeepSeek-V3 with gbs=1, seqlen=4096 in an environment with torch installed.

---

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