|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
|
|
|
