feat: track policy training compute throughput

[ybgao-nvidia]

What does this PR do ?

Allows tracking of FLOPS (floating point operations per second) for policy training, which can be used to estimate MFU.

This PR adds:

1. The FLOPCounter class: maintains a sum of FLOPs incurred through model training at a batch granularity. After processing each batch, the trainer passes in a tuple (batch size, sequence length) into the counter, which computes and accumulates the total FLOPs.
2. Invocation of the tracker in the DTensor worker.
3. Additions in the GRPO, DPO and SFT training scripts to track the time consumed and compute the compute throughput (FLOPs per second) for each training iteration.
4. Training FLOPs calculation for Qwen2 (previously only Qwen3 is included)
5. MFU metric logging in console and wandb (as train/train_fp_utilization if enabled)

Limitations:

1. We have currently only implemented support for DTensor worker.
2. The conversion from HF's PretrainedConfig to FLOPSConfig only supports Qwen2Config and LlamaConfig

Issues

This PR resolves #59.

Usage

When training with a supported model, we will see the FLOPS in addition to the training time upon completion of each epoch as follows:

📊 Training Results:
  • Loss: 0.0121
  • Avg Reward: 0.3809
  • Mean Generation Length: 285.2734
  • Training FLOPS: 774.91 TFLOPS (96.86 TFLOPS per rank)
  • Training Model Floating Point Utilization: 9.79%

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

[jiemingz]

could you also make sure your changes are covered by tests?

RL/tests/unit/models/generation/test_vllm_generation.py

Line 345 in a607b41

def test_vllm_policy_generation(policy, test_input_data, tokenizer):