Separated out the partial gradient functionality to a separate PR

pytorch · bohnstingl · Nov 6, 2024 · Nov 6, 2024 · Nov 7, 2024 · Nov 7, 2024
commit 67d62fb4a09ca9468b747cd3e44f8d6f4bf6906d
diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py
@@ -2517,16 +2517,17 @@ def setUp(self):
 
     def _check_autograd(self, result, result_exp, autograd_param):
         grad_param = [p for p in autograd_param if p.requires_grad]
-        result_masked = [r for r in result if r.requires_grad]
-        result_exp_masked = [r for r in result_exp if r.requires_grad]
+
+        result_flatten, _ = pytree.tree_flatten(result)
+        result_exp_flatten, _ = pytree.tree_flatten(result_exp)
+        result_flatten = [r for r in result_flatten if r.requires_grad]
+        result_exp_flatten = [r for r in result_exp_flatten if r.requires_grad]
 
         # Check the result and parameter lists
-        assert len(result_masked) == len(
-            result_exp_masked
+        assert len(result_flatten) == len(
+            result_exp_flatten
         ), "The number of elements requiring gradients is different for the results and the expected results"
 
-        result_flatten, _ = pytree.tree_flatten(result_masked)
-        result_exp_flatten, _ = pytree.tree_flatten(result_exp_masked)
         grad_exp_init = [torch.ones_like(el) for el in result_exp_flatten]
         expected_grads = torch.autograd.grad(
             result_exp_flatten, grad_param, grad_exp_init
@@ -2871,7 +2872,7 @@ def second_chain_fct(scan_fct, inp, **kwargs):
             model=AssociativeScanModels.ChainFn(**kwargs),
             model_fake=AssociativeScanModels.ChainFn(**kwargs_fake),
             inputs=inp,
-            autograd_param=None if not autograd else inp,
+            autograd_param=None if not autograd else (inp,),
         )
 
     @unittest.skipIf(not SM70OrLater, "triton")
@@ -2915,7 +2916,7 @@ def second_chain_fct(scan_fct, inp, **kwargs):
             model=AssociativeScanModels.ChainFn(**kwargs),
             model_fake=AssociativeScanModels.ChainFn(**kwargs_fake),
             inputs=inp,
-            autograd_param=None if not autograd else inp,
+            autograd_param=None if not autograd else (inp,),
         )
 
     @unittest.skipIf(not SM70OrLater, "triton")
@@ -2968,7 +2969,7 @@ def second_chain_fct(scan_fct, inp, **kwargs):
             model=AssociativeScanModels.ChainFn(**kwargs),
             model_fake=AssociativeScanModels.ChainFn(**kwargs_fake),
             inputs=inp,
-            autograd_param=None if not autograd else inp,
+            autograd_param=None if not autograd else (inp,),
         )
 
     # TODO: Does not work because of the usage of vmap witin associative_scan
@@ -3068,7 +3069,7 @@ def body_fn(ind, loop_val):
             model=AssociativeScanModels.CombineFn(**kwargs),
             model_fake=AssociativeScanModels.CombineFn(**kwargs_fake),
             inputs=inp,
-            autograd_param=None if not autograd else inp,
+            autograd_param=None if not autograd else (inp,),
         )
 
     # TODO: Does not work because of the usage of vmap witin associative_scan
@@ -3135,7 +3136,7 @@ def combine_fn(x, y):
             model=AssociativeScanModels.CombineFn(**kwargs),
             model_fake=AssociativeScanModels.CombineFn(**kwargs_fake),
             inputs=inp,
-            autograd_param=None if not autograd else inp,
+            autograd_param=None if not autograd else (inp,),
         )
 
     # TODO: Does not work because of the usage of vmap witin associative_scan
@@ -3203,7 +3204,7 @@ def body(x):
             model=AssociativeScanModels.CombineFn(**kwargs),
             model_fake=AssociativeScanModels.CombineFn(**kwargs_fake),
             inputs=inp,
-            autograd_param=None if not autograd else inp,
+            autograd_param=None if not autograd else (inp,),
         )
 
     @unittest.skipIf(not SM70OrLater, "triton")
@@ -3281,107 +3282,6 @@ def test_associative_scan_binary_operator(
             autograd_param=None if not autograd else elements,
         )
 
-    @unittest.skipIf(not SM70OrLater, "triton")
-    @requires_cuda
-    @parametrize("combine_mode", ["pointwise", "generic"])
-    @parametrize("compile_mode", ["none", "compile", "compile_dynamic_shape"])
-    @parametrize("reverse", [False, True])
-    @parametrize("device", [torch.device("cpu"), torch.device("cuda")])
-    # Skipping the combination of combine_mode=pointwise and device=cpu
-    # as the current implementation of pointwise does only support CUDA device
-    @decorateIf(
-        unittest.skip,
-        lambda params: (
-            params["combine_mode"] == "pointwise"
-            and (params["device"] == torch.device("cpu") or torch.version.hip)
-        ),
-    )
-    def test_associative_scan_partial_grad(
-        self, combine_mode, compile_mode, reverse, device
-    ):
-        import random
-
-        fct_cmp = compile_mode_helper(associative_scan, compile_mode)
-
-        n_params = 6
-        autograds = []
-        autograds.append([True, True, True, True, True, True])
-        autograds.append([False, False, False, False, False, False])
-        autograds.append([False, True, False, False, False, False])
-        for _ in range(5):
-            autograds.append([bool(random.randint(0, 1)) for _ in range(n_params)])
-
-        def mul2(x, y):
-            return (*[xv * yv for xv, yv in zip(x, y)],)
-
-        for a_grads in autograds:
-            inp = tuple(
-                [
-                    torch.randn(10, 3, 2, device=device, requires_grad=a_grads[n])
-                    for n in range(n_params)
-                ]
-            )
-
-            kwargs = {
-                "dim": 0,
-                "reverse": reverse,
-                "compile_mode": compile_mode,
-                "combine_fn": mul2,
-                "combine_mode": combine_mode,
-            }
-            kwargs_fake = self._prepare_fake_kwargs(kwargs)
-            self._run_test(
-                model=AssociativeScanModels.CombineFn(**kwargs),
-                model_fake=AssociativeScanModels.CombineFn(**kwargs_fake),
-                inputs=inp,
-                autograd_param=inp,
-            )
-
-    @unittest.skipIf(not SM70OrLater, "triton")
-    @requires_cuda
-    @parametrize("combine_mode", ["pointwise", "generic"])
-    @parametrize("compile_mode", ["none", "compile", "compile_dynamic_shape"])
-    @parametrize("reverse", [False, True])
-    @parametrize("device", [torch.device("cpu"), torch.device("cuda")])
-    # Skipping the combination of combine_mode=pointwise and device=cpu
-    # as the current implementation of pointwise does only support CUDA device
-    @decorateIf(
-        unittest.skip,
-        lambda params: (
-            params["combine_mode"] == "pointwise"
-            and (params["device"] == torch.device("cpu") or torch.version.hip)
-        ),
-    )
-    def test_associative_scan_partial_grad_no_grad(
-        self, combine_mode, compile_mode, reverse, device
-    ):
-        fct_cmp = compile_mode_helper(associative_scan, compile_mode)
-
-        def mul_single_nograd(x, y):
-            xy1 = x[0] * y[0]
-            with torch.no_grad():
-                xy2 = x[1] * y[1]
-            return xy1, xy2
-
-        inp = tuple(
-            [torch.randn(10, 3, 2, device=device, requires_grad=True) for n in range(2)]
-        )
-
-        kwargs = {
-            "dim": 0,
-            "reverse": reverse,
-            "compile_mode": compile_mode,
-            "combine_fn": mul_single_nograd,
-            "combine_mode": combine_mode,
-        }
-        kwargs_fake = self._prepare_fake_kwargs(kwargs)
-        self._run_test(
-            model=AssociativeScanModels.CombineFn(**kwargs),
-            model_fake=AssociativeScanModels.CombineFn(**kwargs_fake),
-            inputs=inp,
-            autograd_param=inp[0:1],
-        )
-
     @unittest.skipIf(not SM70OrLater, "triton")
     @requires_cuda
     def test_associative_scan_sparse_tensor(self):

@@ -26,20 +26,12 @@
     track_tensor_tree,
 )
 
-from .utils import _from_fun, _maybe_reenter_make_fx, create_fw_bw_graph
+from .utils import _from_fun, create_fw_bw_graph
 
 
 aten = torch._ops.ops.aten
 
 
-def get_gradient_mask(tensor_list):
-    return [True if v is not None and v.requires_grad else False for v in tensor_list]
-
-
-def mask_gradient(grads, mask):
-    return [g for g, m in zip(grads, mask) if m]
-
-
 def wrap_combine_fn_flat(*args, combine_fn, spec, num_leaves):
     assert len(args) == 2 * num_leaves
     lhs = pytree.tree_unflatten(args[:num_leaves], spec)
@@ -96,6 +88,15 @@ def create_fw_bw_graph_combinefn(combine_fn, xs):
             fw_xs_2 = [first_slice_copy(pytree.tree_map(_from_fun, x)) for x in xs]
             outs = combine_fn(*fw_xs_1, *fw_xs_2)
 
+            # TODO: Support partial Autograd for associative_scan later
+            if pytree.tree_any(
+                lambda t: not t.requires_grad,  # type: ignore[union-attr]
+                (outs),
+            ):
+                raise RuntimeError(
+                    "gradient flags of the combine_fn output differ from the xs. Consider checking for `torch.no_grad()` statements?"
+                )
+
             fw_outputs = [pytree.tree_map(_from_fun, o) for o in outs]
             if any(not isinstance(out, torch.Tensor) for out in fw_outputs):
                 raise RuntimeError(
@@ -110,18 +111,7 @@ def create_fw_bw_graph_combinefn(combine_fn, xs):
                 (*fw_outputs,),
             )
 
-            gradient_mask_post = get_gradient_mask(outs)
-
-            def wrapper(*args):
-                grads = joint_graph(*args)
-                grads = mask_gradient(grads, gradient_mask_post * 2)
-                return (*grads,)
-
-            joint_graph = _maybe_reenter_make_fx(wrapper)(
-                *fw_outputs, *fw_xs_1, *fw_xs_2
-            )
-
-        return fw_graph, joint_graph, gradient_mask_post
+        return fw_graph, joint_graph
 
 
 class AssociativeScanOp(HigherOrderOperator):
@@ -454,20 +444,14 @@ def forward(
         ctx,
         fw_graph,
         joint_graph,
-        gradient_mask,
-        *flat_args,
+        *xs,
     ):
-        xs = flat_args
-
         ctx._joint_graph = joint_graph
-        ctx._dim = 0
-
         num_xs = len(xs)
         ctx._num_xs = num_xs
 
         scan_length = xs[0].shape[0]
         ctx._scan_length = scan_length
-        ctx._gradient_mask = gradient_mask
         ctx._mapped_joint_graph = torch.vmap(joint_graph, 0, 0)
 
         with torch._C._AutoDispatchBelowAutograd():
@@ -477,7 +461,7 @@ def forward(
         return (*outs,)
 
     @staticmethod
-    def backward(ctx, *flat_grads_unmasked):
+    def backward(ctx, *flat_grads):
         r"""
         This function computes the gradients of the scan operation.
         It does so by factorizing the components of the chainrule into
@@ -574,8 +558,6 @@ def g(x: torch.Tensor, y: torch.Tensor):
         dim = 0
         scan_length = ctx._scan_length
         num_xs = ctx._num_xs
-        gradient_mask = ctx._gradient_mask
-        num_xs_masked = sum(gradient_mask)
 
         # Extract the inputs to the forward path and outputs from the forward path
         flat_args = ctx.saved_tensors
@@ -587,27 +569,12 @@ def g(x: torch.Tensor, y: torch.Tensor):
         # vmap joint graph over scan dimension
         mapped_joint_graph = ctx._mapped_joint_graph
 
-        # TODO: This function can be shared with scan
-        def expand_grads_with_None(real_grads, mask):
-            g_list = []
-            true_cnt = 0
-            for m in mask:
-                if m:
-                    g_list.append(real_grads[true_cnt])
-                    true_cnt += 1
-                else:
-                    g_list.append(None)
-            return g_list
-
         with torch._C._AutoDispatchBelowAutograd():
             """Step 1"""
             shifted_outs = [
                 torch.concat([ones, aten.slice(o, 0, 0, -1, 1)], 0) for o in outs
             ]
 
-            # Mask the gradients for the variables that do not require gradients for partial gradient support
-            flat_grads = [fg for fg, m in zip(flat_grads_unmasked, gradient_mask) if m]
-
             # Function to compute the gradients with respect
             # *) to the inputs (xs) -> grads_xs
             # *) to the previosus outputs -> grads_hs
@@ -637,8 +604,8 @@ def compute_part_grads(flat_grad_ind):
             # Compute the grads_xs and grads_hs by collecting all the partial gradients
             grads_intermediate = compute_grad_hs_xs()
             grads_h_parts, grads_x_parts = (
-                grads_intermediate[:num_xs_masked],
-                grads_intermediate[num_xs_masked:],
+                grads_intermediate[:num_xs],
+                grads_intermediate[num_xs:],
             )
 
             # Helper variables to generate the gradient matrix
@@ -762,10 +729,7 @@ def stack_and_reduce_parts(mat_fact):
                 0,
             )
 
-            # Mask all gradients for input variables that do not require gradients
-            grads = expand_grads_with_None(grads, gradient_mask)
-
-            return *[None] * 3, *grads
+            return *[None] * 2, *grads
 
 
 @associative_scan_op.py_impl(DispatchKey.Autograd)
@@ -781,6 +745,15 @@ def associative_scan_autograd(combine_fn, xs):
         with torch._C._AutoDispatchBelowAutograd():
             return associative_scan_op(combine_fn, xs)
 
+    # TODO: Support partial Autograd for associative_scan later
+    if pytree.tree_any(
+        lambda t: not t.requires_grad,  # type: ignore[union-attr]
+        (xs),
+    ):
+        raise RuntimeError(
+            "associative_scan currently only supports Autograd if all xs require gradients."
+        )
+
     # TODO: The create_fw_bw is always invoked twice:
     # Once in the forward path and
     # once in the backward path, where it should only be invoked for the grad grad case.
@@ -794,13 +767,11 @@ def associative_scan_autograd(combine_fn, xs):
     (
         fw_graph,
         joint_graph,
-        gradient_mask,
     ) = create_fw_bw_graph_combinefn(combine_fn, xs)
 
     flat_out = ScanAutogradOp.apply(
         fw_graph,
         joint_graph,
-        gradient_mask,
         *xs,
     )
     return (*flat_out,)