PyTorch CUDA allocator optimization for dynamic batch shape dataloading in ASR
State: closed | Pull request: yes
Labels: name: common

name: Run CICD

# What does this PR do ?

I was profiling a particularly unlucky run that had dynamic batch shapes and operated close to max GPU RAM. The profile revealed it was re-allocating memory for every mini-batch, generating about 30% overhead in training. This can be resolved gracefully by turning on `expandable_segments` option in PyTorch CUDA allocator which instead of reallocating, extends existing allocations as needed, removing this significant overhead. 

In this PR I'm proposing to automatically set this option during dataloader instantiation. It can be disabled via configuration.

For documentation purposes, the profile before the change (red blocks in CUDA API timelines indicates malloc/free):
<img width="2140" alt="image" src="https://github.com/NVIDIA/NeMo/assets/15930688/783a4b0d-4f48-4958-b3f1-bc05c839abbe">

and the profile after the fix:

<img width="2140" alt="image" src="https://github.com/NVIDIA/NeMo/assets/15930688/5730fb24-605e-4723-926a-e8e8cf2ff1c0">

The blue bars at the top of the profiles in CUDA HW kernel utilization timelines are more condensed in the new profile, indicating improved GPU utilization.

**Collection**: ASR

# Changelog 
- PyTorch CUDA allocator optimization for dynamic batch shape dataloading in ASR

# Usage
* You can potentially add a usage example below

```python
# Add a code snippet demonstrating how to use this 
```

# Jenkins CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

There's no need to comment `jenkins` on the PR to trigger Jenkins CI.
The GitHub Actions CI will run automatically when the PR is opened.
To run CI on an untrusted fork, a NeMo user with write access must 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)


> It would be appreciated if you could take a look into whether there is a more robust way to check if this option is already set or not than by checking the environment variable.

I grepped through pytorch 2.3 code. Unfortunately it looks like only `_set_allocator_settings` is exposed to python API, and there is no `_get` equivalent.

> Also, it's not clear to me how this setting might interact with a non-default allocator (for example, the RMM torch allocator provided by nvidia). Presumably setting this config has no effect in this case.

Good point, I didn't know about RMM. It turns out it's available in our containers so I just tested it out on a 1-GPU training run. It seems to "just work" (although I had to decrease the batch size to avoid CUDA OOM after a ~100 steps) so I assume these options are discarded for custom allocators.