feat: upstream wandb script

[youngeunkwon0405]

What does this PR do ?

Upstreaming the wandb script for summarizing the timing metrics of a wandb logged run.

The script was originally written by @guyueh1, but I think it's very useful and better suited for upstream use.

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

• ...

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

[youngeunkwon0405]

/okay to test 3564ceb

[youngeunkwon0405]

Hi @guyueh1, I have been using your script very well :) I think upstreaming this script might be helpful for many folks on the team. I polished a script a little.

Hi @terrykong, this script provides the numbers that we track from our internal performance tracker. I think it will be helpful to upstream this script to share with @kajalj22 @humairafirdowse18 and @parthmannan
The following is a run example.
If wandb url is https://wandb.ai/nvidia/nemo-rl/runs/ap3nbd36

$> python3 tools/wandb_script.py --project nemo-rl --uid ap3nbd36
wandb: [wandb.Api()] Loaded credentials for https://api.wandb.ai from WANDB_API_KEY.
Warning: skipping keys not found in history: ['timing/train/exposed_generation', 'timing/train/weight_sync']
Average history:
train/mean_total_tokens_per_sample                                  722.644922
timing/train/total_step_time                                        103.461639
timing/train/generation                                              71.353614
timing/train/policy_training                                          6.631733
timing/train/policy_and_reference_logprobs                            4.791450
timing/train/prepare_for_generation/total                            16.812624
timing/train/prepare_for_generation/transfer_and_update_weights       7.789963
performance/tokens_per_sec_per_gpu                                   15.945105
performance/generation_tokens_per_sec_per_gpu                        23.120822
performance/training_worker_group_tokens_per_sec_per_gpu            145.360424
performance/policy_training_tokens_per_sec_per_gpu                  251.332359
performance/policy_and_reference_logprobs_tokens_per_sec_per_gpu    345.446787
performance/train_flops_per_gpu                                      57.074949
performance/train_fp_utilization                                      0.023296
timing/train/checkpointing                                                 NaN
dtype: float64
History at step:
train/mean_total_tokens_per_sample                                  740.630859
timing/train/total_step_time                                        102.204233
timing/train/generation                                              70.710803
timing/train/policy_training                                          5.798725
timing/train/policy_and_reference_logprobs                            4.582254
timing/train/prepare_for_generation/total                            17.529307
timing/train/prepare_for_generation/transfer_and_update_weights       7.737464
performance/tokens_per_sec_per_gpu                                   16.363796
performance/generation_tokens_per_sec_per_gpu                        23.651962
performance/training_worker_group_tokens_per_sec_per_gpu            161.107075
performance/policy_training_tokens_per_sec_per_gpu                  288.416701
performance/policy_and_reference_logprobs_tokens_per_sec_per_gpu    364.983959
performance/train_flops_per_gpu                                      65.454019
performance/train_fp_utilization                                      0.026716
timing/train/checkpointing                                                 NaN
Name: 5, dtype: float64

[youngeunkwon0405]

/okay to test f89962a

[youngeunkwon0405]

/okay to test b0fc86d