jamtur01's picture
Upload folder using huggingface_hub
9c6594c verified
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
import torch
from deepspeed.utils.torch import required_torch_version
backward_inputs = []
enabled_patched_func = False
original_grad_fn = None
base_meta = type(torch.autograd.Function)
if required_torch_version(min_version=2.7):
class FunctionMeta(base_meta):
def __new__(cls, name, bases, dct):
if name == "CompiledFunction":
original_backward_impl = dct.get("_backward_impl")
def wrapped_backward_impl(ctx, all_args):
assert original_backward_impl is not None
if enabled_patched_func:
backward_inputs.append(all_args)
wrapped_backward_impl.owner_class.compiled_bw = None
return original_backward_impl(ctx, all_args)
wrapped_backward_impl.owner_class = None
dct["_backward_impl"] = staticmethod(wrapped_backward_impl)
new_class = super().__new__(cls, name, bases, dct)
wrapped_backward_impl.owner_class = new_class
return new_class
return super().__new__(cls, name, bases, dct)
elif required_torch_version(min_version=2.6):
class FunctionMeta(base_meta):
def __new__(cls, name, bases, dct):
if name == "CompiledFunction":
original_backward_prologue = dct.get("_backward_prologue")
def wrapped_backward_prologue(ctx, *grad_outputs):
assert original_backward_prologue is not None
all_args = original_backward_prologue(ctx, *grad_outputs)
if enabled_patched_func:
backward_inputs.append(all_args)
wrapped_backward_prologue.owner_class.compiled_bw = None
return all_args
wrapped_backward_prologue.owner_class = None
dct["_backward_prologue"] = staticmethod(wrapped_backward_prologue)
new_class = super().__new__(cls, name, bases, dct)
wrapped_backward_prologue.owner_class = new_class
return new_class
return super().__new__(cls, name, bases, dct)
def patch_compiled_func():
global enabled_patched_func
enabled_patched_func = True
class PatchedFunction(torch.autograd.Function, metaclass=FunctionMeta):
pass
global original_grad_fn
original_grad_fn = torch.autograd.Function
torch.autograd.Function = PatchedFunction
return backward_inputs
def unpatch_compiled_func():
global enabled_patched_func
enabled_patched_func = False
global original_grad_fn
torch.autograd.Function = original_grad_fn
def get_backward_inputs():
return backward_inputs