Avoid graph breaks by disabling sourceless calls in instrument_w_nvtx (…

…deepspeedai#7081)

This PR is a continuation of the efforts to improve Deepspeed
performance when using PyTorch compile.

The `instrument_w_nvtx` decorator is used to instrument code with NVIDIA
Tools Extension (NVTX) markers for profiling and visualizing code
execution on GPUs.

Along with executing the function itself, `instrument_w_nvtx` makes
calls to `nvtx.range_push` and `nvtx.range_pop` which can't be traced by
Dynamo.

That's why this decorator causes a graph break.
The impact on performance can be significant due to numerous uses of the
decorator throughout the code.

We propose a simple solution: Don't invoke the sourceless functions when
torch is compiling.

---------

Signed-off-by: Max Kovalenko <[email protected]>
Co-authored-by: Logan Adams <[email protected]>

deepspeed/utils/nvtx.py

@@ -4,19 +4,21 @@
 # DeepSpeed Team
 
 from deepspeed.accelerator import get_accelerator
+from deepspeed.runtime.compiler import is_compiling
 
 enable_nvtx = True
 
 
 def instrument_w_nvtx(func):
-    """decorator that causes an NVTX range to be recorded for the duration of the
-    function call."""
+    """Decorator that records an NVTX range for the duration of the function call.
+       Skips NVTX instrumentation when torch.compile is active to avoid graph breaks.
+    """
 
     def wrapped_fn(*args, **kwargs):
-        if enable_nvtx:
+        if enable_nvtx and not is_compiling():
             get_accelerator().range_push(func.__qualname__)
         ret_val = func(*args, **kwargs)
-        if enable_nvtx:
+        if enable_nvtx and not is_compiling():
             get_accelerator().range_pop()
         return ret_val