OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity
53cf46b8c6
pytorch/test/dynamo/test_autograd_function.py
Zhanghan Wang 53cf46b8c6 Fix names conflict when lifting (#129817) 
## Bug description
When pending args that are potentially to be lift [here](58f346c874/torch/_dynamo/output_graph.py (L1866)) having same base name, like `contiguous` and `contiguous_1`, the call into [create_graph_input](58f346c874/torch/_dynamo/output_graph.py (L2081)) can finally create a name ([here](58f346c874/torch/fx/graph.py (L1008))) that overwrite args to lift. And thus causing a wrong output of graph.

## Reproducing
Below is an reproduceable example,
```python
import logging
from typing import List

import torch
from functorch.compile import aot_module_simplified, make_boxed_func

@torch.library.custom_op("mylib::somefunc_forward", mutates_args=())
def somefunc_forward(
    input_: torch.Tensor,
    weight: torch.Tensor,
    shape: List[int],
) -> torch.Tensor:
    return torch.ones_like(input_)

@somefunc_forward.register_fake
def _(input_, shape, weight):
    return torch.empty_like(input_)

@torch.library.custom_op("mylib::somefunc_backward", mutates_args=())
def somefunc_backward(
    grad_output: torch.Tensor,
    input_: torch.Tensor,
    weight: torch.Tensor,
    shape: List[int],
) -> torch.Tensor:
    print(f"backward.{grad_output.shape=}")
    print(f"backward.{input_.shape=}")
    print(f"backward.{weight.shape=}")
    print(f"backward.{shape=}")
    assert list(weight.shape) == shape
    return torch.ones_like(weight)

@somefunc_backward.register_fake
def _(grad_output, input_, weight, shape):
    return torch.empty_like(weight)

def a_func(grad_output, input_, weight_, shape):
    return torch.ones_like(input_.sum() * weight_)

class SomeFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight, normalized_shape):
        ctx.normalized_shape = normalized_shape
        input_ = input.contiguous()
        weight_ = weight.contiguous()
        output = somefunc_forward(input_, weight_, ctx.normalized_shape)
        ctx.save_for_backward(input_, weight_)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        input_, weight_ = ctx.saved_tensors
        # grad_weight = a_func(grad_output, input_, weight_, ctx.normalized_shape)
        grad_weight = somefunc_backward(
            grad_output.contiguous(),
            input_,
            weight_,
            ctx.normalized_shape,
        )
        return None, grad_weight, None

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.ones(7))

    def forward(self, x):
        return SomeFunc.apply(x, self.weight, [7])

model = MyModel()
torch._logging.set_logs(dynamo=logging.DEBUG, aot=logging.DEBUG, graph_code=True)

def aot_print_backend(gm, sample_inputs):
    # Forward compiler capture
    def fw(gm, sample_inputs):
        print(f"----- fw")
        gm.print_readable()
        return make_boxed_func(gm.forward)

    # Backward compiler capture
    def bw(gm, sample_inputs):
        print(f"----- bw")
        gm.print_readable()
        return make_boxed_func(gm.forward)

    # Call AOTAutograd
    gm_forward = aot_module_simplified(
        gm, sample_inputs, fw_compiler=fw, bw_compiler=bw
    )
    return gm_forward

model = torch.compile(
    model,
    backend=aot_print_backend,
    dynamic=False,
)
out = model(torch.rand((128, 4, 7)))
out.mean().backward()
```

I can see log that showing calling into create_graph_input like
```log
V0629 02:08:46.839914 8200981504 torch/_dynamo/output_graph.py:2042] [0/0] create_graph_input contiguous (none)
V0629 02:08:46.839998 8200981504 torch/_dynamo/output_graph.py:2042] [0/0] create_graph_input contiguous_1 (none)
```

And the backward graph generate will be like
```log
class GraphModule(torch.nn.Module):
    def forward(self, function_ctx, somefunc_forward_default: "f32[128, 4, 7]", contiguous: "f32[128, 4, 7]", contiguous_1: "f32[7]"):
        contiguous_1 = contiguous
        contiguous_2 = contiguous_1

        # No stacktrace found for following nodes
        _set_grad_enabled = torch._C._set_grad_enabled(False)

         # File: /Users/bytedance/testtorch/test_custom_op_bug.py:61 in backward, code: grad_output.contiguous(),
        contiguous: "f32[128, 4, 7]" = somefunc_forward_default.contiguous();  somefunc_forward_default = None

         # File: /opt/tiger/pytorch/torch/_library/custom_ops.py:506 in __call__, code: return self._opoverload(*args, **kwargs)
        somefunc_backward_default: "f32[7]" = torch.ops.mylib.somefunc_backward.default(contiguous, contiguous_1, contiguous_2, [7]);  contiguous = contiguous_1 = contiguous_2 = None

        # No stacktrace found for following nodes
        _set_grad_enabled_1 = torch._C._set_grad_enabled(True)
        return (None, somefunc_backward_default)
```

The original code of `somefunc_backward` takes a input list of `grad_output`, `input_`, `weight` and `shape`, where `weight` should be shape of `torch.Size([7])`. However, in the graph, `contiguous1` and `contiguous_2` are assigned with `contiguous`, this leads to assertion failure I added in `somefunc_backward`.

## Environment
```log
Collecting environment information...
PyTorch version: 2.5.0a0+git0b7e8df
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.5 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.3.9.4)
CMake version: version 3.26.4
Libc version: N/A

Python version: 3.9.19 (main, May  6 2024, 14:39:30)  [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.5-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M3 Pro

Versions of relevant libraries:
[pip3] numpy==2.0.0
[pip3] optree==0.11.0
[pip3] torch==2.5.0a0+git0b7e8df
[pip3] torchgraph==0.0.1
[conda] numpy                     2.0.0                    pypi_0    pypi
[conda] optree                    0.11.0                   pypi_0    pypi
[conda] torch                     2.5.0a0+git0b7e8df           dev_0    <develop>
[conda] torchgraph                0.0.1                     dev_0    <develop>
```

## How to fix?

I put a naive fix that add the potential args to lift into the used_names. This visits private variables, will fix that if this issue makes sense to you.

@zou3519 @oulgen

Co-authored-by: rzou <zou3519@gmail.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/129817
Approved by: https://github.com/zou3519
2024-07-15 18:49:12 +00:00

1269 lines
39 KiB
Python
Raw Blame History

# Owner(s): ["module: dynamo"]
# flake8: noqa: B950
import copy
import math
from dataclasses import dataclass
from typing import List
import torch
import torch._dynamo.test_case
import torch._dynamo.testing
import torch._dynamo.utils
from functorch.compile import aot_module_simplified
from torch.testing._internal.triton_utils import HAS_CUDA, requires_cuda
if HAS_CUDA:
import triton
from torch.testing._internal.triton_utils import add_kernel
class CustomFunc1(torch.autograd.Function):
@staticmethod
def forward(ctx, foo):
return foo + foo
@staticmethod
def backward(ctx, grad_output):
return grad_output
class CustomFunc3(torch.autograd.Function):
# Test there is graph break in forward function
@staticmethod
def forward(ctx, foo):
result = foo + foo
torch._dynamo.graph_break()
result = result + foo
ctx.save_for_backward(result)
return result
@staticmethod
def backward(ctx, grad_output):
(result,) = ctx.saved_tensors
return grad_output * math.sqrt(result.numel())
class Module1(torch.nn.Module):
def forward(self, foo):
return CustomFunc1().apply(foo)
class Module2(torch.nn.Module):
def __init__(self):
super().__init__()
self.fn = CustomFunc1.apply
def forward(self, foo):
return self.fn(foo)
class Module3(torch.nn.Module):
def forward(self, foo):
return CustomFunc1().apply(foo)
class Module4(torch.nn.Module):
def __init__(self):
super().__init__()
self.fn = CustomFunc1.apply
def forward(self, foo):
return self.fn(foo)
class Module5(torch.nn.Module):
def forward(self, foo):
return CustomFunc3().apply(foo)
class Module6(torch.nn.Module):
def __init__(self):
super().__init__()
self.fn = CustomFunc3.apply
def forward(self, foo):
return self.fn(foo)
class LinearFunction(torch.autograd.Function):
# Note that forward, setup_context, and backward are @staticmethods
@staticmethod
def forward(input, weight, bias):
output = input.mm(weight.t())
if bias is not None:
output += bias.unsqueeze(0).expand_as(output)
return output
@staticmethod
# inputs is a Tuple of all of the inputs passed to forward.
# output is the output of the forward().
def setup_context(ctx, inputs, output):
input, weight, bias = inputs
ctx.save_for_backward(input, weight, bias)
# This function has only a single output, so it gets only one gradient
@staticmethod
def backward(ctx, grad_output):
input, weight, bias = ctx.saved_tensors
grad_input = grad_weight = grad_bias = None
if ctx.needs_input_grad[0]:
grad_input = grad_output.mm(weight)
if ctx.needs_input_grad[1]:
grad_weight = grad_output.t().mm(input)
if bias is not None and ctx.needs_input_grad[2]:
grad_bias = grad_output.sum(0)
return grad_input, grad_weight, grad_bias
class ModuleLinear(torch.nn.Module):
def forward(self, input, weight, bias=None):
return LinearFunction.apply(input, weight, bias)
class MaterializingGradFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
ctx.set_materialize_grads(False)
return x.clone(), x.clone()
@staticmethod
def backward(ctx, grad_out1, grad_out2):
return grad_out1, grad_out2
class MaterializingGradModule(torch.nn.Module):
def forward(self, x):
return MaterializingGradFunction.apply(x)
class CustomFuncBwdPrintGraphBreak(torch.autograd.Function):
@staticmethod
def forward(ctx, foo):
return torch.add(foo, foo)
@staticmethod
def backward(ctx, grad_output):
print("graph break!")
return grad_output
class CustomFuncBwdPrintModule(torch.nn.Module):
def forward(self, x):
return CustomFuncBwdPrintGraphBreak.apply(x)
class CustomFuncStrideBwd(torch.autograd.Function):
@staticmethod
def forward(ctx, foo):
return torch.add(foo, foo)
@staticmethod
def backward(ctx, grad_output):
return grad_output.stride()
class CustomFuncStrideModule(torch.nn.Module):
def forward(self, x):
return CustomFuncStrideBwd.apply(x)
class CustomFuncSaveForBwd(torch.autograd.Function):
@staticmethod
def forward(ctx, foo):
result = foo + foo
result = result + foo
ctx.save_for_backward(result)
return result
@staticmethod
def backward(ctx, grad_output):
(result,) = ctx.saved_tensors
return grad_output * math.sqrt(result.numel())
class SaveForBwdModule(torch.nn.Module):
def forward(self, foo):
return CustomFuncSaveForBwd().apply(foo)
class ContextSaveAndMark(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
with torch.no_grad():
ctx.save_for_backward(x)
ctx.mark_non_differentiable(x)
return x
@staticmethod
def backward(ctx, grad_output):
return grad_output
class ContextMarkAndSave(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
with torch.no_grad():
ctx.mark_non_differentiable(x)
ctx.save_for_backward(x)
return x
@staticmethod
def backward(ctx, grad_output):
return grad_output
class ModuleWithGradFunc(torch.nn.Module):
def __init__(self, func):
super().__init__()
self.f = func.apply
def forward(self, x):
return self.f(x)
@torch.library.custom_op("_torch_testing::custom_op_forward", mutates_args=())
def custom_op_forward(
foo: torch.Tensor,
bar: torch.Tensor,
shape: List[int],
) -> torch.Tensor:
return torch.ones_like(foo)
@custom_op_forward.register_fake
def _(foo, bar, weight):
return torch.empty_like(foo)
@torch.library.custom_op("_torch_testing::custom_op_backward", mutates_args=())
def custom_op_backward(
grad_output: torch.Tensor,
foo: torch.Tensor,
bar: torch.Tensor,
shape: List[int],
) -> torch.Tensor:
assert list(bar.shape) == shape
return torch.ones_like(bar)
@custom_op_backward.register_fake
def _(grad_output, foo, bar, shape):
return torch.empty_like(bar)
class CustomOpFunc(torch.autograd.Function):
@staticmethod
def forward(ctx, input, weight, normalized_shape):
ctx.normalized_shape = normalized_shape
input_ = input.contiguous()
weight_ = weight.contiguous()
output = custom_op_forward(input_, weight_, ctx.normalized_shape)
ctx.save_for_backward(input_, weight_)
return output
@staticmethod
def backward(ctx, grad_output):
input_, weight_ = ctx.saved_tensors
# grad_weight = a_func(grad_output, input_, weight_, ctx.normalized_shape)
grad_weight = custom_op_backward(
grad_output.contiguous(),
input_,
weight_,
ctx.normalized_shape,
)
return None, grad_weight, None
class CustomOpModule(torch.nn.Module):
def __init__(self, shape):
super().__init__()
self.shape = shape
self.weight = torch.nn.Parameter(torch.ones(self.shape))
def forward(self, x):
return CustomOpFunc.apply(x, self.weight, self.shape)
class AutogradFunctionTests(torch._dynamo.test_case.TestCase):
# Sound behaviors, tested for working capture
def test_autograd_function_equivalence(self):
for grad in [True, False]:
for i in range(1, 5):
torch._dynamo.reset()
model = globals()[f"Module{i}"]()
opt_model = torch._dynamo.optimize("eager")(model)
self.assertTrue(
torch.allclose(
opt_model(torch.ones(2, 3, requires_grad=grad)),
torch.tensor([2.0], requires_grad=grad),
)
)
def test_autograd_function_has_graph_break(self):
for grad in [True, False]:
x = torch.randn(10, requires_grad=grad)
for model in [Module5(), Module6()]:
torch._dynamo.reset()
cnts = torch._dynamo.testing.CompileCounter()
opt_model = torch._dynamo.optimize(cnts)(model)
for _ in range(3):
ref = model(x)
res = opt_model(x)
self.assertTrue(torch.allclose(ref, res))
self.assertEqual(cnts.frame_count, 2)
def test_linear_setup_context(self):
model = ModuleLinear()
opt_model = torch._dynamo.optimize("eager", nopython=True)(model)
input = torch.randn(2, 2, dtype=torch.double, requires_grad=True)
weight = torch.randn(3, 2, dtype=torch.double, requires_grad=True)
eager_result = model(input, weight)
optim_result = opt_model(input, weight)
self.assertEqual(optim_result, eager_result)
def test_materialize_grad(self):
model = MaterializingGradModule()
opt_model = torch._dynamo.optimize("eager")(model)
x = torch.randn(2, 2, dtype=torch.double, requires_grad=True)
optim_result = opt_model(x)
eager_result = model(x)
self.assertEqual(optim_result, eager_result)
def test_print_in_bwd(self):
model = CustomFuncBwdPrintModule()
opt_model = torch._dynamo.optimize("eager", nopython=True)(model)
x = torch.randn(2, 2, dtype=torch.double, requires_grad=True)
with self.assertRaisesRegex(torch._dynamo.exc.Unsupported, "builtin: print"):
opt_model(x)
def test_stride_in_bwd(self):
torch._dynamo.utils.counters.clear()
cnt = torch._dynamo.testing.CompileCounter()
model = CustomFuncStrideModule()
opt_model = torch.compile(backend=cnt)(model)
x = torch.randn(2, 2, dtype=torch.double, requires_grad=True)
ref = model(x)
res = opt_model(x)
self.assertEqual(ref, res)
self.assertEqual(cnt.frame_count, 1)
# graph break: Illegal getattr invocation stride in strict mod.
self.assertEqual(
list(torch._dynamo.utils.counters["graph_break"].values()), [1]
)
def test_enum_arg(self):
from enum import Enum
class SomeEnum(Enum):
A = 0
B = 1
class Foo(torch.autograd.Function):
@staticmethod
def forward(ctx, x, e):
if e is SomeEnum.A:
return x.sin()
else:
return x.cos()
@staticmethod
def backward(ctx, g):
return g
@torch.compile(backend="eager", fullgraph=True)
def f(x, enum):
output = Foo.apply(
x,
enum,
)
return output
x = torch.tensor([[1.0, 2, 3], [4, 5, 6]], requires_grad=True)
y = f(x, SomeEnum.A)
self.assertEqual(y, x.sin())
def test_save_for_bwd(self):
model = SaveForBwdModule()
opt_model = torch._dynamo.optimize("eager", nopython=True)(model)
x = torch.randn(2, 2, dtype=torch.double, requires_grad=True)
opt_model(x)
def test_allow_in_graph(self):
torch._dynamo.utils.counters.clear()
cnt = torch._dynamo.testing.CompileCounter()
@torch._dynamo.allow_in_graph
class AllowInGraphFunc(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
torch._dynamo.graph_break()
ctx.x0 = x.size(0)
return x * 2
@staticmethod
def backward(ctx, grad_out):
return grad_out * ctx.x0
@torch.compile(backend=cnt, fullgraph=True)
def fn(x):
return AllowInGraphFunc.apply(x)
x = torch.rand(2, 3, requires_grad=True)
result = fn(x)
self.assertEqual(result, AllowInGraphFunc.apply(x))
self.assertEqual(cnt.frame_count, 1)
def test_once_differentiable(self):
from torch.autograd.function import once_differentiable
torch._dynamo.utils.counters.clear()
cnt = torch._dynamo.testing.CompileCounter()
class ScaleGradient(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
return x
@staticmethod
@once_differentiable
def backward(ctx, grad):
return grad * 0.5
@torch.compile(backend=cnt, fullgraph=True)
def fn(x):
return ScaleGradient.apply(x)
x = torch.randn(3, requires_grad=True)
result = fn(x)
self.assertEqual(result, ScaleGradient.apply(x))
self.assertEqual(cnt.frame_count, 1)
def test_classmethod(self):
class Shake(torch.autograd.Function):
@classmethod
def forward(cls, ctx, foo):
return foo + foo
@classmethod
def backward(cls, ctx, grad_output):
return grad_output
def f(x):
return Shake.apply(x)
x = torch.randn(4, 4, 4, 4, requires_grad=True)
opt_m = torch.compile(backend="eager")(f)
opt_m(x)
def test_function_context_save_and_mark(self):
mod = ModuleWithGradFunc(ContextSaveAndMark)
args, kwargs = ([torch.rand([1])], {})
before = mod(*args, **kwargs)
torch._dynamo.reset()
compiled_model = torch._dynamo.optimize("eager")(mod)
after = compiled_model(*args, **kwargs)
self.assertEqual(before, after)
def test_function_context_mark_and_save(self):
mod = ModuleWithGradFunc(ContextMarkAndSave)
args, kwargs = ([torch.rand([1])], {})
before = mod(*args, **kwargs)
torch._dynamo.reset()
compiled_model = torch._dynamo.optimize("eager")(mod)
after = compiled_model(*args, **kwargs)
self.assertEqual(before, after)
def test_multi_output(self):
torch._dynamo.utils.counters.clear()
cnt = torch._dynamo.testing.CompileCounter()
class Foo(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
return x.clone(), x.clone()
@staticmethod
def backward(ctx, grad1, grad2):
return grad1 + grad2
@torch.compile(backend=cnt, fullgraph=True)
def f(x):
return Foo.apply(x)
x = torch.randn(3, requires_grad=True)
result = f(x)
self.assertEqual(result, Foo.apply(x))
self.assertEqual(cnt.frame_count, 1)
def test_amp_custom_fwd_bwd(self):
torch._dynamo.utils.counters.clear()
cnt = torch._dynamo.testing.CompileCounter()
class MyMM(torch.autograd.Function):
@staticmethod
@torch.amp.custom_fwd(device_type="cuda")
def forward(ctx, a, b):
ctx.save_for_backward(a, b)
return a.mm(b)
@staticmethod
@torch.amp.custom_bwd(device_type="cuda")
def backward(ctx, grad):
a, b = ctx.saved_tensors
return grad.mm(b.t()), a.t().mm(grad)
@torch.compile(backend=cnt, fullgraph=True)
def fn(a, b):
return MyMM.apply(a, b)
a = torch.randn([64, 64], dtype=torch.float32, requires_grad=True)
grad = a.clone()
res = fn(a, a)
res.backward(grad)
self.assertEqual(res, MyMM.apply(a, a))
self.assertEqual(cnt.frame_count, 1)
def test_user_defined_object_as_input(self):
cnt = torch._dynamo.testing.CompileCounterWithBackend("aot_eager")
@dataclass
class Weird:
x: int
b: torch.Tensor
c: torch.Tensor
class Foo(torch.autograd.Function):
@staticmethod
def forward(ctx, x: torch.Tensor, weird: Weird, z: torch.Tensor):
ctx.save_for_backward(weird.b, weird.c)
return weird.b * weird.c * x.clone()
@staticmethod
def backward(ctx, grad):
b, c = ctx.saved_tensors
return grad * b * c, None, grad * 2
@torch.compile(backend=cnt, fullgraph=True)
def f(x, weird, z):
return Foo.apply(x, weird, z)
x = torch.tensor(2.0, requires_grad=True)
weird = Weird(1.2, torch.tensor(2.5, requires_grad=True), torch.tensor(3.5))
z = torch.tensor(3.0, requires_grad=True)
result = f(x, weird, z)
result.sum().backward()
self.assertEqual(result, Foo.apply(x, weird, z))
self.assertEqual(x.grad, 2.5 * 3.5)
self.assertEqual(z.grad, 2.0)
self.assertEqual(weird.b.grad, None)
# check Dynamo captured graph is correct!
actual_graph = torch._dynamo.testing.normalize_gm(
cnt.graphs[0].print_readable(print_output=False)
)
self.assertExpectedInline(
actual_graph,
"""\
class GraphModule(torch.nn.Module):
def forward(self, L_x_: "f32[]", L_z_: "f32[]", L_weird_b: "f32[]", L_weird_c: "f32[]"):
l_x_ = L_x_
l_z_ = L_z_
l_weird_b = L_weird_b
l_weird_c = L_weird_c
function_ctx = torch.autograd.function.FunctionCtx()
fwd_body_0 = self.fwd_body_0
bwd_body_0 = self.bwd_body_0
autograd_function_apply: "f32[]" = torch._functorch.autograd_function.autograd_function_apply(fwd_body_0, bwd_body_0, l_x_, l_z_, l_weird_b, l_weird_c, args_tensor_mask = [True, False, True]); fwd_body_0 = bwd_body_0 = l_x_ = l_z_ = l_weird_b = l_weird_c = None
return (autograd_function_apply,)
class GraphModule(torch.nn.Module):
def forward(self, ctx, x: "f32[]", z: "f32[]", l_weird_b: "f32[]", l_weird_c: "f32[]"):
ctx_1 = ctx
x_1 = x
z_1 = z
l_weird_b_1 = l_weird_b
l_weird_c_1 = l_weird_c
mul: "f32[]" = l_weird_b_1 * l_weird_c_1
clone: "f32[]" = x_1.clone(); x_1 = None
mul_1: "f32[]" = mul * clone; mul = clone = None
return (mul_1, [l_weird_b_1, l_weird_c_1])
class GraphModule(torch.nn.Module):
def forward(self, ctx, grad: "f32[]", l_weird_b: "f32[]", l_weird_c: "f32[]"):
ctx_1 = ctx
grad_1 = grad
l_weird_b_1 = l_weird_b
l_weird_c_1 = l_weird_c
_set_grad_enabled = torch._C._set_grad_enabled(False)
mul: "f32[]" = grad_1 * l_weird_b_1; l_weird_b_1 = None
mul_1: "f32[]" = mul * l_weird_c_1; mul = l_weird_c_1 = None
mul_2: "f32[]" = grad_1 * 2; grad_1 = None
_set_grad_enabled_1 = torch._C._set_grad_enabled(True)
return (mul_1, mul_2)
""",
)
def test_tensor_list_as_input(self):
class Foo(torch.autograd.Function):
@staticmethod
def forward(ctx, x, tl):
ctx.save_for_backward(tl[0], tl[1])
return x.clone() * (tl[0] + tl[1])
@staticmethod
def backward(ctx, grad):
tl0, tl1 = ctx.saved_tensors
return grad * (tl0 + tl1), None
@torch.compile(backend="aot_eager", fullgraph=True)
def f(x, tl):
return Foo.apply(x, tl)
x = torch.tensor(2.0, requires_grad=True)
tl = [
torch.tensor(3.0, requires_grad=True),
torch.tensor(4.0, requires_grad=True),
]
result = f(x, tl)
result.sum().backward()
self.assertEqual(result, Foo.apply(x, tl))
self.assertEqual(x.grad, 7.0)
self.assertEqual(tl[0].grad, None)
self.assertEqual(tl[1].grad, None)
def test_multiple_different_non_tensor_inputs(self):
@dataclass
class Weird:
x: int
b: torch.Tensor
c: torch.Tensor
class Foo(torch.autograd.Function):
@staticmethod
def forward(ctx, x, weird, z, tl):
ctx.save_for_backward(weird.b, weird.c, tl[0], tl[1])
return x.clone() * weird.b * weird.c * tl[0]
@staticmethod
def backward(ctx, grad):
b, c, tl0, _ = ctx.saved_tensors
return grad * b * c * tl0, None, grad * 2, None
@torch.compile(backend="aot_eager", fullgraph=True)
def f(x, weird, z, tl):
return Foo.apply(x, weird, z, tl)
x = torch.tensor(2.0, requires_grad=True)
weird = Weird(
1.2,
torch.tensor(2.5, requires_grad=True),
torch.tensor(3.5, requires_grad=True),
)
z = torch.tensor(3.0, requires_grad=True)
tl = [
torch.tensor(0.5, requires_grad=True),
torch.tensor(0.6, requires_grad=True),
]
result = f(x, weird, z, tl)
result.sum().backward()
self.assertEqual(result, Foo.apply(x, weird, z, tl))
self.assertEqual(x.grad, 2.5 * 3.5 * 0.5)
self.assertEqual(z.grad, 2.0)
self.assertEqual(weird.b.grad, None)
self.assertEqual(weird.c.grad, None)
self.assertEqual(tl[0].grad, None)
self.assertEqual(tl[1].grad, None)
def test_backward_returns_none_for_tensor_input(self):
class Foo(torch.autograd.Function):
@staticmethod
def forward(ctx, x, y):
ctx.save_for_backward(y)
return x.clone() * y
@staticmethod
def backward(ctx, grad):
(y,) = ctx.saved_tensors
return grad * y, None
@torch.compile(backend="aot_eager", fullgraph=True)
def f(x, y):
return Foo.apply(x, y)
x = torch.tensor(2.0, requires_grad=True)
y = torch.tensor(3.0, requires_grad=True)
result = f(x, y)
result.sum().backward()
self.assertEqual(result, Foo.apply(x, y))
self.assertEqual(x.grad, 3.0)
self.assertEqual(y.grad, None)
def test_function_with_bound_free_variable(self):
class LowerBound(torch.autograd.Function):
@staticmethod
def forward(ctx, inputs, bound):
ctx.save_for_backward(inputs, inputs.new_ones(1) * bound)
return inputs.clamp(min=bound)
@staticmethod
def backward(ctx, grad_output):
inputs, bound = ctx.saved_tensors
return (inputs >= bound) * grad_output, None
class MyMod(torch.nn.Module):
def __init__(self):
super().__init__()
self.gamma = torch.nn.Parameter(torch.rand([4, 128, 32, 32]))
def forward(self, x):
gamma = LowerBound.apply(self.gamma, 1)
return x + gamma
mod = MyMod()
args, kwargs = ([torch.rand([4, 128, 32, 32])], {})
before = mod(*args, **kwargs)
compiled_model = torch._dynamo.optimize("eager")(mod)
after = compiled_model(*args, **kwargs)
self.assertEqual(before, after)
# I pulled all of these test cases from test_autograd.py
# In the future, we should make the Dynamo test suite actually
# run on test_autograd.py (it's disabled right now) and delete these.
def test_smoke_from_test_autograd(self):
class Func(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
out0 = x.clone()
out1 = x.clone()
ctx.mark_non_differentiable(out1)
ctx._materialize_non_diff_grads = False
return out0, out1
@staticmethod
def backward(ctx, g0, g1):
assert g1 is None
return g0
def mult1(x):
return x.prod(dim=-1).prod(dim=-1)
class Mult(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
y = mult1(x)
ctx.save_for_backward(x, y)
return y
@staticmethod
def backward(ctx, grad_output):
x, y = ctx.saved_tensors
return (grad_output * y)[:, None, None] / x
mult2 = Mult.apply
class Double(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
y = x**2
ctx.save_for_backward(x, y)
return y
@staticmethod
def backward(ctx, grad_output):
x, _ = ctx.saved_tensors
return grad_output * 2 * x
# this is equivalent, but uses the output of .forward() in .backward()
class Double2(Double):
@staticmethod
def backward(ctx, grad_output):
x, y = ctx.saved_tensors
return grad_output * 2 * y / x
double = Double.apply
double2 = Double2.apply
class Identity(torch.autograd.Function):
@staticmethod
def forward(ctx, a, b):
return a, a + b
@staticmethod
def backward(ctx, grad_a, grad_b):
return grad_a + grad_b, grad_b
class MyFunc2(torch.autograd.Function):
@staticmethod
def forward(ctx, inp):
return inp.clone()
@staticmethod
def backward(ctx, gO):
return torch.tensor(float("nan")).expand(10, 10)
def run_fn(a):
out = MyFunc2.apply(a)
return out.sum()
class MyFn(torch.autograd.Function):
@staticmethod
def forward(ctx, inp):
return inp.view_as(inp)
@staticmethod
def backward(ctx, grad):
return grad
class MyAdder(torch.autograd.Function):
@staticmethod
def forward(ctx, a, b):
a.add_(b)
ctx.mark_dirty(a)
return a
@staticmethod
def backward(ctx, grad):
return grad, grad
class InplaceMul(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
result = x.mul_(2)
ctx.mark_dirty(result)
return result
@staticmethod
def backward(ctx, grad_output):
pass
@staticmethod
def jvp(ctx, x_t):
if jvp_err: # noqa: F821
return x_t
else:
return x_t.mul_(2)
class MyFn2(torch.autograd.Function):
@staticmethod
def forward(ctx, x, y):
return x + y, x
@staticmethod
def vjp(ctx, gO1, gO2):
return gO1 + gO2, gO1
@staticmethod
def jvp(ctx, x_t, y_t):
return x_t + y_t, fn(x_t) # noqa: F821
class MyFn3(torch.autograd.Function):
@staticmethod
def forward(ctx, inp, inplace):
view = inp.clone()[:3]
if inplace:
view += 2
return view
@staticmethod
def backward(ctx, grad):
return grad, None
def test():
a = torch.tensor(1.0, requires_grad=True)
out = Func.apply(a)[0]
out.backward()
x = torch.ones(2, 4, 4).requires_grad_()
mult2(x)
x = torch.tensor(2).double().requires_grad_()
double(x)
double2(x)
x = torch.randn(5, 5, requires_grad=True)
y = torch.randn(5, 5, requires_grad=True)
q, p = Identity.apply(x, y)
a = torch.rand(1, 2)
b = torch.rand(1, requires_grad=True)
view_a = MyFn.apply(a)
a = torch.ones(2, requires_grad=True)
b = torch.ones(2, requires_grad=True)
c = MyAdder.apply(a.clone(), b)
c.sum().backward()
z = torch.tensor(1.0, requires_grad=True)
x = z.clone()
y = InplaceMul.apply(x)
a = torch.tensor(1.0, dtype=torch.double, requires_grad=True)
b = torch.tensor(1.0, dtype=torch.double, requires_grad=True)
c = torch.tensor(1.0, dtype=torch.double)
d = torch.tensor(1.0, dtype=torch.double)
MyFn2.apply(a, b)
MyFn2.apply(c, d)
base = torch.rand(10, requires_grad=True)
foo = MyFn3.apply(base, False)
test()
opt_test = torch._dynamo.optimize("eager")(test)
opt_test()
def test_tensor_subclass_intermediary_input(self):
class FooTensor(torch.Tensor):
@staticmethod
def __new__(cls, data, config, scale):
self = torch.Tensor._make_wrapper_subclass(
cls,
config[0],
strides=config[1],
storage_offset=config[2],
dtype=config[3],
layout=config[4],
requires_grad=config[5],
device=data.device,
)
self._data = data
self._config = config
self._scale = scale
return self
def __repr__(self):
return "FooTensor"
def __tensor_flatten__(self):
return ("_data",), (
self._config,
self._scale,
)
@staticmethod
def __tensor_unflatten__(tensors, metadatas, outer_size, outer_stride):
return FooTensor(tensors["_data"], metadatas[0], metadatas[1])
@classmethod
def __torch_dispatch__(cls, func, types, args, kwargs=None):
# handling clone and view is so dynamo fakefication passes, it's not
# intended to be handling user code
if func == torch.ops.aten.clone.default:
return FooTensor(
args[0]._data.clone(), args[0]._config, args[0]._scale
)
elif func == torch.ops.aten.view.default:
new_data = args[0]._data.view(*args[1:])
return FooTensor(new_data, args[0]._config, args[0]._scale)
raise NotImplementedError
class foo_autograd_fn(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
# access some data from `x`, where `x` is a tensor subclass
x2 = x._data + 1.0
# create and return a tensor subclass from within a torch.autograd.Function
x3 = FooTensor(x2, x._config, x._scale)
return x3._data
@staticmethod
def backward(ctx, g):
return g
x_ref = torch.randn(4, 4).requires_grad_(True)
x = copy.deepcopy(x_ref)
scale = torch.tensor(1.0)
# Weird that this is needed, but not having this breaks a lot of things
torch._dynamo.allow_in_graph(FooTensor)
def foo(x, scale):
config = (
x.size(),
x.stride(),
x.storage_offset(),
x.dtype,
x.layout,
x.requires_grad,
)
x = FooTensor(x, config, scale)
x = foo_autograd_fn.apply(x)
return x
y_ref = foo(x_ref, scale)
y_ref.sum().backward()
foo_opt = torch.compile(foo, backend="eager")
y = foo_opt(x, scale)
y.sum().backward()
self.assertEqual(y, y_ref)
self.assertEqual(x.grad, x_ref.grad)
def test_smuggle_symint_issue_111031(self):
from torch.autograd import Function
class Foo(Function):
@staticmethod
def forward(ctx, x):
ctx.x0 = x.size(0)
return x * 2
@staticmethod
def backward(ctx, grad_out):
return grad_out * ctx.x0
cnts = torch._dynamo.testing.CompileCounter()
@torch.compile(backend=cnts, fullgraph=True, dynamic=True)
def foo(x):
return Foo.apply(x)
foo(torch.randn(2, requires_grad=True))
self.assertEqual(cnts.frame_count, 1)
def test_needs_input_grad(self):
cnt = torch._dynamo.testing.CompileCounter()
class NeedsInputGradFunc(torch.autograd.Function):
@staticmethod
def forward(ctx, foo):
result = foo + foo
ctx.save_for_backward(result)
return result
@staticmethod
@torch.compile(backend=cnt, fullgraph=True)
def backward(ctx, grad_output):
(result,) = ctx.saved_tensors
if ctx.needs_input_grad[0]:
return grad_output * result.sin()
return None
x = torch.randn(10, requires_grad=True)
NeedsInputGradFunc.apply(x).sum().backward()
self.assertEqual(x.grad.shape, x.shape)
self.assertEqual(cnt.frame_count, 1)
self.assertEqual(cnt.op_count, 2)
def test_repeated_save_for_backward_calls(self):
from torch.autograd import Function
class Foo(Function):
@staticmethod
def forward(ctx, x, y):
ctx.save_for_backward(x)
ctx.save_for_backward(x, y)
return x * y
@staticmethod
def backward(ctx, grad_out):
x, y = ctx.saved_tensors
return grad_out * x, grad_out * y
cnts = torch._dynamo.testing.CompileCounter()
def foo(x, y):
return Foo.apply(x, y)
x_ref = torch.randn(2, requires_grad=True)
y_ref = torch.randn(2, requires_grad=True)
x_test = x_ref.clone().detach().requires_grad_()
y_test = y_ref.clone().detach().requires_grad_()
out_ref = foo(x_ref, y_ref)
out_ref.sum().backward()
out_test = torch.compile(foo, backend=cnts)(x_test, y_test)
out_test.sum().backward()
self.assertEqual(cnts.frame_count, 1)
self.assertEqual(out_ref, out_test)
self.assertEqual(x_ref.grad, x_test.grad)
self.assertEqual(y_ref.grad, y_test.grad)
def test_smuggle_tensor_and_complex_structures(self):
from torch.autograd import Function
class Foo(Function):
@staticmethod
def forward(ctx, x):
ctx.x0 = x
ctx.x1 = [1, 2, 3]
return x * 2
@staticmethod
def backward(ctx, grad_out):
x0mul = grad_out * ctx.x0
for i in ctx.x1:
x0mul = (x0mul * i) + x0mul
return x0mul
cnts = torch._dynamo.testing.CompileCounter()
@torch.compile(backend=cnts, fullgraph=True, dynamic=True)
def foo(x):
return Foo.apply(x)
foo(torch.randn(2, requires_grad=True))
self.assertEqual(cnts.frame_count, 1)
def test_default_values(self):
from torch.autograd import Function
class Foo(Function):
@staticmethod
def forward(ctx, x, alpha=0.99):
return x
@staticmethod
def backward(ctx, grad_out):
return grad_out
@torch.compile
def foo(x):
return Foo.apply(x)
# Make sure guards for default values do not crash
foo(torch.randn(2))
foo(torch.randn(2, requires_grad=True))
def test_tuple_arg(self):
cnt = torch._dynamo.testing.CompileCounter()
class TupleArgFunc(torch.autograd.Function):
@staticmethod
def forward(ctx, x, shape):
ctx.save_for_backward(torch.randn(shape))
return x + 1
@staticmethod
def backward(ctx, grad_output):
(result,) = ctx.saved_tensors
return result, None
@torch.compile(backend=cnt, fullgraph=True)
def fn():
return TupleArgFunc.apply(x, shape)
shape = (10, 10)
x = torch.randn(shape, requires_grad=True)
out = fn()
out.sum().backward()
self.assertEqual(out, x + 1)
self.assertEqual(x.grad.shape, shape)
self.assertEqual(cnt.frame_count, 1)
self.assertEqual(cnt.op_count, 2)
def test_custom_op(self):
shape = [7]
x = torch.rand(128, shape[0])
model = CustomOpModule(shape)
out = model(x)
def backend(gm, example_inputs):
return aot_module_simplified(
gm, example_inputs, fw_compiler=lambda gm, _: gm
)
opt_model = torch.compile(model, backend=backend)
opt_out = opt_model(x)
opt_out.mean().backward()
self.assertEqual(out, opt_out)
@requires_cuda
def test_triton_kernel_basic(self):
class Add(torch.autograd.Function):
@staticmethod
def forward(ctx, x, y):
ctx.save_for_backward(x, y)
output = torch.zeros_like(x)
n_elements = output.numel()
grid = lambda meta: ( # noqa: E731
triton.cdiv(n_elements, meta["BLOCK_SIZE"]),
)
add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=16)
return output
@staticmethod
def backward(ctx, grad_output):
x, y = ctx.saved_tensors
return x * grad_output, y * grad_output
@torch.compile(fullgraph=True, backend="inductor")
def f(x, y):
z = Add.apply(x, y)
return z
x = torch.randn(10, device="cuda", requires_grad=True)
y = torch.randn(10, device="cuda", requires_grad=True)
z = f(x, y)
loss = z.sum()
loss.backward()
self.assertEqual(x + y, z)
@requires_cuda
def test_triton_kernel_multiple_out(self):
class Add(torch.autograd.Function):
@staticmethod
def forward(ctx, x, y):
ctx.save_for_backward(x, y)
ctx.t1 = x
ctx.t2 = y
output = torch.zeros_like(x)
n_elements = output.numel()
grid = lambda meta: ( # noqa: E731
triton.cdiv(n_elements, meta["BLOCK_SIZE"]),
)
add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=16)
return output, x
@staticmethod
def backward(ctx, grad_output, old_x):
x, y = ctx.saved_tensors
x1 = ctx.t1
y1 = ctx.t2
return old_x * x * x1 * grad_output, y * y1 * grad_output
@torch.compile(fullgraph=True, backend="inductor")
def f(x, y):
z = Add.apply(x, y)
return z
x = torch.randn(10, device="cuda", requires_grad=True)
y = torch.randn(10, device="cuda", requires_grad=True)
z, _ = f(x, y)
loss = z.sum()
loss.backward()
self.assertEqual(x + y, z)
if __name__ == "__main__":
from torch._dynamo.test_case import run_tests
run_tests()
Reference in New Issue View Git Blame Copy Permalink