mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
78bff1e8c1
See #145101 for details. Pull Request resolved: https://github.com/pytorch/pytorch/pull/145139 Approved by: https://github.com/bobrenjc93
778 lines
28 KiB
Python
778 lines
28 KiB
Python
# mypy: allow-untyped-defs
|
|
from typing import NamedTuple
|
|
|
|
import torch
|
|
import torch.utils._pytree as pytree
|
|
from torch._C._functorch import (
|
|
_unwrap_for_grad,
|
|
_wrap_for_grad,
|
|
current_level,
|
|
TransformType,
|
|
)
|
|
from torch._functorch.apis import vmap
|
|
from torch._functorch.utils import enable_single_level_autograd_function
|
|
from torch._functorch.vmap import (
|
|
_add_batch_dim,
|
|
_broadcast_to_and_flatten,
|
|
restore_vmap,
|
|
unwrap_batched,
|
|
wrap_batched,
|
|
)
|
|
from torch._ops import HigherOrderOperator
|
|
from torch.autograd.forward_ad import _set_fwd_grad_enabled
|
|
|
|
|
|
# autograd.Function technically runs before the regular PyTorch dispatcher.
|
|
# This is how features like autocast and torch_dispatch (e.g. PythonTLSSnapshot)
|
|
# work with it. One day we might decide to change this, but until then,
|
|
# we need to give the illusion that autograd.Function runs before those things.
|
|
#
|
|
# We do this by using creating a custom HigherOrderOperator that only functorch
|
|
# dispatches specially.
|
|
class CustomFunctionHigherOrderOperator(HigherOrderOperator):
|
|
def __init__(self) -> None:
|
|
super().__init__("custom_function_call")
|
|
|
|
def __call__(self, autograd_function, *args, **kwargs):
|
|
# When custom_function_call is done dispatching through functorch,
|
|
# it should just invoke the autograd.Function. This is consistent
|
|
# with the autograd.Function behavior of being invoked before the
|
|
# PyTorch dispatcher.
|
|
#
|
|
# This will lead us into trouble later down the line, but this is
|
|
# pre-existing. There is an invariant that a function traced by
|
|
# make_fx should have the same behavior when provided the same
|
|
# Tensor. However, make_fx sees autograd.Function as a composite
|
|
# (because autograd.Function happens before the Python dispatch key)
|
|
# and only traces the forward pass.
|
|
if torch._C._are_functorch_transforms_active():
|
|
return super().__call__(autograd_function, *args, **kwargs)
|
|
return autograd_function.apply(*args, **kwargs)
|
|
|
|
|
|
# "custom_function_call"
|
|
# This is the mechanism for an autograd.Function that works with functorch transforms.
|
|
# It wraps an autograd.Function; interactions with functorch transforms are defined
|
|
# via PyDispatcher and HigherOrderOperator rather than through the traditional PyTorch
|
|
# dispatcher.
|
|
custom_function_call = CustomFunctionHigherOrderOperator()
|
|
|
|
|
|
# The grad rule for custom_function_call is to construct a new _SingleLevelFunction
|
|
# (autograd.Function that only works with a single layer (level) of functorch) that:
|
|
# - unwraps the inputs
|
|
# - redispatches to custom_function_call
|
|
# - wraps the outputs
|
|
# and whose backward pass calls the original autograd.Function's backward.
|
|
#
|
|
# Why do we need to redispatch to custom_function_call?
|
|
# -----------------------------------------------------
|
|
# This is consistent with how ATen operators work with functorch's grad transform:
|
|
# they always redispatch to the original operator.
|
|
# Consider torch.sin, and let's say we do grad0(grad1(torch.sin))(x)
|
|
#
|
|
# grad1 will:
|
|
# - set up the autograd graph
|
|
# - unwrap the inputs
|
|
# - redispatch to at::sin (*)
|
|
# - rewrap the outputs on the return
|
|
#
|
|
# On the redispatch in (*), grad0 will:
|
|
# - set up the autograd graph
|
|
# - unwrap the inputs
|
|
# - redispatch to at::sin
|
|
# - rewrap the outputs on the return
|
|
#
|
|
# To "set up the autograd graph", we generate a _SingleLevelFunction
|
|
# and apply it.
|
|
@custom_function_call.py_impl(TransformType.Grad)
|
|
@custom_function_call.py_impl(TransformType.Jvp)
|
|
def custom_function_call_grad(interpreter, autograd_function, *operands):
|
|
Generated = generate_single_level_function(interpreter, autograd_function)
|
|
with enable_single_level_autograd_function():
|
|
flat_out = Generated.apply(*operands)
|
|
return flat_out
|
|
|
|
|
|
def generate_single_level_function(interpreter, autograd_function):
|
|
level = interpreter.level()
|
|
|
|
def forward(*operands):
|
|
unwrapped_operands = pytree.tree_map_only(
|
|
torch.Tensor, lambda x: _unwrap_for_grad(x, level), operands
|
|
)
|
|
# Both enable_grad() and _set_fwd_grad_enabled() are necessary no matter
|
|
# the transform. _SingleLevelFunction will turn off both fwd and bwd
|
|
# gradient computation and we need to turn it back on here.
|
|
with torch.enable_grad(), _set_fwd_grad_enabled(True), interpreter.lower():
|
|
unwrapped_output = custom_function_call(
|
|
autograd_function, *unwrapped_operands
|
|
)
|
|
|
|
# See NOTE [mark_dirty object identity check]
|
|
def wrap_fn(output):
|
|
return _wrap_for_grad(output, level)
|
|
|
|
return wrap_outputs_maintaining_identity(
|
|
unwrapped_output, unwrapped_operands, operands, wrap_fn
|
|
)
|
|
|
|
def setup_context(ctx, inputs, output):
|
|
return autograd_function.setup_context(ctx, inputs, output)
|
|
|
|
# backward is only used if the transform is TransformType.Grad
|
|
def backward(ctx, *grads):
|
|
result = autograd_function.backward(ctx, *grads)
|
|
return result
|
|
|
|
# jvp is only used if the transform is TransformType.Jvp
|
|
def jvp(ctx, *tangents):
|
|
result = autograd_function.jvp(ctx, *tangents)
|
|
return result
|
|
|
|
# This is the sequence of magic words to dynamically generate a Subclass with
|
|
# a given name. A Tensor's .grad_fn field has a class name that is the original
|
|
# autograd.Function's name + Backward, so we do this to generate some
|
|
# meaningful name.
|
|
name = f"{autograd_function.__name__}Generated"
|
|
Generated = type(
|
|
name,
|
|
(torch.autograd.function._SingleLevelFunction,),
|
|
{
|
|
"forward": staticmethod(forward),
|
|
"backward": staticmethod(backward),
|
|
"jvp": staticmethod(jvp),
|
|
"setup_context": staticmethod(setup_context),
|
|
},
|
|
)
|
|
return Generated
|
|
|
|
|
|
# wrap_outputs_maintaining_identity handles outputs from the vmap,
|
|
# backward (vjp), and jvp staticmethod. The way it distinguishes
|
|
# between the vmap case and the {backward, jvp} case is if the out_dims
|
|
# are specified or not.
|
|
#
|
|
# NB: we cannot use out_dims=None as the deciding factor. This because
|
|
# out_dims=None can still happen in the vmap staticmethod! What the
|
|
# user is saying in that case is that their output does not have a
|
|
# dimension that is being vmapped over, which is valid.
|
|
NO_OUT_DIMS = "not specified"
|
|
|
|
|
|
# NOTE [mark_dirty object identity check]
|
|
# autograd.Function's ctx.mark_dirty expect a returned input
|
|
# to have the same object identity as the input.
|
|
# Mode-only functorch will greatly simplify this logic.
|
|
def wrap_outputs_maintaining_identity(
|
|
outputs, unwrapped_inputs, orig_inputs, wrap_fn, out_dims=NO_OUT_DIMS
|
|
):
|
|
flat_unwrapped_inputs = pytree.arg_tree_leaves(*unwrapped_inputs)
|
|
flat_orig_inputs = pytree.arg_tree_leaves(*orig_inputs)
|
|
|
|
unwrapped_input_to_orig_input = {
|
|
id(unwrapped): orig
|
|
for unwrapped, orig in zip(flat_unwrapped_inputs, flat_orig_inputs)
|
|
}
|
|
|
|
flat_outputs, spec = pytree.tree_flatten(outputs)
|
|
result = []
|
|
|
|
out_dims_specified = out_dims != NO_OUT_DIMS
|
|
|
|
if out_dims_specified:
|
|
flat_out_dims = _broadcast_to_and_flatten(out_dims, spec)
|
|
# _broadcast_to_and_flatten returns None if it is unable to broadcast.
|
|
# TODO: update following link from master to stable once that's out
|
|
if flat_out_dims is None:
|
|
raise RuntimeError(
|
|
f"The autograd.Function's vmap staticmethod returned an "
|
|
f"incompatible (output, out_dims) tuple. "
|
|
f"Expected out_dims={out_dims} "
|
|
f"to be compatible with the structure of `output`. "
|
|
f"out_dims has structure {pytree.tree_flatten(out_dims)[1]} "
|
|
f"but output has structure {spec}. "
|
|
f"For more details, please see "
|
|
f"https://pytorch.org/docs/main/notes/extending.func.html"
|
|
)
|
|
|
|
for i, output in enumerate(flat_outputs):
|
|
if not isinstance(output, torch.Tensor):
|
|
result.append(output)
|
|
continue
|
|
if id(output) in unwrapped_input_to_orig_input:
|
|
result.append(unwrapped_input_to_orig_input[id(output)])
|
|
continue
|
|
if out_dims_specified:
|
|
result.append(wrap_fn(output, flat_out_dims[i])) # type: ignore[possibly-undefined, index]
|
|
else:
|
|
result.append(wrap_fn(output))
|
|
|
|
return pytree.tree_unflatten(result, spec)
|
|
|
|
|
|
# NOTE: [functorch vjp and autograd interaction]
|
|
# There's an edge case with the functorch vjp and autograd interaction
|
|
# that will eventually be fixed by mode-only functorch.
|
|
# The TL;DR is that there's no way to unwrap a dead GradTensorWrapper,
|
|
# so we (the framework) need to do it manually. Regular PyTorch operators
|
|
# automatically do so this is consistent.
|
|
#
|
|
# class MyExp(torch.autograd.Function):
|
|
# @staticmethod
|
|
# def forward(x):
|
|
# return x.exp()
|
|
#
|
|
# @staticmethod
|
|
# def setup_context(ctx, inputs, output):
|
|
# y = output
|
|
# ctx.save_for_backward(y)
|
|
#
|
|
# @staticmethod
|
|
# def backward(gy):
|
|
# y, = ctx.saved_tensors()
|
|
# return MyMul.apply(gy, y)
|
|
#
|
|
# x = torch.randn([], requires_grad=True)
|
|
# gy = torch.randn([], requires_grad=True)
|
|
# _, vjp_fn = vjp(MySin.apply, x)
|
|
# result = vjp_fn(gy)
|
|
#
|
|
# MyMul is an autograd.Function that is not shown here.
|
|
# It saves a `y` for backward (since gy requires grad).
|
|
#
|
|
# in vjp_fn(gy), we get:
|
|
# > MyMul.apply(gy, GradTensorWrapper(y, level=dead))
|
|
# Because the y that is saved for backward by MyExp is a GradTensorWrapper
|
|
# but is now dead since we are outside the vjp context.
|
|
#
|
|
# PyTorch dispatcher operations, upon seeing a dead GradTensorWrapper,
|
|
# will automatically unwrap the GradTensorWrapper when applied.
|
|
# But since autograd.Function technically sits above the regular PyTorch
|
|
# dispatcher, it doesn't get this treatment. So we manually do
|
|
# the unwrapping to be consistent with regular PyTorch dispatcher operations.
|
|
|
|
|
|
class VmapInfo(NamedTuple):
|
|
batch_size: int
|
|
randomness: str
|
|
|
|
|
|
def has_overriden_vmap_rule(autograd_function):
|
|
return autograd_function.vmap is not torch.autograd.Function.vmap
|
|
|
|
|
|
def validate_vmap_returns_tuple_of_two_elements(result):
|
|
base_error_msg = (
|
|
"Expected the vmap staticmethod to have two returns, an output "
|
|
"and out_dims with pytree structure compatible with the output. "
|
|
)
|
|
if not isinstance(result, tuple):
|
|
raise RuntimeError(base_error_msg + f"Got a {type(result)} instead")
|
|
if not len(result) == 2:
|
|
raise RuntimeError(base_error_msg + f"Got {len(result)} returns instead")
|
|
|
|
|
|
@custom_function_call.py_impl(TransformType.Vmap)
|
|
def custom_function_call_vmap(interpreter, autograd_function, *operands, **kwargs):
|
|
if any(
|
|
isinstance(val, torch.Tensor)
|
|
for val in torch.utils._pytree.tree_flatten(kwargs)[0]
|
|
):
|
|
raise NotImplementedError(
|
|
f"Run vmap on autograd.Function with kwarg-only Tensor args. "
|
|
f"Please do not pass kwarg-only Tensors to autograd.Function. "
|
|
f"Got: {kwargs}"
|
|
)
|
|
|
|
if autograd_function.generate_vmap_rule:
|
|
if has_overriden_vmap_rule(autograd_function):
|
|
# TODO: Update link to stable once that's out
|
|
# https://github.com/pytorch/pytorch/issues/92029
|
|
raise RuntimeError(
|
|
f"You tried to vmap over {autograd_function.__name__}, but "
|
|
f"it has both generate_vmap_rule=True and an overriden vmap "
|
|
f"staticmethod. Please set generate_vmap_rule=False or delete "
|
|
f"the overriden vmap staticmethod to avoid ambiguity. "
|
|
f"For more details, please see "
|
|
f"https://pytorch.org/docs/main/notes/extending.func.html"
|
|
)
|
|
return custom_function_call_vmap_generate_rule(
|
|
interpreter, autograd_function, *operands
|
|
)
|
|
|
|
if not has_overriden_vmap_rule(autograd_function):
|
|
# TODO: Update link to stable once that's out
|
|
# https://github.com/pytorch/pytorch/issues/92029
|
|
raise RuntimeError(
|
|
f"You tried to vmap over {autograd_function.__name__}, but "
|
|
f"it does not have vmap support. Please override and implement the "
|
|
f"vmap staticmethod or set generate_vmap_rule=True. "
|
|
f"For more details, please see "
|
|
f"https://pytorch.org/docs/main/notes/extending.func.html"
|
|
)
|
|
|
|
return custom_function_call_vmap_helper(
|
|
interpreter, autograd_function.vmap, autograd_function, *operands, **kwargs
|
|
)
|
|
|
|
|
|
def custom_function_call_vmap_helper(
|
|
interpreter, vmap_function, op, *operands, **kwargs
|
|
):
|
|
current_level = interpreter.level()
|
|
info = VmapInfo(
|
|
batch_size=interpreter.batch_size(),
|
|
randomness=interpreter.randomness(),
|
|
)
|
|
# We're either in the autograd.Function case (vmap staticmethod)
|
|
# or the torch.library.register_vmap case.
|
|
autograd_function_case = isinstance(op, torch.autograd.function.FunctionMeta)
|
|
|
|
def lower_to_next():
|
|
if autograd_function_case:
|
|
return interpreter.lower()
|
|
else:
|
|
return torch._C._ExcludeDispatchKeyGuard(
|
|
torch._C.DispatchKeySet(torch._C.DispatchKey.FuncTorchBatched)
|
|
)
|
|
|
|
unwrapped_operands, in_dims = unwrap_batched(operands, current_level)
|
|
# If none of the tensors are batched at the current level, then we skip the
|
|
# current level. This saves the user from needing to handle this case in
|
|
# their vmap staticmethod (and is consistent with our C++ batching rule API)
|
|
if pytree.tree_all(lambda dim: dim is None, in_dims):
|
|
with lower_to_next():
|
|
if autograd_function_case:
|
|
return custom_function_call(op, *operands)
|
|
else:
|
|
return op(*operands, **kwargs)
|
|
|
|
with lower_to_next():
|
|
result = vmap_function(info, in_dims, *unwrapped_operands, **kwargs)
|
|
validate_vmap_returns_tuple_of_two_elements(result)
|
|
unwrapped_output, out_dims = result
|
|
|
|
# See NOTE [mark_dirty object identity check]
|
|
def wrap_fn(output, out_dim):
|
|
return (
|
|
output
|
|
if out_dim is None
|
|
else _add_batch_dim(output, out_dim, current_level)
|
|
)
|
|
|
|
return wrap_outputs_maintaining_identity(
|
|
unwrapped_output, unwrapped_operands, operands, wrap_fn, out_dims=out_dims
|
|
)
|
|
|
|
|
|
def unpack_outputs(outputs):
|
|
out_dims = outputs[-1]
|
|
if isinstance(out_dims, tuple):
|
|
outputs = outputs[:-1]
|
|
else:
|
|
outputs = outputs[0]
|
|
return outputs, out_dims
|
|
|
|
|
|
def custom_function_call_vmap_generate_rule(interpreter, autograd_function, *operands):
|
|
unwrapped_operands, in_dims = unwrap_batched(operands, interpreter.level())
|
|
vmapped_function = vmapify_autograd_function(
|
|
autograd_function, in_dims, interpreter.batch_size(), interpreter.randomness()
|
|
)
|
|
with interpreter.lower():
|
|
outputs = custom_function_call(vmapped_function, *unwrapped_operands)
|
|
|
|
assert isinstance(outputs, tuple)
|
|
outputs, out_dims = unpack_outputs(outputs)
|
|
return wrap_batched(outputs, out_dims, interpreter.level())
|
|
|
|
|
|
@custom_function_call.py_impl(TransformType.Functionalize)
|
|
def custom_function_call_functionalize(
|
|
interpreter, autograd_function, generate_vmap_rule, *operands
|
|
):
|
|
raise RuntimeError("NYI: Functionalize rule for custom_function_call")
|
|
|
|
|
|
def vmapify_autograd_function(autograd_function, in_dims, batch_size, randomness):
|
|
def forward(*operands):
|
|
outputs, out_dims = restore_vmap(
|
|
autograd_function.forward, in_dims, batch_size, randomness
|
|
)(*operands)
|
|
if isinstance(outputs, torch.Tensor):
|
|
return outputs, out_dims
|
|
else:
|
|
return *outputs, out_dims
|
|
|
|
def setup_context(ctx, inputs, outputs):
|
|
outputs, out_dims = unpack_outputs(outputs)
|
|
key = id(Generated)
|
|
|
|
def inner(inputs, outputs):
|
|
# wrapped_ctx.save_for_backward will:
|
|
# - unwrap batchedtensors into (tensor, bdim)
|
|
# - save_for_backward(*unwrapped_tensors)
|
|
# - assign the bdims to wrapped_ctx._pt_saved_tensors_bdims
|
|
wrapped_ctx = CtxCustomSave(ctx, current_level())
|
|
autograd_function.setup_context(wrapped_ctx, inputs, outputs)
|
|
|
|
# input_shapes are used for reductify later to reduce expanded gradients
|
|
# to the correct shape.
|
|
# See NOTE: [Why can't we rely on autograd to reduce expanded gradients?]
|
|
# for more details
|
|
input_shapes = tuple(
|
|
inp.shape if isinstance(inp, torch.Tensor) else None for inp in inputs
|
|
)
|
|
if not hasattr(ctx, "_pt_input_shapes"):
|
|
ctx._pt_input_shapes = {}
|
|
ctx._pt_input_shapes.update({key: input_shapes})
|
|
|
|
if not hasattr(ctx, "_pt_saved_tensors_bdims_stack"):
|
|
ctx._pt_saved_tensors_bdims_stack = {}
|
|
ctx._pt_saved_tensors_bdims_stack.update(
|
|
{key: (wrapped_ctx._pt_saved_tensors_bdims)}
|
|
)
|
|
|
|
# See NOTE: [Why do we need to run setup_context under a vmap?]
|
|
restore_vmap(
|
|
inner,
|
|
(in_dims, out_dims),
|
|
batch_size,
|
|
randomness,
|
|
)(inputs, outputs)
|
|
|
|
if not hasattr(ctx, "_pt_out_dims"):
|
|
ctx._pt_out_dims = {}
|
|
ctx._pt_out_dims.update({key: out_dims})
|
|
|
|
def jvp(ctx, *tangents):
|
|
key = id(Generated)
|
|
|
|
def jvp_no_context(saved_tensors, tangents):
|
|
wrapped_ctx = CtxWithSavedTensors(ctx, saved_tensors)
|
|
return autograd_function.jvp(wrapped_ctx, *tangents)
|
|
|
|
tangent_in_dims = get_tangents_in_dims(in_dims, tangents)
|
|
out_tangents, out_tangents_dims = restore_vmap(
|
|
jvp_no_context,
|
|
(ctx._pt_saved_tensors_bdims_stack[key], tangent_in_dims),
|
|
batch_size,
|
|
randomness,
|
|
)(ctx.saved_tensors, tangents)
|
|
|
|
result = reductify(
|
|
out_tangents, out_tangents_dims, ctx._pt_out_dims[key], batch_size
|
|
)
|
|
if isinstance(result, torch.Tensor):
|
|
return result, None
|
|
else:
|
|
return *result, None
|
|
|
|
def backward(ctx, *grad_outputs):
|
|
key = id(Generated)
|
|
grad_outputs_ = grad_outputs[:-1]
|
|
grad_outputs_in_dims = ctx._pt_out_dims[key]
|
|
|
|
if not isinstance(grad_outputs_in_dims, tuple):
|
|
grad_outputs_in_dims = (grad_outputs_in_dims,)
|
|
|
|
grad_outputs_in_dims = tuple(
|
|
in_dim if grad_output is not None else None
|
|
for grad_output, in_dim in zip(grad_outputs_, grad_outputs_in_dims)
|
|
)
|
|
|
|
def backward_no_context(inputs):
|
|
saved_tensors, grad_outputs = inputs
|
|
wrapped_ctx = CtxWithSavedTensors(ctx, saved_tensors)
|
|
return autograd_function.backward(wrapped_ctx, *grad_outputs)
|
|
|
|
grad_ins, grad_ins_dims = restore_vmap(
|
|
backward_no_context,
|
|
((ctx._pt_saved_tensors_bdims_stack[key], grad_outputs_in_dims),),
|
|
batch_size,
|
|
randomness,
|
|
)((ctx.saved_tensors, grad_outputs_))
|
|
result = reductify(
|
|
grad_ins, grad_ins_dims, in_dims, batch_size, ctx._pt_input_shapes[key]
|
|
)
|
|
return result
|
|
|
|
name = f"Vmapped{autograd_function.__name__}"
|
|
Generated = type(
|
|
name,
|
|
(torch.autograd.Function,),
|
|
{
|
|
"forward": staticmethod(forward),
|
|
"backward": staticmethod(backward),
|
|
"jvp": staticmethod(jvp),
|
|
"setup_context": staticmethod(setup_context),
|
|
"generate_vmap_rule": True,
|
|
},
|
|
)
|
|
|
|
return Generated
|
|
|
|
|
|
# tangents might be None, so we need to replace
|
|
# the corresponding in_dims with None.
|
|
def get_tangents_in_dims(input_dims, tangents):
|
|
flat_in_dims, spec = pytree.tree_flatten(input_dims)
|
|
flat_tangents = pytree.arg_tree_leaves(*tangents)
|
|
result = [
|
|
None if tangent is None else in_dim
|
|
for in_dim, tangent in zip(flat_in_dims, flat_tangents)
|
|
]
|
|
return pytree.tree_unflatten(result, spec)
|
|
|
|
|
|
# NOTE: [Why do we need to run setup_context under a vmap?]
|
|
# Consider the following autograd.Function
|
|
#
|
|
# class Sum(torch.autograd.Function):
|
|
# @staticmethod
|
|
# def forward(x):
|
|
# return x.sum()
|
|
# @staticmethod
|
|
# def setup_context(ctx, inputs, outputs):
|
|
# ctx.x_shape = inputs[0]
|
|
# @staticmethod
|
|
# def backward(ctx, gy):
|
|
# return gy.expand(ctx.x_shape)
|
|
#
|
|
# x = torch.randn(B, 4)
|
|
# in_dims = 0
|
|
# vmap(Sum.apply, in_dims)(x)
|
|
#
|
|
# Let's assume for a moment that we didn't vmap setup_context in VmappedSum:
|
|
#
|
|
# class VmappedSum(torch.autograd.Function):
|
|
# @staticmethod
|
|
# def forward(x):
|
|
# return vmap(Sum.forward, in_dims)(x)
|
|
#
|
|
# @staticmethod
|
|
# def setup_context(ctx, inputs, outputs):
|
|
# Sum.setup_context(ctx, inputs, outputs)
|
|
#
|
|
# @staticmethod
|
|
# def backward(ctx, gy):
|
|
# def backward_no_context(gy):
|
|
# return gy.expand(ctx.x_shape)
|
|
#
|
|
# dims = (0,)
|
|
# gx = vmap(backward_no_context, dims)(gy)
|
|
# return gx
|
|
#
|
|
# We end up saving [B, 4] as x_shape. In the backward, gy has shape [B],
|
|
# and we're doing:
|
|
#
|
|
# def backward_no_context(gy):
|
|
# return gy.expand([B, 4])
|
|
#
|
|
# gx = vmap(backward_no_context, dims)(gy: "Tensor[B]")
|
|
#
|
|
# This gives us the wrong result (gx has shape [B, B, 4], but it should
|
|
# have shape [4]). Performing vmap over setup_context means the shape
|
|
# saved has shape [4] and leads to a correct result shape for gx.
|
|
|
|
|
|
# Wraps a ctx object. Forwards all attr accesses to the underlying object
|
|
# except for the attrs in _pt_attrs
|
|
class WrappedCtx:
|
|
_pt_reserved_attrs: tuple[str, ...] = ("_pt_reserved_attrs", "_pt_inner_ctx")
|
|
|
|
def __init__(self, ctx):
|
|
if not isinstance(ctx, WrappedCtx):
|
|
reserved_attrs = type(self)._pt_reserved_attrs
|
|
for name in reserved_attrs:
|
|
if not hasattr(ctx, name):
|
|
continue
|
|
raise RuntimeError(
|
|
f"PyTorch reserves the {reserved_attrs} field on ctx. "
|
|
"Please name your fields on ctx something else to avoid name "
|
|
"collision."
|
|
)
|
|
self._pt_inner_ctx = ctx
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._pt_inner_ctx, name)
|
|
|
|
def __setattr__(self, name, value):
|
|
if name in type(self)._pt_reserved_attrs:
|
|
self.__dict__[name] = value
|
|
return
|
|
return setattr(self._pt_inner_ctx, name, value)
|
|
|
|
|
|
# Wraps ctx to create a new ctx object that overrides saved_tensors.
|
|
class CtxWithSavedTensors(WrappedCtx):
|
|
_pt_reserved_attrs = ("_pt_new_saved_tensors", *WrappedCtx._pt_reserved_attrs)
|
|
|
|
def __init__(self, ctx, new_saved_tensors):
|
|
super().__init__(ctx)
|
|
self._pt_new_saved_tensors = new_saved_tensors
|
|
|
|
@property
|
|
def saved_tensors(self):
|
|
return self._pt_new_saved_tensors
|
|
|
|
|
|
class CtxCustomSave(WrappedCtx):
|
|
_pt_reserved_attrs = (
|
|
"_pt_saved_tensors_bdims",
|
|
"_pt_current_level",
|
|
*WrappedCtx._pt_reserved_attrs,
|
|
)
|
|
|
|
def __init__(self, ctx, current_level):
|
|
super().__init__(ctx)
|
|
self._pt_saved_tensors_bdims = ()
|
|
self._pt_current_level = current_level
|
|
|
|
def save_for_backward(self, *tensors):
|
|
unwrapped_tensors, bdims = unwrap_batched(tensors, self._pt_current_level)
|
|
self._pt_inner_ctx.save_for_backward(*unwrapped_tensors)
|
|
self._pt_saved_tensors_bdims = bdims
|
|
|
|
def save_for_forward(self, *tensors):
|
|
unwrapped_tensors, bdims = unwrap_batched(tensors, self._pt_current_level)
|
|
self._pt_inner_ctx.save_for_forward(*unwrapped_tensors)
|
|
self._pt_saved_tensors_bdims = bdims
|
|
|
|
|
|
def reductify(
|
|
grad_input,
|
|
grad_input_bdim,
|
|
input_bdim,
|
|
batch_size,
|
|
target_shape_without_bdim_to_reduce_to=None,
|
|
):
|
|
if not isinstance(grad_input, tuple):
|
|
grad_input = (grad_input,)
|
|
if not isinstance(grad_input_bdim, tuple):
|
|
grad_input_bdim = (grad_input_bdim,)
|
|
if not isinstance(input_bdim, tuple):
|
|
input_bdim = (input_bdim,)
|
|
|
|
if target_shape_without_bdim_to_reduce_to is None:
|
|
target_shape_without_bdim_to_reduce_to = len(grad_input) * (None,)
|
|
result = tuple(
|
|
reductify_leaf(gi, gi_bdim, i_bdim, batch_size, maybe_ishape)
|
|
for gi, gi_bdim, i_bdim, maybe_ishape in zip(
|
|
grad_input,
|
|
grad_input_bdim,
|
|
input_bdim,
|
|
target_shape_without_bdim_to_reduce_to,
|
|
)
|
|
)
|
|
return result
|
|
|
|
|
|
def reductify_leaf(
|
|
grad_input,
|
|
grad_input_bdim,
|
|
input_bdim,
|
|
batch_size,
|
|
target_shape_without_bdim_to_reduce_to=None,
|
|
):
|
|
if grad_input is None:
|
|
return None
|
|
|
|
if grad_input_bdim is None and input_bdim is None:
|
|
return grad_input
|
|
|
|
if grad_input_bdim is not None and input_bdim is None:
|
|
return grad_input.sum(grad_input_bdim)
|
|
|
|
# NOTE: [Why can't we rely on autograd to reduce expanded gradients?]
|
|
# For reverse-mode AD,
|
|
# given a grad_input and input, it is valid for the user to return a
|
|
# grad_input that has a broadcasted shape when compared to the input.
|
|
# In this situation, autograd automatically reduces the grad_input to
|
|
# the shape of the input.
|
|
#
|
|
# However, when input_bdim is not None, we have problems.
|
|
#
|
|
# [example 1]
|
|
# grad_input: Tensor[3, 4], input: Tensor[B, 4]
|
|
# We can expand grad_input to Tensor[B, 3, 4], but that isn't broadcastable
|
|
# from [B, 4].
|
|
#
|
|
# [example 2]
|
|
# grad_input: Tensor[3, B, 4], input: Tensor[B, 4]
|
|
# We can swizzle grad_input to Tensor[B, 3, 4], but that isn't broadcastable
|
|
# from [B, 4].
|
|
#
|
|
# This means that we need to also reduce the grad_input to the shape of the
|
|
# input. This behavior is controlled by the `target_shape_without_bdim_to_reduce_to` flag;
|
|
# if not-None then we do the reducing manually, otherwise, we do not do a reduction.
|
|
assert input_bdim is not None
|
|
|
|
if grad_input_bdim is None:
|
|
grad_input = grad_input.unsqueeze(input_bdim)
|
|
new_shape = list(grad_input.shape)
|
|
new_shape[input_bdim] = batch_size
|
|
grad_input = grad_input.expand(new_shape)
|
|
grad_input_bdim = input_bdim
|
|
|
|
if target_shape_without_bdim_to_reduce_to is not None:
|
|
return vmap(
|
|
torch.Tensor.sum_to_size,
|
|
in_dims=(grad_input_bdim, None),
|
|
out_dims=input_bdim,
|
|
)(grad_input, target_shape_without_bdim_to_reduce_to)
|
|
|
|
if input_bdim != grad_input_bdim:
|
|
grad_input = grad_input.movedim(grad_input_bdim, input_bdim)
|
|
return grad_input
|
|
|
|
|
|
def autograd_function_forward_rewritten(original_forward, original_setup_context):
|
|
def new_forward(ctx, *args, **kwargs):
|
|
output = original_forward(*args, **kwargs)
|
|
original_setup_context(ctx, args, output)
|
|
return output
|
|
|
|
return new_forward
|
|
|
|
|
|
class AutogradFunctionApply(HigherOrderOperator):
|
|
def __init__(self) -> None:
|
|
super().__init__("autograd_function_apply")
|
|
|
|
def __call__(self, fwd, bwd, *fwd_args, **fwd_kwargs):
|
|
saved_values = None
|
|
args_tensor_mask = fwd_kwargs["args_tensor_mask"]
|
|
non_differentiable_idx = fwd_kwargs["non_differentiable_idx"]
|
|
length_of_tensor_args = sum(args_tensor_mask)
|
|
# Filter out the original tensor args from fwd_args,
|
|
# lifted freevars should not be args of ApplyTemplate.apply
|
|
# since we don't need to calculate the gradients of them.
|
|
new_fwd_args = fwd_args[:length_of_tensor_args]
|
|
|
|
class ApplyTemplate(torch.autograd.Function):
|
|
@staticmethod
|
|
def forward(ctx, *args):
|
|
nonlocal saved_values
|
|
output, saved_values = fwd(None, *fwd_args)
|
|
|
|
# If users call ctx.mark_non_differentiable() in the original fwd function.
|
|
if len(non_differentiable_idx) > 0:
|
|
non_differentiable_output = []
|
|
for i, x in enumerate(output):
|
|
if i in non_differentiable_idx:
|
|
non_differentiable_output.append(x)
|
|
ctx.mark_non_differentiable(*non_differentiable_output)
|
|
|
|
return output
|
|
|
|
@staticmethod
|
|
def backward(ctx, *grad):
|
|
return bwd(None, *grad, *saved_values)
|
|
|
|
return ApplyTemplate.apply(*new_fwd_args)
|
|
|
|
|
|
autograd_function_apply = AutogradFunctionApply()
|