OOMptimizer: bucketing batch size profiles to make GPUs go 🔥
State: closed | Pull request: yes
Labels: name: ASR

name: CI

name: common

name: Run CICD

# What does this PR do ?

Major contributions:
* Canary-1B can be trained with 5x larger batch sizes compared to our earlier baseline. It maxes out GPU utilization (memory, compute, and power consumption wise). As a result the mean training step time is 2.75x longer, resulting in a training throughput of 5x / 2.75x ~= 180% of the original recipe. I managed to reproduce Canary-1B in about 40k training steps on the same number of GPUs, changing only bucketing/batch size settings using new features in this PR.
    * Update: actually reproduces Canary-1B in half of the training time with slightly improved WERs.
    * Update 2: also reproduced Canary-1B in the original training time using 4x less GPUs.
* Note: these tools are applicable to all ASR models and can easily be made applicable to any (audio|text)->(audio|text) model.
* OOMptimizer script that given a model config and bucket bins, finds the optimal batch sizes for each bucket bin. Optimal = maximum GPU utilization.
* 2D bucketing with a dedicated estimation script and dataloading support. Allows to stratify sampling by input and output sequence lengths, resulting in improved training throughput for encoder-decoder models.
* Ability to filter out examples exceeding a tokens-per-second threshold during training (e.g. some datasets have very severe outliers, 20x more tokens than median).
* Enables concurrent bucketing, speeding up the start of the training loop (in my experiments reduces the wait time from ~2min to ~20 seconds). With this setting, the bucketing buffer is filled asynchronously and lets the sampler draw batches when it's at least 10% filled, so the training can start faster.
* Documentation and examples of usage of the new features.
* Unit and integration for the new features.

**Collection**: ASR

# Changelog 
- Add specific line by line info of high level changes in this PR.

# Usage
* You can potentially add a usage example below

```python
# Add a code snippet demonstrating how to use this 
```

# GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

# Before your PR is "Ready for review"
**Pre checks**:
- [ ] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [x] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)


I realized we can make a post-processing pass on the max_batch_size list and merge buckets with identical batch sizes. Merging buckets will improve randomization. If the original num_buckets was large enough to trigger merging, that approach will lead us to an optimal number of buckets. 

"Optimal = maximum GPU utilization." in your original comment: Well, optimal here means that the maximum amount of GPU memory (the highest possible such that you won't get an OOM) is used right?

@pzelasko regarding more brrr for asr models in nemo, there is also a low-hanging fruit of using PyTorch core's SDPA in conformer attention:
- https://github.com/NVIDIA/NeMo/pull/9590

> @pzelasko regarding more brrr for asr models in nemo, there is also a low-hanging fruit of using PyTorch core's SDPA in conformer attention:
> 
> * [Use torch sdpa implementation in ASR mha #9590](https://github.com/NVIDIA/NeMo/pull/9590)

Thanks for bringing it to my attention. I'll try to review this soon, also CC @galv you might find it interesting. 

> I'll try to review this soon

Thanks a bunch :)

Also, maybe https://pytorch.org/blog/flexattention/ will accelerate in the future SDPA with attn_bias :)

Note: I will only merge this PR after #10086; there are enough changes to Canary's training workflow here that we need these new integration and unit tests to be certain that we introduce no regression.