[Fix] Hanging for Fully Randomized Bucketing
State: closed | Pull request: yes

Signed-off-by: stevehuang52 <heh@nvidia.com>

# What does this PR do ?

This PR fixes the problem of the program hanging during training with fully randomized bucketing.

# Why hanging?

When training RNNT models with fully randomized bucketing, and enabling fused batch, the program will hang when calculating training WER: [https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/asr/modules/rnnt.py#L931](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/asr/modules/rnnt.py#L931)

The bug happens when the GPU with largest batch size is still calculating the forward fused-batch while the other GPUs with smaller batch are already finished:
[https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/asr/modules/rnnt.py#L853-L933](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/asr/modules/rnnt.py#L853-L933)

This is due to the fact that, when `compute()` is called, PTL will perform synchronization across all devices, which eventually calls `distributed.barrier()`:
[https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/metric.py#L113](https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/metric.py#L113), [https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/metric.py#L485-L513](https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/metric.py#L485-L513), [https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/utilities/distributed.py#L122
](https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/utilities/distributed.py#L122
)

In fully randomized bucketing, different GPUs tend to have different batch sizes, and when some GPUs have finished their batch  earlier than others, the unfinished ones will keep hanging at `distributed.barrier()` to wait for the synchronization that will never happen.

# How we fix it?

The simple fix is to move the `compute()` function call outside of the loop of sub-batches. Note that, theoretically this modification have the same output as the original one, and here's the logic.

Suppose that the there are 3 GPUs and each has 2 sub-batches, where the score produced by each GPU is fixed as its rank:
|                   | GPU 1 | GPU 2 | GPU 3 |
|-------------|--------|--------|---------|
|sub-batch 1|    1      |   2      |    3       |
|sub-batch 2|    1      |     2    |     3      |


In the original code, where `compute()` is called inside the for loop of sub-batches, it's actually syncing across all GPUs in each step, then aggregate across all sub-batches. In other words, it first perform all_reduce() to sync each column for each row, then sum over the rows:

|                   | GPU 1 | GPU 2 | GPU 3 |           
|-------------|--------|--------|---------|   
|sub-batch 1|    1      |   2      |    3       | 

Call `compute()` on each GPU and all_reduce &#8595;

|                   | GPU 1 | GPU 2 | GPU 3 |           
|-------------|--------|--------|---------|   
|sub-batch 1|    6      |   6      |    6      | 

 Run sub-batch 2 &#8595;

|                   | GPU 1 | GPU 2 | GPU 3 |           
|-------------|--------|--------|---------|   
|sub-batch 1|    6      |   6      |    6      | 
|sub-batch 2|    1      |   2      |    3       | 

Call `compute()` on each GPU and all_reduce &#8595;
|                   | GPU 1 | GPU 2 | GPU 3 |           
|-------------|--------|--------|---------|   
|sub-batch 1|    6      |   6      |    6      | 
|sub-batch 2|    6      |   6      |    6       | 

Sum for each GPU &#8595;
|                   | GPU 1 | GPU 2 | GPU 3 |           
|-------------|--------|--------|---------|   
|batch total |   12      |   12      |    12      | 

On the other hand, this PR first aggregate the scores per GPU, then sync across all GPUs after the for loop of sub-batches. In other words, we sum the rows for each column, and then perform all_reduce to  sync across all rows:
|                   | GPU 1 | GPU 2 | GPU 3 |           
|-------------|--------|--------|---------|   
|sub-batch 1|    1      |   2      |    3       | 

update on each GPU &#8595; 

|                   | GPU 1 | GPU 2 | GPU 3 |           
|-------------|--------|--------|---------|   
|batch total|    1      |   2     |    3      | 

run sub-batch 2 &#8595;

|                   | GPU 1 | GPU 2 | GPU 3 |           
|-------------|--------|--------|---------|   
|sub-batch 1|    1     |   2      |    3      | 
|sub-batch 2|    1      |   2      |    3       | 

update on each GPU &#8595;
|                   | GPU 1 | GPU 2 | GPU 3 |           
|-------------|--------|--------|---------|   
batch total |    2      |   4      |    6      | 

Call `compute()` and and all_reduce &#8595;
|                   | GPU 1 | GPU 2 | GPU 3 |           
|-------------|--------|--------|---------|   
|batch total |   12      |   12      |    12      | 


We can see that, two approaches actually have the same results, and it doesn't matter whether we perform the synchronization in or after the for loop of sub-batches. By moving `compute()` outside the for loop, we can guarantee that there is no hanging even when GPUs have different batch sizes.