Add MemoryProfileCallback
State: closed | Pull request: yes
Labels: name: Run CICD

# What does this PR do ?

This callback enables recording a timeline of memory allocations during training. 

**Collection**: [Note which collection this PR will affect]

# 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**:
- [ ] 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)


Currently, this will dump the snapshot for all ranks. Should we add an option to only save a snapshot from rank 0? or a specifiable rank?

> Currently, this will dump the snapshot for all ranks. Should we add an option to only save a snapshot from rank 0? or a specifiable rank?

Addressed in latest commit


also, can we allow disabling the profiling, so users can only run the cycle detector?

Yeah I was just being cautious to avoid side effects of the import but I don't think this import should have any, so it's probably fine.

Get Outlook for iOS<https://aka.ms/o0ukef>
________________________________
From: Shriya Rishab ***@***.***>
Sent: Friday, August 23, 2024 10:07:29 AM
To: NVIDIA/NeMo ***@***.***>
Cc: Hemil Desai ***@***.***>; Comment ***@***.***>
Subject: Re: [NVIDIA/NeMo] Add MemoryProfileCallback (PR #10166)


@ShriyaPalsamudram commented on this pull request.

________________________________

In nemo/lightning/pytorch/callbacks/memory_profiler.py<https://github.com/NVIDIA/NeMo/pull/10166#discussion_r1729126364>:

> @@ -0,0 +1,78 @@
+import os
+
+import torch
+from pytorch_lightning.callbacks.callback import Callback
+from torch.utils.viz._cycles import warn_tensor_cycles


Isn't it better to fail early if there is an issue with this import than fail inside the callback?

—
Reply to this email directly, view it on GitHub<https://github.com/NVIDIA/NeMo/pull/10166#discussion_r1729126364>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/AB6Q25H6ZIEEJZOY5WCDAWDZS5F3DAVCNFSM6AAAAABMS6X6TCVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZDENJXGQ3TONZVGY>.
You are receiving this because you commented.Message ID: ***@***.***>