-
Notifications
You must be signed in to change notification settings - Fork 3.8k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add some default NVTX ranges #7633
base: master
Are you sure you want to change the base?
Conversation
Codecov Report
@@ Coverage Diff @@
## master #7633 +/- ##
==========================================
- Coverage 90.44% 89.75% -0.69%
==========================================
Files 455 456 +1
Lines 26334 26406 +72
==========================================
- Hits 23818 23702 -116
- Misses 2516 2704 +188
... and 33 files with indirect coverage changes 📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more |
Signed-off-by: Serge Panev <[email protected]>
for more information, see https://pre-commit.ci
Signed-off-by: Serge Panev <[email protected]>
Signed-off-by: Serge Panev <[email protected]>
Signed-off-by: Serge Panev <[email protected]>
LGTM! |
"0") == "1" and torch.cuda.is_available(): | ||
self._nvtx_handles = dict() | ||
|
||
def get_hooks_for(func_name): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think this is not an issue if NVIDIA_NVTX_RANGES
is not meant to be used by the community, but I prefer to define the hook at the top-level to avoid PicklingError
in a multiprocessing setting. https://docs.python.org/3/library/pickle.html#what-can-be-pickled-and-unpickled
This PR leverages torch's NVTX API to add default nvtx ranges to help profiling with nsys. Those can be enabled with the
NVIDIA_NVTX_RANGES
environment variable.How to use:
Result
It adds ranges for the following functions:
Dataloaders
'scollate_fn
Dataset
's__getitem__
MessagePassing
's main member functionsNeighborSampler
's