wehub-resource-sync
114 lines
2.7 KiB
Python
114 lines
2.7 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
import contextlib
|
|
import functools
|
|
from deepspeed.utils.torch import required_torch_version
|
|
from deepspeed.accelerator import get_accelerator
|
|
|
|
try:
|
|
from torch.compiler import is_compiling as torch_is_compiling
|
|
except ImportError:
|
|
try:
|
|
from torch._dynamo.external_utils import is_compiling as torch_is_compiling
|
|
except ImportError:
|
|
# Torch does not have compiler support
|
|
torch_is_compiling = lambda: False
|
|
|
|
try:
|
|
if required_torch_version(min_version="2.6.0a"):
|
|
from torch._dynamo.compiled_autograd import _enable as compiled_autograd_enable
|
|
else:
|
|
from torch._dynamo.compiled_autograd import enable as compiled_autograd_enable
|
|
|
|
_COMPILED_AUTOGRAD_AVAILABLE = True
|
|
except (ImportError, ModuleNotFoundError):
|
|
_COMPILED_AUTOGRAD_AVAILABLE = False
|
|
|
|
|
|
def is_compile_supported():
|
|
return required_torch_version(min_version=2.1)
|
|
|
|
|
|
def disable(func):
|
|
if is_compile_supported():
|
|
return torch.compiler.disable(func)
|
|
return func
|
|
|
|
|
|
def enable(min_version=None):
|
|
"""
|
|
Decorator factory to enable compiling of a function if the minimum PyTorch version requirement is met.
|
|
|
|
Args:
|
|
min_version (str, optional): Minimum PyTorch version required (e.g., "2.7.0").
|
|
If None, the function is always enabled.
|
|
|
|
Returns:
|
|
Callable: A decorator that wraps the function.
|
|
|
|
Examples:
|
|
@enable("2.7.0")
|
|
def my_function():
|
|
pass
|
|
|
|
@enable
|
|
def another_function():
|
|
pass
|
|
"""
|
|
|
|
def decorator(func):
|
|
if not is_compiling():
|
|
return func
|
|
|
|
@functools.wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
if min_version is None or required_torch_version(min_version=min_version):
|
|
return func(*args, **kwargs)
|
|
return disable(func)(*args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
# Called with no arguments
|
|
if callable(min_version):
|
|
func = min_version
|
|
min_version = None
|
|
return decorator(func)
|
|
|
|
return decorator
|
|
|
|
|
|
def is_compiling():
|
|
return torch_is_compiling()
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def compiled_autograd(enabled: bool, kwargs: dict):
|
|
if not enabled or not _COMPILED_AUTOGRAD_AVAILABLE:
|
|
yield
|
|
return
|
|
|
|
if torch_is_compiling():
|
|
yield
|
|
return
|
|
|
|
compiler_fn = torch.compile(backend=get_accelerator().get_compile_backend(), **kwargs)
|
|
|
|
with compiled_autograd_enable(compiler_fn):
|
|
yield
|
|
|
|
|
|
def dummy_decorator(func):
|
|
return func
|
|
|
|
|
|
# robust version of @torch.compile
|
|
def compile():
|
|
if hasattr(torch, "compile"):
|
|
return torch.compile
|
|
else:
|
|
return dummy_decorator
|