chore: import upstream snapshot with attribution

deepspeed/runtime/compiler.py

@@ -0,0 +1,113 @@
+# 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