wehub-resource-sync
74 lines
2.1 KiB
Python
74 lines
2.1 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import logging
|
|
import sys
|
|
|
|
from packaging import version as pkg_version
|
|
|
|
import torch
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
def required_torch_version(min_version=None, max_version=None):
|
|
assert min_version or max_version, "Must provide a min_version or max_version argument"
|
|
|
|
torch_version = pkg_version.parse(torch.__version__)
|
|
|
|
if min_version and pkg_version.parse(str(min_version)) > torch_version:
|
|
return False
|
|
|
|
if max_version and pkg_version.parse(str(max_version)) < torch_version:
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def is_functorch_transforming():
|
|
"""True when called under torch.func.grad / vmap / jacrev / etc."""
|
|
# peek_interpreter_stack is private; getattr keeps this working if it moves.
|
|
_functorch = getattr(torch._C, "_functorch", None)
|
|
if _functorch is None:
|
|
return False
|
|
peek = getattr(_functorch, "peek_interpreter_stack", None)
|
|
return peek is not None and peek() is not None
|
|
|
|
|
|
def register_grad_hook(param, hook):
|
|
if required_torch_version(min_version=2.1):
|
|
return param.register_post_accumulate_grad_hook(hook)
|
|
else:
|
|
param_tmp = param.expand_as(param)
|
|
grad_acc = param_tmp.grad_fn.next_functions[0][0]
|
|
return grad_acc.register_hook(hook)
|
|
|
|
|
|
def jit_script_compat(fn):
|
|
fn_name = getattr(fn, "__qualname__", getattr(fn, "__name__", repr(fn)))
|
|
|
|
can_try_compile = (required_torch_version(min_version=2.0) and hasattr(torch, "compile")
|
|
and not (sys.version_info >= (3, 12) and not required_torch_version(min_version=2.4)))
|
|
|
|
if can_try_compile:
|
|
try:
|
|
return torch.compile(fn)
|
|
except Exception:
|
|
_logger.debug(
|
|
"torch.compile failed for %s, falling back to torch.jit.script",
|
|
fn_name,
|
|
exc_info=True,
|
|
)
|
|
|
|
try:
|
|
return torch.jit.script(fn)
|
|
except Exception:
|
|
_logger.debug(
|
|
"torch.jit.script failed for %s, returning unmodified function",
|
|
fn_name,
|
|
exc_info=True,
|
|
)
|
|
return fn
|