feat: Add deepseek flops tracker (original NVIDIA-NeMo#1250) (NVIDIA-NeMo#1305)

Signed-off-by: Guyue Huang <guyueh@nvidia.com>
Signed-off-by: Guyue Huang <140554423+guyueh1@users.noreply.github.com>

Author: guyueh1

nemo_rl/utils/flops_tracker.py

@@ -25,7 +25,7 @@
 from transformers.models.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig
 
 from nemo_rl.models.policy.utils import sliding_window_overwrite
-from nemo_rl.utils.flops_formulas import FLOPSConfig, llama, qwen2, qwen3
+from nemo_rl.utils.flops_formulas import FLOPSConfig, deepseekv3, llama, qwen2, qwen3
 
 
 def get_default_hf_config(model_name: str) -> PretrainedConfig:
@@ -77,6 +77,27 @@ def convert_config_to_flops_config(
             attention_heads=config.num_attention_heads,
             vocab_size=config.vocab_size,
         ), llama
+    elif config.__class__.model_type == "deepseek_v3":
+        return FLOPSConfig(
+            gbs=0,
+            hs=config.hidden_size,
+            layers=config.num_hidden_layers,
+            ffn_hs=config.intermediate_size,
+            attention_heads=config.num_attention_heads,
+            moe_router_topk=config.num_experts_per_tok,
+            query_groups=config.num_key_value_heads,
+            vocab_size=config.vocab_size,
+            q_lora_rank=config.q_lora_rank,
+            kv_lora_rank=config.kv_lora_rank,
+            qk_head_dim=config.qk_nope_head_dim,
+            qk_pos_emb_head_dim=config.qk_rope_head_dim,
+            v_head_dim=config.v_head_dim,
+            moe_layer_freq=1,
+            moe_shared_expert_intermediate_size=config.moe_intermediate_size,
+            moe_ffn_hidden_size=config.moe_intermediate_size,
+            mtp_num_layers=0,
+            causal_self_attn=True,
+        ), deepseekv3
     else:
         raise ValueError(f"Unsupported config type: {type(config)}")
 

tests/unit/utils/test_flops_counter.py

@@ -28,6 +28,7 @@
         ("meta-llama/Llama-3.1-405B-Instruct", 128, 8192, 2.65e18),
         ("Qwen/Qwen3-30B-A3B", 128, 4096, 9.37e15),
         ("Qwen/Qwen3-235B-A22B", 128, 4096, 6.21e16),
+        ("deepseek-ai/DeepSeek-V3", 1, 4096, 1.023e15),
     ],
 )
 def test_flops_counter(model_name, gbs, seqlen, expected_flops):