File size: 2,856 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
# 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
|