[REQUEST] Put NVTX instrumentation in its own domain (and potentially categories)

Is your feature request related to a problem? Please describe.
When analyzing profiles and traces of DL applications it is very useful to use custom instrumentation, in my case with NVTX and Nsight Systems. But DeepSpeed happens to have its own instrumentation which is good, but because it does not have a domain, it gets assigned to NVTX's default domain. This makes it get mixed with other libraries' and developer's own instrumentation, and sometimes make it the traces increase the size because of the large amount of small regions. Overall, it limits the general optionality and flexibility given to the performance analyst.

Describe the solution you'd like
I propose that all the NVTX regions added by DeepSpeed fall inside its own NVTX domain. Therefore, they will appear separately on Nsight Systems, and can be easily filtered, classified and processed by other performance tools that obtain this data.

The proposed solution has two basic implementation actions, although one can be considered a downside. The current API of NVTX exposed by pytorch does not mirror exactly the full features of the NVTX library offered by NVIDIA. Precisely, the current available interface only accepts a message. Thus, to be able to classify the instrumentation inside a Domain we would need to include the NVTX library python package as a dependency. This is the difficult part for the project, I guess.

After this is done, the implementation side would only require changing the following line of code:
https://github.com/deepspeedai/DeepSpeed/blob/5f7b687018bd1e0340c661859820fd97aa80a616/deepspeed/utils/nvtx.py#L19
In a simple implementation, adding the domain argument to the push_range call (as in here) would do the job. A more sophisticated solution that I think would also be great could be:

• Declaring previously the domain and instrumentation context, to avoid re-registering the strings every time.
• Also using the categories argument to internally classify the DeepSpeed's instrumentation (I don't have a proposal for that).

Describe alternatives you've considered
It is clear that there is a step that would make all this job easier, and it is PyTorch implementing all the NVTX library API. This would make this issue straightforward. I am considering opening an issue on PyTorch about this also.

Other alternatives are manually classifying in postprocessing based on the string DeepSpeed appearing on the name, but this is neigher always the case of this string appearing in all region's names. And this alternative is only limited to postprocessing limits, but still does not solve the possibility of filtering this at the level of tracing or analyzing.

Additional context
My use case is very specific, but is gaining relevance on the area of AI at HPC performance engineering. I know it is very tightened to the methodologies and tools used to analyze performance: it is very interesting to me because we do a deep post-processing and translation of Nsight Systems traces into a different tool.

I can implement myself the PR for this, but IDK the impact of adding a new dependency.

Start with deepspeed/utils/nvtx.py at the linked line and compare the current PyTorch NVTX interface with the NVIDIA NVTX Python package. Determine the dependency impact and whether domain or category support is in scope. Done means DeepSpeed regions can be separated from other NVTX instrumentation, with appropriate validation of the tracing behavior.