Skip to content
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

Open
wants to merge 7 commits into
base: master
Choose a base branch
from
Open

Conversation

Kh4L
Copy link
Contributor

@Kh4L Kh4L commented Jun 23, 2023

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:

NVIDIA_NVTX_RANGES=1 nsys profile -s none -t cuda,nvtx,osrt,cudnn,cublas --force-overwrite=true python train.py 

Result

nsys_profile_example_pyg

It adds ranges for the following functions:

  • Dataloaders's collate_fn
  • Dataset's __getitem__
  • MessagePassing's main member functions
  • NeighborSampler's

@codecov
Copy link

codecov bot commented Jun 23, 2023

Codecov Report

Merging #7633 (b01f5ef) into master (0c0d9c1) will decrease coverage by 0.69%.
Report is 14 commits behind head on master.
The diff coverage is 50.00%.

@@            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     
Files Changed Coverage Δ
torch_geometric/nn/conv/message_passing.py 94.68% <ø> (ø)
torch_geometric/loader/__init__.py 61.90% <20.00%> (-38.10%) ⬇️
torch_geometric/data/dataset.py 94.70% <100.00%> (+0.02%) ⬆️
torch_geometric/nn/conv/pna_conv.py 100.00% <100.00%> (ø)
torch_geometric/sampler/neighbor_sampler.py 93.47% <100.00%> (+0.02%) ⬆️

... and 33 files with indirect coverage changes

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

@Kh4L Kh4L force-pushed the default_nvtx_pyg branch from b39b36b to da89da4 Compare June 28, 2023 10:52
@Kh4L Kh4L force-pushed the default_nvtx_pyg branch from 84c7a6d to 1827f62 Compare July 13, 2023 09:43
@Kh4L Kh4L marked this pull request as ready for review July 13, 2023 09:49
@Kh4L Kh4L requested review from mananshah99, a team and EdisonLeeeee as code owners July 13, 2023 09:49
Kh4L and others added 5 commits July 27, 2023 19:39
Signed-off-by: Serge Panev <[email protected]>
Signed-off-by: Serge Panev <[email protected]>
Signed-off-by: Serge Panev <[email protected]>
Signed-off-by: Serge Panev <[email protected]>
@Kh4L Kh4L force-pushed the default_nvtx_pyg branch from 1827f62 to 0061570 Compare July 27, 2023 10:43
@puririshi98
Copy link
Contributor

LGTM!

"0") == "1" and torch.cuda.is_available():
self._nvtx_handles = dict()

def get_hooks_for(func_name):
Copy link
Member

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants