feat(flops): add NemotronH MoE-aware FLOPs tracking for MFU

[dafu-wu]

What does this PR do ?

Add MoE-aware FLOPs tracking for NemotronH models (e.g. Nemotron-3 Nano V3) so that training runs report Training FLOPS and Model Floating Point Utilization (MFU).

NemotronH is a hybrid Mamba/attention/MLP model whose layer layout is given by hybrid_override_pattern, which for the MoE variants is dominated by E (MoE FFN) layers. The FLOPs tracker previously raised ValueError for NemotronHConfig, so total_flops/theoretical_tflops were never populated and no FLOPS/MFU line was logged.

This wires up a nemotron_h branch in convert_config_to_flops_config and extends _hybrid_model_flops to account for E layers via a new _moe_layer_flops helper (routed topk + shared experts, non-gated ReLU^2). Without the E handling the dominant FFN cost would be omitted and MFU would be severely under-reported.

Builds on #1971 (which only added the wiring, not MoE-layer FLOPs).

• flops_formulas: add gated_linear_unit/moe_shared_expert_num fields, _moe_layer_flops, and E handling in _hybrid_model_flops
• flops_tracker: import nemotronh and add nemotron_h config branch
• tests: cover MoE-layer counting and shared-expert/topk scaling

Issues

n/a

Usage

from nemo_rl.utils.flops_tracker import FLOPTracker, get_default_hf_config

# Nemotron-3 Nano V3 (MoE hybrid) now returns a valid tracker instead of raising.
tracker = FLOPTracker.from_config(model_name, get_default_hf_config(model_name))
tracker.track_batch(input_lengths)          # per-sequence packed tokens
print(tracker.total_flops)                   # > 0, MoE 'E' layers included

With this, an SFT/GRPO run on a NemotronH model prints:

  • Training FLOPS: <x> TFLOPS (<x/ranks> per rank)
  • Training Model Floating Point Utilization: <MFU>%

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

• MoE FLOPs use moe_ffn_hidden_size * (moe_router_topk + moe_shared_expert_num) as the effective FFN width. NemotronH uses a non-gated ReLU^2 activation, so gated_linear_unit=False (factor 1); gated (SwiGLU-style) models would use factor 2.
• Validated against a first-principles 6 * N_active * tokens estimate (formula/rule-of-thumb ratio ≈ 0.76, expected since the rule-of-thumb ignores attention/mamba sublinear terms) and confirmed each E in the pattern increases total FLOPs by exactly one _moe_layer_flops.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.