feat: add throughput/prompt_length/total_num_tokens metrics

[ZhiyuLi-Nvidia]

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.
As attached. Add the following metrics:

• train/mean_prompt_length: mean token length of prompt
• train/total_num_tokens: total num of tokens (including prompt and generation tokens)
• train/throughput_per_gpu: num of tokens per second (including prompt and generation tokens)

https://wandb.ai/nvidia/grpo-dev/runs/33pb40kk?nw=nwuserzhiyul

As a quick verification/test:

• train/total_num_tokens = num_prompts_per_step x num_generations_per_prompt x (train/mean_prompt_length + train/mean_total_tokens_per_sample)
  • actual: 2869598 vs calculated 2869596.16 (numerical difference)
• train/throughput = train/total_num_tokens / step time / total_num_gpus
  • actual: 1199 vs 1199

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

• ...

[ZhiyuLi-Nvidia]

@terrykong @SahilJain314 could you take a look?