• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# mypy: ignore-errors
2
3import torch
4import torch.utils._pytree as pytree
5from torch.testing._utils import wrapper_set_seed
6from functorch.compile import compiled_function, min_cut_rematerialization_partition, nop
7from .make_fx import randomize
8import re
9
10
11class assert_raises_regex:
12    def __init__(self, exception_cls, regex):
13        self.exception_cls = exception_cls
14        self.regex = regex
15
16    def __enter__(self):
17        pass
18
19    def __exit__(self, exc_type, exc_val, traceback):
20        if exc_type == self.exception_cls:
21            msg = str(exc_val)
22            if not re.search(self.regex, msg):
23                raise AssertionError(
24                    f"Expected exception to match regex. regex: {self.regex}, exception: {msg}")
25            return True  # Squashes the exception
26        if exc_type is not None:
27            raise AssertionError(
28                f"Expected {self.exception_cls} to be raised, instead got exception {exc_type}")
29        raise AssertionError("Expected exception to be raised but none was")
30
31
32def aot_autograd_check(
33        func,
34        args,
35        kwargs,
36        dynamic,
37        assert_raises_regex_fn=assert_raises_regex,
38        assert_equals_fn=torch.testing._comparison.assert_close,
39        check_gradients=True,
40        try_check_data_specialization=False):
41    """Compares func(*args, **kwargs) in eager-mode to under AOTAutograd.
42
43    Compares outputs and (if check_gradients=True) gradients produced by
44    AOTAutograd against eager-mode PyTorch.
45
46    We assume that func(*args, **kwargs) succeeds in eager-mode PyTorch.
47
48    """
49    flat_args, args_spec = pytree.tree_flatten((args, kwargs))
50    args_is_tensor = [isinstance(arg, torch.Tensor) for arg in flat_args]
51    args = [arg for arg in flat_args if isinstance(arg, torch.Tensor)]
52
53    # We construct a new function that only accepts Tensors as inputs
54    def func_no_tensors(args):
55        reconstructed_flat_args = []
56        args = iter(args)
57        for v in flat_args:
58            if isinstance(v, torch.Tensor):
59                reconstructed_flat_args.append(next(args))
60            else:
61                reconstructed_flat_args.append(v)
62
63        c_args, c_kwargs = pytree.tree_unflatten(reconstructed_flat_args, args_spec)
64        return func(*c_args, **c_kwargs)
65
66    compiled_f = compiled_function(
67        func_no_tensors, nop, nop, dynamic=dynamic, partition_fn=min_cut_rematerialization_partition)
68
69    out = wrapper_set_seed(func_no_tensors, args)
70    if check_gradients == "auto":
71        any_tensor_requires_grad = pytree.tree_any_only(torch.Tensor, lambda x: x.requires_grad, args)
72        any_output_requires_grad = pytree.tree_any_only(torch.Tensor, lambda x: x.requires_grad, out)
73        check_gradients = any_tensor_requires_grad and any_output_requires_grad
74    if not check_gradients:
75        compiled_out = wrapper_set_seed(compiled_f, args)
76        assert_equals_fn(compiled_out, out, msg=outputs_msg)
77        return
78    _test_aot_autograd_forwards_backwards_helper(
79        func_no_tensors, compiled_f, args, assert_raises_regex_fn, assert_equals_fn,
80        try_check_data_specialization)
81
82outputs_msg = (
83    "Outputs of the operator are different in eager-mode PyTorch vs "
84    "AOTAutograd. This means the operator will have incorrect output "
85    "underneath torch.compile. This could be because the operator's "
86    "implementation not traceable or that there is a bug in AOTAutograd."
87)
88
89
90def _test_aot_autograd_forwards_backwards_helper(
91        f, compiled_f, args, assert_raises_regex_fn, assert_equals_fn,
92        try_check_data_specialization):
93    # Verify grads are equal between compiled and non-compiled versions of f.
94
95    def call_forwards_backwards(f, args):
96        flat_args = pytree.arg_tree_leaves(*args)
97        diff_args = [arg for arg in flat_args if isinstance(arg, torch.Tensor) and
98                     arg.requires_grad]
99        out = wrapper_set_seed(f, args)
100        flat_out = pytree.tree_leaves(out)
101
102        sm = 0
103        for i in flat_out:
104            if isinstance(i, torch.Tensor):
105                # We need to call .abs() because it is possible that the output of the
106                # operator is a complex Tensor and autograd will yell at autograd.grad
107                # on a complex Tensor unless we manually provide the grad_output flag.
108                sm += i.sum().abs()
109        assert isinstance(sm, torch.Tensor)
110        return out, torch.autograd.grad(sm, diff_args, allow_unused=True)
111
112    def check(args, ignore_failure=False):
113        try:
114            orig_out, orig_grad = call_forwards_backwards(f, args)
115        except Exception:
116            if ignore_failure:
117                return
118            raise
119
120        # See https://github.com/pytorch/pytorch/pull/98960#issuecomment-1505962215
121        tensor_args = [x for x in pytree.tree_flatten(args)[0] if isinstance(x, torch.Tensor)]
122        any_non_leaves = any(x.grad_fn is not None for x in tensor_args)
123        if all(x is None for x in orig_grad) and any_non_leaves:
124            with assert_raises_regex_fn(RuntimeError, 'does not require grad and does not have a grad_fn'):
125                call_forwards_backwards(compiled_f, args)
126            return
127
128        msg = (
129            "Gradients of the operator are different in eager-mode PyTorch vs "
130            "AOTAutograd. This means the operator will have incorrect gradients "
131            "underneath torch.compile. This could be because the operator's "
132            "backward is incorrectly registered or not traceable or that there "
133            "is a bug in AOTAutograd."
134        )
135
136        compiled_out, compiled_grad = call_forwards_backwards(compiled_f, args)
137        assert_equals_fn(compiled_out, orig_out, msg=outputs_msg)
138        assert_equals_fn(compiled_grad, orig_grad, msg=msg)
139
140    check(args, ignore_failure=False)
141
142    # Randomize the data and run the traced graph with it, to catch bugs
143    # where we may have baked in Tensor data into the trace.
144    # This is not guaranteed to succeed, because `f` might have preconditions
145    # on the values of the inputs, so we just ignore if this test fails.
146    if try_check_data_specialization:
147        args = randomize(args)
148        check(args, ignore_failure=True)
149