update is_comm_kernel check to work with newer nccl versions (and thus generate expected comm/compute overlap numbers)

[lessw2020]

What does this PR do?

This PR fixes a common error hit with newer versions of Nccl.
Most traces currently break on computing the comm/compute overlap numbers b/c it does not account for the kernel name being prefaced with 'ncclDev' instead of only 'ncclKernel'.

The error is

Runtime Warning: invalid value encountered in scalar divide return (shifted_overlap["time_y"] - shifted_overlap["time_x"].sum())

The fix is simple, check for both to get the proper comm/compute overlap numbers in hta utility function "is comm kernel".

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos, doc improvements)
  • N/A
• Did you write any new necessary tests?
  • N/A
• Did you make sure to update the docs?
  • N/A
• Did you update the changelog?
  • N/A

[anupambhatnagar]

@lessw2020 can you please add a test case for this? also please resolve the pre-commit errors.

[amazloumi]

I come across this issue when working on distributed torch profiling, get_gpu_kernel_breakdown() does not show any communication breakdown. Do you have any plan to merging the fix of this PR soon?

I have two different types of nccl related items on my torch trace file (collected using torch.profiler.tensorboard_trace_handler) that none of them are captured as COMM category in TraceAnalysis .get_gpu_kernel_breakdown(). The items are as follows:

1. The items with "cat": "user_annotation" and "name" starting with "nccl:" such as nccl:_all_gather_base, nccl:_reduce_scatter_base, nccl:reduce and nccl:all_reduce
2. The items with "cat": "kernel" and "name" starting with "ncclDevKernel_" such as ncclDevKernel_AllReduce_Sum_f32_RING_LL, ncclDevKernel_Reduce_Sum_f32_RING_LL, ncclDevKernel_ReduceScatter_Sum_f32_RING_LL and ncclDevKernel_AllGather_RING_LL

I am not sure if the first item also needs to be considered as kernel communication cost while checking tensorboard I believe torch.profiler plug-in on tensorboard is considering those on item 1.