Performance issues for SFT with megatron backend

[Kipok]

Describe the bug

The speed of SFT jobs with megatron backend is much slower than expected.

Two issues we've identified so far:

1. The step itself is much slower than corresponding configuration in nemo-aligner. Seems some problem with dp comms
2. The GPUs are idle between steps. This takes from 25% to 50% of total job time with no GPU activity.

Steps/Code to reproduce bug

We see this in multiple SFT jobs, but for nsys above the configuration is like this

1.5b model with tp=1, cp=16, 2 nodes, 128k context length, sequence packing enabled with megatron backend.

We also see this with 16 nodes and 48k context length.

Can provide full repro script internally.

[guyueh1]

We have a couple features/bugfixes in progress to solve this issue

• solve exposed DP time: remove empty_cache in the training loop (PR ETA 9/5); this reduce 32node training time by ~25%
• solve exposed synchronization at the beginning due to Ray overhead: WIP no ETA yet, will reduce 32 node training time by 13%

[guyueh1]

What is remained for this ticket? is it the qwen-14B performance check?
If the original ask (DP time) is resolved, let's close this issue; we can start a new on on 14B and ask kate to triage

@wedu-nvidia / @Kipok

[wedu-nvidia]

@guyueh1 I am not sure if it is fully resolved or not?
For this

solve exposed DP time: remove empty_cache in the training loop (PR ETA 9/5); this reduce 32node training time by ~25%
I used this pr, but I only got 10% time reduced.

[guyueh1]

@wedu-nvidia This comes with this PR 3a1ca3f
to use it you have to set policy.megatron_cfg.empty_unused_memory_level=0 in your yaml or command line

[wedu-nvidia]

@wedu-nvidia This comes with this PR 3a1ca3f to use it you have to set policy.megatron_cfg.empty_unused_memory_level=0 in your yaml or command line

Sorry, I forgot to add one sentence in my previous reply. I used the pr to empty cache and got 10% speed up, not 25%.