wehub-resource-sync
39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.runtime.compiler import is_compiling
|
|
|
|
enable_nvtx = True
|
|
DEEPSPEED_NVTX_DOMAIN = "DeepSpeed"
|
|
|
|
|
|
def _range_push(accelerator, msg):
|
|
if getattr(accelerator, "supports_nvtx_domain", False):
|
|
return accelerator.range_push(msg, domain=DEEPSPEED_NVTX_DOMAIN)
|
|
return accelerator.range_push(msg)
|
|
|
|
|
|
def _range_pop(accelerator):
|
|
if getattr(accelerator, "supports_nvtx_domain", False):
|
|
return accelerator.range_pop(domain=DEEPSPEED_NVTX_DOMAIN)
|
|
return accelerator.range_pop()
|
|
|
|
|
|
def instrument_w_nvtx(func):
|
|
"""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 and not is_compiling():
|
|
_range_push(get_accelerator(), func.__qualname__)
|
|
ret_val = func(*args, **kwargs)
|
|
if enable_nvtx and not is_compiling():
|
|
_range_pop(get_accelerator())
|
|
return ret_val
|
|
|
|
return wrapped_fn
|