From 86ea4fc9f7c9b34547e592e0d2e4b164c1a21331 Mon Sep 17 00:00:00 2001 From: Thomas Bohnstingl Date: Tue, 17 Sep 2024 22:16:54 +0200 Subject: [PATCH 01/12] WIP: wrapper function to allow partial gradients --- test/functorch/test_control_flow.py | 31 +++++---- torch/_higher_order_ops/scan.py | 101 ++++++++++++++++++++++------ 2 files changed, 98 insertions(+), 34 deletions(-) diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py index 621a640f374ff..a1257e61371dc 100644 --- a/test/functorch/test_control_flow.py +++ b/test/functorch/test_control_flow.py @@ -2902,13 +2902,14 @@ def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): dim = 1 scan_fct = compile_mode_helper(scan, compile_mode) autograds = [] - autograds.append([True, True, True, True, True, True]) - autograds.append([False, False, True, True, True, True]) - autograds.append([True, True, False, False, False, False]) - autograds.append([True, 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(6)]) + # autograds.append([True, True, True, True, True, True]) + autograds.append([True, True, True, False, False, True]) + # autograds.append([False, False, True, True, True, True]) + # autograds.append([True, True, False, False, False, False]) + # autograds.append([True, 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(6)]) for autograd in autograds: x = torch.randn(3, 10, 5, device=device, requires_grad=autograd[0]) @@ -2926,11 +2927,17 @@ def RNN(x: torch.Tensor, y: torch.Tensor): return c_new, h_new if False in autograd: - with self.assertRaisesRegex( - RuntimeError, - "scan currently only supports Autograd if all.*", - ): - result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) + # with self.assertRaisesRegex( + # RuntimeError, + # "scan currently only supports Autograd if all.*", + # ): + # result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) + result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) + result_exp = _fake_scan(RNN, h, x, dim=dim, reverse=reverse) + self.assertEqual(result, result_exp) + + if autograd: + self.check_autograd(result, result_exp, params) else: result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) result_exp = _fake_scan(RNN, h, x, dim=dim, reverse=reverse) diff --git a/torch/_higher_order_ops/scan.py b/torch/_higher_order_ops/scan.py index 972428d7273f1..c666a62874f32 100644 --- a/torch/_higher_order_ops/scan.py +++ b/torch/_higher_order_ops/scan.py @@ -112,14 +112,14 @@ def wrapper_fwd_combine_fn(*args): wrapper_fwd_combine_fn(*fw_init, *fw_xs, *fw_additional_inputs), num_init, ) - # TODO: Support this in the future - if pytree.tree_any( - lambda t: not t.requires_grad, # type: ignore[union-attr] - combine_fn(*fw_init, *fw_xs, *fw_additional_inputs), - ): - raise RuntimeError( - "scan currently only supports Autograd if all init, xs and lifted parameters require gradients." - ) + # # TODO: Support this in the future + # if pytree.tree_any( + # lambda t: not t.requires_grad, # type: ignore[union-attr] + # combine_fn(*fw_init, *fw_xs, *fw_additional_inputs), + # ): + # raise RuntimeError( + # "scan currently only supports Autograd if all init, xs and lifted parameters require gradients." + # ) fw_carry, fw_outputs = [pytree.tree_map(_from_fun, c) for c in carry], [ pytree.tree_map(_from_fun, o) for o in outs @@ -151,13 +151,27 @@ def wrapper_fwd_combine_fn(*args): (*fw_init, *fw_xs, *fw_additional_inputs), (*fw_carry, *fw_outputs[num_init:]), ) + + # Get mask to mask None gradients + g_c, g_xs = _extract_carry_and_out( + joint_graph(*fw_carry, + *fw_outputs[num_init:], + *fw_init, + *fw_xs, + *fw_additional_inputs,), num_init + ) + carry_mask = [True if g is not None and callable(g.grad_fn) else False for g in g_c] + xs_mask = [True if g is not None and callable(g.grad_fn) else False for g in g_xs[: len(g_xs) - num_additional_inputs]] + additional_inputs_mask = [True if g.requires_grad else False for g in additional_inputs] def wrapper_bwd_combine_fn(*args): - carried_g_additional_input = args[:num_additional_inputs] + carried_g_additional_input = [add_inp for add_inp, add_inp_m in zip(args[:num_additional_inputs], additional_inputs_mask) if add_inp_m] g_c, g_xs = _extract_carry_and_out( joint_graph(*args[num_additional_inputs:]), num_init ) current_g_additional_inputs = g_xs[len(g_xs) - num_additional_inputs :] + current_g_additional_inputs = [add_inp for add_inp, add_inp_m in zip(current_g_additional_inputs, additional_inputs_mask) if add_inp_m] + new_g_additional_inputs = [ carr_g + curr_g for carr_g, curr_g in zip( @@ -165,6 +179,10 @@ def wrapper_bwd_combine_fn(*args): ) ] g_xs = g_xs[: len(g_xs) - num_additional_inputs] + + g_c = [g for g, g_m in zip(g_c, carry_mask) if g_m] + g_xs = [g for g, g_m in zip(g_xs, xs_mask) if g_m] + return [*new_g_additional_inputs, *g_c, *g_xs] new_joint_graph = _maybe_reenter_make_fx(wrapper_bwd_combine_fn)( @@ -509,6 +527,10 @@ def forward( list(flat_args), num_leaves_init, num_leaves_xs ) ctx._num_leaves_additional_inputs = len(additional_inputs) + + ctx._leaves_init_reqg_mask = [True if g.requires_grad else False for g in init] + ctx._leaves_xs_reqg_mask = [True if g.requires_grad else False for g in xs] + ctx._leaves_additional_inputs_reqg_mask = [True if g.requires_grad else False for g in additional_inputs] with torch._C._AutoDispatchBelowAutograd(): carry, carries_outs = _extract_carry_and_out( @@ -606,9 +628,15 @@ def prepare_initial_gradients( joint_graph = ctx._joint_graph dim = ctx._dim reverse = ctx._reverse - num_leaves_init = ctx._num_leaves_init + num_leaves_init = ctx._num_leaves_init + leaves_init_reqg_mask = ctx._leaves_init_reqg_mask + num_leaves_init_reqg = sum(leaves_init_reqg_mask) num_leaves_xs = ctx._num_leaves_xs + leaves_xs_reqg_mask = ctx._leaves_xs_reqg_mask + num_leaves_xs_reqg = sum(leaves_xs_reqg_mask) num_leaves_additional_inputs = ctx._num_leaves_additional_inputs + leaves_additional_inputs_reqg_mask = ctx._leaves_additional_inputs_reqg_mask + num_leaves_additional_inputs_reqg = sum(leaves_additional_inputs_reqg_mask) num_leaves_ys = ctx._num_leaves_ys # The results from the forward scan are always stacked on dim 0 @@ -649,6 +677,8 @@ def prepare_initial_gradients( ) xs_bwd = [*g_ys, *carries, *xs] + # import pdb + # pdb.set_trace() g_outs = scan_op( joint_graph, [*g_additional_inputs, *g_c_T], @@ -657,14 +687,41 @@ def prepare_initial_gradients( True, additional_inputs, ) - new_g_additional_inputs = g_outs[:num_leaves_additional_inputs] + # pdb.set_trace() + new_g_additional_inputs = g_outs[:num_leaves_additional_inputs_reqg] g_init = g_outs[ - num_leaves_additional_inputs : num_leaves_additional_inputs - + num_leaves_init + num_leaves_additional_inputs_reqg : num_leaves_additional_inputs_reqg + + num_leaves_init_reqg ] - g_xs = g_outs[-num_leaves_xs:] + g_xs = g_outs[len(g_outs) -num_leaves_xs_reqg:] g_xs = prepare_final_gradients_xs(g_xs, dim, reverse) + # pdb.set_trace() + g_list = [] + ind = 0 + for g_m in leaves_additional_inputs_reqg_mask: + if g_m: + g_list.append(new_g_additional_inputs[ind]) + else: + g_list.append(None) + new_g_additional_inputs = g_list + g_list = [] + ind = 0 + for g_m in leaves_init_reqg_mask: + if g_m: + g_list.append(g_init[ind]) + else: + g_list.append(None) + g_init = g_list + g_list = [] + ind = 0 + for g_m in leaves_xs_reqg_mask: + if g_m: + g_list.append(g_xs[ind]) + else: + g_list.append(None) + g_xs = g_list + # pdb.set_trace() return *[None] * 6, *g_init, *g_xs, *new_g_additional_inputs @@ -681,14 +738,14 @@ def scan_autograd(combine_fn, init, xs, dim, reverse, additional_inputs): with torch._C._AutoDispatchBelowAutograd(): return scan_op(combine_fn, init, xs, dim, reverse, additional_inputs) - # TODO: Support this in the future - if pytree.tree_any( - lambda t: not t.requires_grad, # type: ignore[union-attr] - (init, xs, additional_inputs), - ): - raise RuntimeError( - "scan currently only supports Autograd if all init, xs and lifted parameters require gradients." - ) + # # TODO: Support this in the future + # if pytree.tree_any( + # lambda t: not t.requires_grad, # type: ignore[union-attr] + # (init, xs, additional_inputs), + # ): + # raise RuntimeError( + # "scan currently only supports Autograd if all init, xs and lifted parameters require gradients." + # ) # TODO: The create_fw_bw is always invoked twice: # Once in the forward path and From a4a5f75e761521144eed4ed6551415ebf72527f3 Mon Sep 17 00:00:00 2001 From: Thomas Bohnstingl Date: Sat, 21 Sep 2024 00:21:46 +0200 Subject: [PATCH 02/12] Working version of partial additional_input gradients --- test/functorch/test_control_flow.py | 64 ++++++++++++++------ torch/_higher_order_ops/scan.py | 94 ++++++++++++++--------------- 2 files changed, 88 insertions(+), 70 deletions(-) diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py index a1257e61371dc..866ad6bf73495 100644 --- a/test/functorch/test_control_flow.py +++ b/test/functorch/test_control_flow.py @@ -2890,20 +2890,50 @@ def RNN(x: torch.Tensor, y: torch.Tensor): grads = grads[:2] self.assertEqual(grads, expected_grads) self.assertEqual(add_input_grads, expected_add_input_grads) + + @unittest.skipIf(not SM70OrLater, "triton") + @requires_cuda + @parametrize("reverse", [False, True]) + @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) + @parametrize("autograd", [False, True]) + def test_scan_closure_RNN_parameters_as_inputs(self, reverse, compile_mode, device, autograd): + x = torch.randn(3, 5, 10, device=device, requires_grad=autograd) + h = torch.randn(3, 7, device=device, requires_grad=autograd) + W_ih = torch.randn(5, 7, device=device, requires_grad=autograd) + b_ih = torch.randn(7, device=device, requires_grad=autograd) + W_hh = torch.randn(7, 7, device=device, requires_grad=autograd) + b_hh = torch.randn(7, device=device, requires_grad=autograd) + + def RNN(x: torch.Tensor, y: torch.Tensor): + c_new = y[0] @ W_ih[1] + b_ih[2] + h_new = torch.tanh(c_new + x @ W_hh[3] + b_hh[4]) + return c_new, h_new + + with self.assertRaisesRegex( + RuntimeError, + "All xs leaves must at least have.*", + ): + result = scan(RNN, h, [x, W_ih, b_ih, W_hh, b_hh], dim=2, reverse=reverse) @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda @parametrize("reverse", [False, True]) + # @parametrize("reverse", [False]) @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) - @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) + # @parametrize("compile_mode", ["eager"]) + # @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) + @parametrize("device", [torch.device("cpu")]) def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): import random dim = 1 scan_fct = compile_mode_helper(scan, compile_mode) autograds = [] + autograds.append([True, True, True, False, True, False]) + autograds.append([True, True, True, False, False, False]) + autograds.append([True, True, False, False, False, False]) + # autograds.append([True, True, True, True, True, True]) - autograds.append([True, True, True, False, False, True]) # autograds.append([False, False, True, True, True, True]) # autograds.append([True, True, False, False, False, False]) # autograds.append([True, False, False, False, False, False]) @@ -2926,25 +2956,19 @@ def RNN(x: torch.Tensor, y: torch.Tensor): h_new = torch.tanh(c_new + x @ W_hh + b_hh) return c_new, h_new - if False in autograd: - # with self.assertRaisesRegex( - # RuntimeError, - # "scan currently only supports Autograd if all.*", - # ): - # result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) - result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) - result_exp = _fake_scan(RNN, h, x, dim=dim, reverse=reverse) - self.assertEqual(result, result_exp) - - if autograd: - self.check_autograd(result, result_exp, params) - else: - result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) - result_exp = _fake_scan(RNN, h, x, dim=dim, reverse=reverse) - self.assertEqual(result, result_exp) + # if False in autograd: + # with self.assertRaisesRegex( + # RuntimeError, + # "scan currently only supports Autograd if all.*", + # ): + # result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) + # else: + result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) + result_exp = _fake_scan(RNN, h, x, dim=dim, reverse=reverse) + self.assertEqual(result, result_exp) - if autograd: - self.check_autograd(result, result_exp, params) + if autograd: + self.check_autograd(result, result_exp, params) @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda diff --git a/torch/_higher_order_ops/scan.py b/torch/_higher_order_ops/scan.py index c666a62874f32..db8ecfa812901 100644 --- a/torch/_higher_order_ops/scan.py +++ b/torch/_higher_order_ops/scan.py @@ -65,6 +65,9 @@ def shift_source_dim_to_target_dim(t, from_dim: int, to_dim: int): dims.insert(to_dim, from_dim) return t.permute(*dims) +def get_gradient_mask(tensor_list): + return [True if v.requires_grad else False for v in tensor_list] + # Internal functions for scan.py def wrap_combine_fn_flat( @@ -160,18 +163,23 @@ def wrapper_fwd_combine_fn(*args): *fw_xs, *fw_additional_inputs,), num_init ) - carry_mask = [True if g is not None and callable(g.grad_fn) else False for g in g_c] - xs_mask = [True if g is not None and callable(g.grad_fn) else False for g in g_xs[: len(g_xs) - num_additional_inputs]] - additional_inputs_mask = [True if g.requires_grad else False for g in additional_inputs] + carry_mask = get_gradient_mask(g_c) + xs_mask = get_gradient_mask(g_xs[: len(g_xs) - num_additional_inputs]) + additional_inputs_mask = get_gradient_mask(additional_inputs) + + bw_additional_inputs = [add_inp for add_inp, add_inp_m in zip(bw_additional_inputs, additional_inputs_mask) if add_inp_m] + num_additional_inputs_masked = len(bw_additional_inputs) def wrapper_bwd_combine_fn(*args): - carried_g_additional_input = [add_inp for add_inp, add_inp_m in zip(args[:num_additional_inputs], additional_inputs_mask) if add_inp_m] + carried_g_additional_input = args[:num_additional_inputs_masked] + # carried_g_additional_input = [add_inp for add_inp, add_inp_m in zip(carried_g_additional_input, additional_inputs_mask) if add_inp_m] + g_c, g_xs = _extract_carry_and_out( - joint_graph(*args[num_additional_inputs:]), num_init + joint_graph(*args[num_additional_inputs_masked:]), num_init ) current_g_additional_inputs = g_xs[len(g_xs) - num_additional_inputs :] current_g_additional_inputs = [add_inp for add_inp, add_inp_m in zip(current_g_additional_inputs, additional_inputs_mask) if add_inp_m] - + new_g_additional_inputs = [ carr_g + curr_g for carr_g, curr_g in zip( @@ -370,11 +378,7 @@ def _scan(init, xs): # Compute dummy shapes for the pre-allocation num_init_leaves = len(init) dummy_carry, dummy_out = _extract_carry_and_out( - operator( - *carry, - *[first_slice_copy(elem, dim) for elem in xs], - *additional_inputs, - ), + operator(*carry, *[first_slice_copy(elem, dim) for elem in xs], *additional_inputs,), num_init_leaves, ) @@ -526,11 +530,11 @@ def forward( init, xs, additional_inputs = ScanAutogradOp.extract_init_xs_additional_inputs( list(flat_args), num_leaves_init, num_leaves_xs ) - ctx._num_leaves_additional_inputs = len(additional_inputs) - ctx._leaves_init_reqg_mask = [True if g.requires_grad else False for g in init] - ctx._leaves_xs_reqg_mask = [True if g.requires_grad else False for g in xs] - ctx._leaves_additional_inputs_reqg_mask = [True if g.requires_grad else False for g in additional_inputs] + ctx._carry_mask = get_gradient_mask(init) + ctx._xs_mask = get_gradient_mask(xs) + ctx._additional_inputs_mask = get_gradient_mask(additional_inputs) + ctx._num_leaves_additional_inputs = len(additional_inputs) with torch._C._AutoDispatchBelowAutograd(): carry, carries_outs = _extract_carry_and_out( @@ -624,20 +628,32 @@ def prepare_initial_gradients( # The initial gradients for the additional_inputs are all zeros g_additional_inputs = [torch.zeros_like(ai) for ai in additional_inputs] return g_c_T, g_ys, g_additional_inputs + + def fill_grads_with_mask(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 joint_graph = ctx._joint_graph dim = ctx._dim reverse = ctx._reverse num_leaves_init = ctx._num_leaves_init - leaves_init_reqg_mask = ctx._leaves_init_reqg_mask - num_leaves_init_reqg = sum(leaves_init_reqg_mask) num_leaves_xs = ctx._num_leaves_xs - leaves_xs_reqg_mask = ctx._leaves_xs_reqg_mask - num_leaves_xs_reqg = sum(leaves_xs_reqg_mask) num_leaves_additional_inputs = ctx._num_leaves_additional_inputs - leaves_additional_inputs_reqg_mask = ctx._leaves_additional_inputs_reqg_mask - num_leaves_additional_inputs_reqg = sum(leaves_additional_inputs_reqg_mask) num_leaves_ys = ctx._num_leaves_ys + + carry_mask = ctx._carry_mask + num_carry_mask = sum(carry_mask) + xs_mask = ctx._xs_mask + num_xs_mask = sum(xs_mask) + additional_inputs_mask = ctx._additional_inputs_mask + num_additional_inputs_mask = sum(additional_inputs_mask) # The results from the forward scan are always stacked on dim 0 # The gradients though need to be provided with the correct scan dimension dim @@ -668,6 +684,8 @@ def prepare_initial_gradients( num_leaves_ys, bwd_scan_dim, ) + + g_additional_inputs = [add_inp for add_inp, add_inp_m in zip(g_additional_inputs, additional_inputs_mask) if add_inp_m] # Prepare the inputs for the backward scan. # This involves flipping the input xs if needed as well as @@ -678,7 +696,6 @@ def prepare_initial_gradients( xs_bwd = [*g_ys, *carries, *xs] # import pdb - # pdb.set_trace() g_outs = scan_op( joint_graph, [*g_additional_inputs, *g_c_T], @@ -688,39 +705,16 @@ def prepare_initial_gradients( additional_inputs, ) # pdb.set_trace() - new_g_additional_inputs = g_outs[:num_leaves_additional_inputs_reqg] + new_g_additional_inputs = g_outs[:num_additional_inputs_mask] g_init = g_outs[ - num_leaves_additional_inputs_reqg : num_leaves_additional_inputs_reqg - + num_leaves_init_reqg + num_additional_inputs_mask : num_additional_inputs_mask + + num_carry_mask ] - g_xs = g_outs[len(g_outs) -num_leaves_xs_reqg:] + g_xs = g_outs[len(g_outs) - num_xs_mask:] g_xs = prepare_final_gradients_xs(g_xs, dim, reverse) # pdb.set_trace() - g_list = [] - ind = 0 - for g_m in leaves_additional_inputs_reqg_mask: - if g_m: - g_list.append(new_g_additional_inputs[ind]) - else: - g_list.append(None) - new_g_additional_inputs = g_list - g_list = [] - ind = 0 - for g_m in leaves_init_reqg_mask: - if g_m: - g_list.append(g_init[ind]) - else: - g_list.append(None) - g_init = g_list - g_list = [] - ind = 0 - for g_m in leaves_xs_reqg_mask: - if g_m: - g_list.append(g_xs[ind]) - else: - g_list.append(None) - g_xs = g_list + new_g_additional_inputs = fill_grads_with_mask(new_g_additional_inputs, additional_inputs_mask) # pdb.set_trace() return *[None] * 6, *g_init, *g_xs, *new_g_additional_inputs From ee3e134f65f8bf72f0d2bb144e2524452ec98651 Mon Sep 17 00:00:00 2001 From: Thomas Bohnstingl Date: Sat, 21 Sep 2024 00:49:41 +0200 Subject: [PATCH 03/12] Working version of partial xs --- test/functorch/test_control_flow.py | 34 ++++++++++++------------ torch/_higher_order_ops/scan.py | 40 ++++++++++++++++------------- 2 files changed, 40 insertions(+), 34 deletions(-) diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py index 866ad6bf73495..f5a1b106e6cad 100644 --- a/test/functorch/test_control_flow.py +++ b/test/functorch/test_control_flow.py @@ -2917,10 +2917,10 @@ def RNN(x: torch.Tensor, y: torch.Tensor): @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda - @parametrize("reverse", [False, True]) - # @parametrize("reverse", [False]) - @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) - # @parametrize("compile_mode", ["eager"]) + # @parametrize("reverse", [False, True]) + @parametrize("reverse", [False]) + # @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) + @parametrize("compile_mode", ["eager"]) # @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) @parametrize("device", [torch.device("cpu")]) def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): @@ -2929,9 +2929,10 @@ def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): dim = 1 scan_fct = compile_mode_helper(scan, compile_mode) autograds = [] - autograds.append([True, True, True, False, True, False]) - autograds.append([True, True, True, False, False, False]) - autograds.append([True, True, False, False, False, False]) + autograds.append([False, True, True, True, True, True, True]) + # autograds.append([True, True, True, False, True, False]) + # autograds.append([True, True, True, False, False, False]) + # autograds.append([True, True, False, False, False, False]) # autograds.append([True, True, True, True, True, True]) # autograds.append([False, False, True, True, True, True]) @@ -2943,16 +2944,17 @@ def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): for autograd in autograds: x = torch.randn(3, 10, 5, device=device, requires_grad=autograd[0]) - h = torch.randn(3, 7, device=device, requires_grad=autograd[1]) - W_ih = torch.randn(5, 7, device=device, requires_grad=autograd[2]) - b_ih = torch.randn(7, device=device, requires_grad=autograd[3]) - W_hh = torch.randn(7, 7, device=device, requires_grad=autograd[4]) - b_hh = torch.randn(7, device=device, requires_grad=autograd[5]) + x1 = torch.randn(3, 10, 5, device=device, requires_grad=autograd[1]) + h = torch.randn(3, 7, device=device, requires_grad=autograd[2]) + W_ih = torch.randn(5, 7, device=device, requires_grad=autograd[3]) + b_ih = torch.randn(7, device=device, requires_grad=autograd[4]) + W_hh = torch.randn(7, 7, device=device, requires_grad=autograd[5]) + b_hh = torch.randn(7, device=device, requires_grad=autograd[6]) - params = [p for p, a in zip([x, h, W_ih, b_ih, W_hh, b_hh], autograd) if a] + params = [p for p, a in zip([x, x1, h, W_ih, b_ih, W_hh, b_hh], autograd) if a] def RNN(x: torch.Tensor, y: torch.Tensor): - c_new = y @ W_ih + b_ih + c_new = y[0] @ W_ih + b_ih + y[1] @ W_ih h_new = torch.tanh(c_new + x @ W_hh + b_hh) return c_new, h_new @@ -2963,8 +2965,8 @@ def RNN(x: torch.Tensor, y: torch.Tensor): # ): # result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) # else: - result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) - result_exp = _fake_scan(RNN, h, x, dim=dim, reverse=reverse) + result = scan_fct(RNN, h, (x, x1), dim=dim, reverse=reverse) + result_exp = _fake_scan(RNN, h, (x, x1), dim=dim, reverse=reverse) self.assertEqual(result, result_exp) if autograd: diff --git a/torch/_higher_order_ops/scan.py b/torch/_higher_order_ops/scan.py index db8ecfa812901..2054812a92b67 100644 --- a/torch/_higher_order_ops/scan.py +++ b/torch/_higher_order_ops/scan.py @@ -66,7 +66,7 @@ def shift_source_dim_to_target_dim(t, from_dim: int, to_dim: int): return t.permute(*dims) def get_gradient_mask(tensor_list): - return [True if v.requires_grad else False for v in tensor_list] + return [True if v is not None and v.requires_grad else False for v in tensor_list] # Internal functions for scan.py @@ -187,9 +187,8 @@ def wrapper_bwd_combine_fn(*args): ) ] g_xs = g_xs[: len(g_xs) - num_additional_inputs] - - g_c = [g for g, g_m in zip(g_c, carry_mask) if g_m] g_xs = [g for g, g_m in zip(g_xs, xs_mask) if g_m] + g_c = [g for g, g_m in zip(g_c, carry_mask) if g_m] return [*new_g_additional_inputs, *g_c, *g_xs] @@ -387,19 +386,22 @@ def _scan(init, xs): # idxs -> Index matrix for scatter_ # out: (num_elems, M, N, ...) # idx: (1, M, N) - outs, idxs = zip( - *[ - [ - torch.zeros( - [num_elems] + list(e.size()), - dtype=e.dtype, - device=e.device, - ), - torch.ones_like(e, dtype=torch.int64).unsqueeze(0), + if len(dummy_out) > 0: + outs, idxs = zip( + *[ + [ + torch.zeros( + [num_elems] + list(e.size()), + dtype=e.dtype, + device=e.device, + ), + torch.ones_like(e, dtype=torch.int64).unsqueeze(0), + ] + for i, e in enumerate(dummy_out) ] - for i, e in enumerate(dummy_out) - ] - ) + ) + else: + outs, idxs = [], [] def store_out_in_outs(out, ind): # Store the intermediate out in the outs matrix @@ -695,7 +697,7 @@ def fill_grads_with_mask(real_grads, mask): ) xs_bwd = [*g_ys, *carries, *xs] - # import pdb + import pdb g_outs = scan_op( joint_graph, [*g_additional_inputs, *g_c_T], @@ -713,9 +715,11 @@ def fill_grads_with_mask(real_grads, mask): g_xs = g_outs[len(g_outs) - num_xs_mask:] g_xs = prepare_final_gradients_xs(g_xs, dim, reverse) - # pdb.set_trace() + new_g_additional_inputs = fill_grads_with_mask(new_g_additional_inputs, additional_inputs_mask) - # pdb.set_trace() + pdb.set_trace() + g_xs = fill_grads_with_mask(g_xs, xs_mask) + pdb.set_trace() return *[None] * 6, *g_init, *g_xs, *new_g_additional_inputs From f6bdfe70fe22ac2dd8a3497e5caaeac99a2085a2 Mon Sep 17 00:00:00 2001 From: Thomas Bohnstingl Date: Sat, 21 Sep 2024 11:57:15 +0200 Subject: [PATCH 04/12] Partial init autograd working --- test/functorch/test_control_flow.py | 44 +++++++++++++++------------ torch/_higher_order_ops/scan.py | 46 +++++++++++++++++++++-------- 2 files changed, 58 insertions(+), 32 deletions(-) diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py index f5a1b106e6cad..4d0ef0f73286f 100644 --- a/test/functorch/test_control_flow.py +++ b/test/functorch/test_control_flow.py @@ -2929,10 +2929,12 @@ def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): dim = 1 scan_fct = compile_mode_helper(scan, compile_mode) autograds = [] + autograds.append([True, True, False, True, True, True, True]) autograds.append([False, True, True, True, True, True, True]) - # autograds.append([True, True, True, False, True, False]) - # autograds.append([True, True, True, False, False, False]) - # autograds.append([True, True, False, False, False, False]) + autograds.append([False, False, True, True, True, True, True]) + autograds.append([True, True, True, True, False, True, False]) + autograds.append([True, True, True, True, False, False, False]) + autograds.append([True, True, True, False, False, False, False]) # autograds.append([True, True, True, True, True, True]) # autograds.append([False, False, True, True, True, True]) @@ -2974,10 +2976,14 @@ def RNN(x: torch.Tensor, y: torch.Tensor): @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda - @parametrize("reverse", [False, True]) - @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) - @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) - @parametrize("autograd", [False, True]) + # @parametrize("reverse", [False, True]) + @parametrize("reverse", [False]) + # @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) + @parametrize("compile_mode", ["eager"]) + # @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) + @parametrize("device", [torch.device("cpu")]) + # @parametrize("autograd", [False, True]) + @parametrize("autograd", [True]) def test_scan_closure_combine_fn_with_no_grad( self, reverse, compile_mode, device, autograd ): @@ -2996,19 +3002,19 @@ def RNN(x: torch.Tensor, y: torch.Tensor): h_new = torch.tanh(c_new + x @ W_hh + b_hh) return c_new, h_new - if autograd: - with self.assertRaisesRegex( - RuntimeError, - "scan currently only supports Autograd if all.*", - ): - result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) - else: - result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) - result_exp = _fake_scan(RNN, h, x, dim=dim, reverse=reverse) - self.assertEqual(result, result_exp) + # if autograd: + # with self.assertRaisesRegex( + # RuntimeError, + # "scan currently only supports Autograd if all.*", + # ): + # result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) + # else: + result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) + result_exp = _fake_scan(RNN, h, x, dim=dim, reverse=reverse) + self.assertEqual(result, result_exp) - if autograd: - self.check_autograd(result, result_exp, (x, h, W_ih, b_ih, W_hh, b_hh)) + if autograd: + self.check_autograd(result[0], result_exp[0], (x, W_ih, b_ih)) @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda diff --git a/torch/_higher_order_ops/scan.py b/torch/_higher_order_ops/scan.py index 2054812a92b67..6bf183b305f5b 100644 --- a/torch/_higher_order_ops/scan.py +++ b/torch/_higher_order_ops/scan.py @@ -147,6 +147,15 @@ def wrapper_fwd_combine_fn(*args): fw_graph = _maybe_reenter_make_fx(wrapper_fwd_combine_fn)( *fw_init, *fw_xs, *fw_additional_inputs ) + + # Get mask to mask None gradients + carry_mask = get_gradient_mask(fw_init) + xs_mask = get_gradient_mask(fw_xs) + # additional_inputs_mask = get_gradient_mask(fw_additional_inputs) + + # Enforce that the gradients of the inits are traced + for el in fw_init: + el.requires_grad = True _, joint_graph = create_fw_bw_graph( combine_fn, @@ -163,9 +172,7 @@ def wrapper_fwd_combine_fn(*args): *fw_xs, *fw_additional_inputs,), num_init ) - carry_mask = get_gradient_mask(g_c) - xs_mask = get_gradient_mask(g_xs[: len(g_xs) - num_additional_inputs]) - additional_inputs_mask = get_gradient_mask(additional_inputs) + additional_inputs_mask = get_gradient_mask(g_xs[len(g_xs) - num_additional_inputs :]) bw_additional_inputs = [add_inp for add_inp, add_inp_m in zip(bw_additional_inputs, additional_inputs_mask) if add_inp_m] num_additional_inputs_masked = len(bw_additional_inputs) @@ -188,7 +195,7 @@ def wrapper_bwd_combine_fn(*args): ] g_xs = g_xs[: len(g_xs) - num_additional_inputs] g_xs = [g for g, g_m in zip(g_xs, xs_mask) if g_m] - g_c = [g for g, g_m in zip(g_c, carry_mask) if g_m] + # g_c = [g if g_m else torch.ones_like(gi) for g, g_m, gi in zip(g_c, carry_mask, args[num_additional_inputs_masked:num_additional_inputs_masked+num_init]) ] return [*new_g_additional_inputs, *g_c, *g_xs] @@ -200,7 +207,7 @@ def wrapper_bwd_combine_fn(*args): *fw_xs, *fw_additional_inputs, ) - return fw_graph, new_joint_graph + return fw_graph, new_joint_graph, carry_mask, xs_mask, additional_inputs_mask def scan( @@ -522,6 +529,9 @@ def forward( reverse, num_leaves_init, num_leaves_xs, + carry_mask, + xs_mask, + additional_inputs_mask, *flat_args, ): ctx._joint_graph = joint_graph @@ -532,11 +542,11 @@ def forward( init, xs, additional_inputs = ScanAutogradOp.extract_init_xs_additional_inputs( list(flat_args), num_leaves_init, num_leaves_xs ) - - ctx._carry_mask = get_gradient_mask(init) - ctx._xs_mask = get_gradient_mask(xs) - ctx._additional_inputs_mask = get_gradient_mask(additional_inputs) ctx._num_leaves_additional_inputs = len(additional_inputs) + + ctx._carry_mask = carry_mask + ctx._xs_mask = xs_mask + ctx._additional_inputs_mask = additional_inputs_mask with torch._C._AutoDispatchBelowAutograd(): carry, carries_outs = _extract_carry_and_out( @@ -710,17 +720,21 @@ def fill_grads_with_mask(real_grads, mask): new_g_additional_inputs = g_outs[:num_additional_inputs_mask] g_init = g_outs[ num_additional_inputs_mask : num_additional_inputs_mask - + num_carry_mask + + num_leaves_init ] + # pdb.set_trace() + g_init = [g for g, g_m in zip(g_init, carry_mask) if g_m] + # pdb.set_trace() g_xs = g_outs[len(g_outs) - num_xs_mask:] g_xs = prepare_final_gradients_xs(g_xs, dim, reverse) new_g_additional_inputs = fill_grads_with_mask(new_g_additional_inputs, additional_inputs_mask) - pdb.set_trace() + # pdb.set_trace() g_xs = fill_grads_with_mask(g_xs, xs_mask) - pdb.set_trace() - return *[None] * 6, *g_init, *g_xs, *new_g_additional_inputs + g_init = fill_grads_with_mask(g_init, carry_mask) + # pdb.set_trace() + return *[None] * 9, *g_init, *g_xs, *new_g_additional_inputs @scan_op.py_impl(DispatchKey.Autograd) @@ -764,6 +778,9 @@ def scan_autograd(combine_fn, init, xs, dim, reverse, additional_inputs): ( fw_graph, joint_graph, + carry_mask, + xs_mask, + additional_inputs_mask, ) = create_fw_bw_graph_combinefn(combine_fn, init, xs, dim, additional_inputs) flat_out = ScanAutogradOp.apply( @@ -773,6 +790,9 @@ def scan_autograd(combine_fn, init, xs, dim, reverse, additional_inputs): reverse, num_leaves_init, num_leaves_xs, + carry_mask, + xs_mask, + additional_inputs_mask, *(init + xs + additional_inputs), ) return *flat_out[:num_leaves_init], *flat_out[num_leaves_init:] From 3ffbbad1b8337920043039c804d7e33b2163c76f Mon Sep 17 00:00:00 2001 From: Thomas Bohnstingl Date: Sat, 21 Sep 2024 13:06:06 +0200 Subject: [PATCH 05/12] with torch.no_grad() working --- test/functorch/test_control_flow.py | 130 +++++++++++++++++----------- torch/_higher_order_ops/scan.py | 16 +++- 2 files changed, 92 insertions(+), 54 deletions(-) diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py index 4d0ef0f73286f..0389fb259aa4a 100644 --- a/test/functorch/test_control_flow.py +++ b/test/functorch/test_control_flow.py @@ -2925,65 +2925,79 @@ def RNN(x: torch.Tensor, y: torch.Tensor): @parametrize("device", [torch.device("cpu")]) def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): import random + torch._dynamo.reset() + torch._dynamo.config.cache_size_limit = 100 dim = 1 scan_fct = compile_mode_helper(scan, compile_mode) + + # The first two booleans are the xs + # The second two are the inits + # The last four are the additional_inputs autograds = [] - autograds.append([True, True, False, True, True, True, True]) - autograds.append([False, True, True, True, True, True, True]) - autograds.append([False, False, True, True, True, True, True]) - autograds.append([True, True, True, True, False, True, False]) - autograds.append([True, True, True, True, False, False, False]) - autograds.append([True, True, True, False, False, False, False]) - # autograds.append([True, True, True, True, True, True]) - # autograds.append([False, False, True, True, True, True]) - # autograds.append([True, True, False, False, False, False]) - # autograds.append([True, 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(6)]) + # # Basic test + # autograds.append([True, True, True, True, True, True, True, True]) + + # # xs tests + # autograds.append([True, False, True, True, True, True, True, True]) + # autograds.append([False, False, True, True, True, True, True, True]) + + # # init tests + # autograds.append([True, True, False, True, True, True, True, True]) + # autograds.append([True, True, False, False, True, True, True, True]) + + # # additional input tests + # autograds.append([True, True, True, True, False, True, False, True]) + # autograds.append([True, True, True, True, False, False, False, False]) + + # # Complex cases + # autograds.append([True, False, False, False, False, False, False, True]) + autograds.append([False, False, False, True, False, False, False, True]) + + # # Random tests + # for _ in range(10): + # autograds.append([bool(random.randint(0, 1)) for _ in range(8)]) for autograd in autograds: + print(autograd) + x = torch.randn(3, 10, 5, device=device, requires_grad=autograd[0]) x1 = torch.randn(3, 10, 5, device=device, requires_grad=autograd[1]) h = torch.randn(3, 7, device=device, requires_grad=autograd[2]) - W_ih = torch.randn(5, 7, device=device, requires_grad=autograd[3]) - b_ih = torch.randn(7, device=device, requires_grad=autograd[4]) - W_hh = torch.randn(7, 7, device=device, requires_grad=autograd[5]) - b_hh = torch.randn(7, device=device, requires_grad=autograd[6]) + h_1 = torch.randn(3, 7, device=device, requires_grad=autograd[3]) + W_ih = torch.randn(5, 7, device=device, requires_grad=autograd[4]) + b_ih = torch.randn(7, device=device, requires_grad=autograd[5]) + W_hh = torch.randn(7, 7, device=device, requires_grad=autograd[6]) + b_hh = torch.randn(7, device=device, requires_grad=autograd[7]) - params = [p for p, a in zip([x, x1, h, W_ih, b_ih, W_hh, b_hh], autograd) if a] + params = [p for p, a in zip([x, x1, h, h_1, W_ih, b_ih, W_hh, b_hh], autograd) if a] def RNN(x: torch.Tensor, y: torch.Tensor): - c_new = y[0] @ W_ih + b_ih + y[1] @ W_ih - h_new = torch.tanh(c_new + x @ W_hh + b_hh) - return c_new, h_new - - # if False in autograd: - # with self.assertRaisesRegex( - # RuntimeError, - # "scan currently only supports Autograd if all.*", - # ): - # result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) - # else: - result = scan_fct(RNN, h, (x, x1), dim=dim, reverse=reverse) - result_exp = _fake_scan(RNN, h, (x, x1), dim=dim, reverse=reverse) + c_new_1 = y[0] @ W_ih + y[1] @ W_ih + c_new_2 = b_ih + y[1] @ W_ih + x[1] + h_new = torch.tanh(c_new_1 + x[0] @ W_hh + b_hh) + return (c_new_1, c_new_2), h_new + + result = scan_fct(RNN, (h, h_1), (x, x1), dim=dim, reverse=reverse) + result_exp = _fake_scan(RNN, (h, h_1), (x, x1), dim=dim, reverse=reverse) self.assertEqual(result, result_exp) if autograd: - self.check_autograd(result, result_exp, params) + result_flat = pytree.tree_leaves(result) + result_exp_flat = pytree.tree_leaves(result_exp) + self.check_autograd([r for r in result_flat if r.requires_grad], [r for r in result_exp_flat if r.requires_grad], params) @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda - # @parametrize("reverse", [False, True]) - @parametrize("reverse", [False]) - # @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) - @parametrize("compile_mode", ["eager"]) - # @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) - @parametrize("device", [torch.device("cpu")]) - # @parametrize("autograd", [False, True]) - @parametrize("autograd", [True]) + @parametrize("reverse", [False, True]) + # @parametrize("reverse", [False]) + @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) + # @parametrize("compile_mode", ["eager"]) + @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) + # @parametrize("device", [torch.device("cpu")]) + @parametrize("autograd", [False, True]) + # @parametrize("autograd", [True]) def test_scan_closure_combine_fn_with_no_grad( self, reverse, compile_mode, device, autograd ): @@ -2996,25 +3010,41 @@ def test_scan_closure_combine_fn_with_no_grad( W_hh = torch.randn(7, 7, device=device, requires_grad=autograd) b_hh = torch.randn(7, device=device, requires_grad=autograd) - def RNN(x: torch.Tensor, y: torch.Tensor): + def RNN_1(x: torch.Tensor, y: torch.Tensor): c_new = y @ W_ih + b_ih with torch.no_grad(): h_new = torch.tanh(c_new + x @ W_hh + b_hh) return c_new, h_new + + def RNN_2(x: torch.Tensor, y: torch.Tensor): + h_new = torch.tanh(x + x @ W_hh + b_hh) + with torch.no_grad(): + c_new = y @ W_ih + b_ih + return c_new, h_new + + def RNN_3(x: torch.Tensor, y: torch.Tensor): + with torch.no_grad(): + c_new = y @ W_ih + b_ih + h_new = torch.tanh(c_new + x @ W_hh + b_hh) + return c_new, h_new - # if autograd: - # with self.assertRaisesRegex( - # RuntimeError, - # "scan currently only supports Autograd if all.*", - # ): - # result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) - # else: - result = scan_fct(RNN, h, x, dim=dim, reverse=reverse) - result_exp = _fake_scan(RNN, h, x, dim=dim, reverse=reverse) + result = scan_fct(RNN_1, h, x, dim=dim, reverse=reverse) + result_exp = _fake_scan(RNN_1, h, x, dim=dim, reverse=reverse) self.assertEqual(result, result_exp) if autograd: self.check_autograd(result[0], result_exp[0], (x, W_ih, b_ih)) + + result = scan_fct(RNN_2, h, x, dim=dim, reverse=reverse) + result_exp = _fake_scan(RNN_2, h, x, dim=dim, reverse=reverse) + self.assertEqual(result, result_exp) + + if autograd: + self.check_autograd(result[1], result_exp[1], (h, W_hh, b_hh)) + + result = scan_fct(RNN_3, h, x, dim=dim, reverse=reverse) + result_exp = _fake_scan(RNN_3, h, x, dim=dim, reverse=reverse) + self.assertEqual(result, result_exp) @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda diff --git a/torch/_higher_order_ops/scan.py b/torch/_higher_order_ops/scan.py index 6bf183b305f5b..c8602eacdd316 100644 --- a/torch/_higher_order_ops/scan.py +++ b/torch/_higher_order_ops/scan.py @@ -148,14 +148,20 @@ def wrapper_fwd_combine_fn(*args): *fw_init, *fw_xs, *fw_additional_inputs ) - # Get mask to mask None gradients - carry_mask = get_gradient_mask(fw_init) - xs_mask = get_gradient_mask(fw_xs) + # # Get mask to mask None gradients + # carry_mask = get_gradient_mask(fw_init) + # xs_mask = get_gradient_mask(fw_xs) # additional_inputs_mask = get_gradient_mask(fw_additional_inputs) # Enforce that the gradients of the inits are traced for el in fw_init: el.requires_grad = True + for el in fw_carry: + el.requires_grad = True + + # fw_xs = [pytree.tree_map(_from_fun, x) for x in fw_xs] + # for el in fw_xs: + # el.requires_grad = True _, joint_graph = create_fw_bw_graph( combine_fn, @@ -172,6 +178,8 @@ def wrapper_fwd_combine_fn(*args): *fw_xs, *fw_additional_inputs,), num_init ) + carry_mask = get_gradient_mask(g_c) + xs_mask = get_gradient_mask(g_xs[: len(g_xs) - num_additional_inputs]) additional_inputs_mask = get_gradient_mask(g_xs[len(g_xs) - num_additional_inputs :]) bw_additional_inputs = [add_inp for add_inp, add_inp_m in zip(bw_additional_inputs, additional_inputs_mask) if add_inp_m] @@ -195,7 +203,7 @@ def wrapper_bwd_combine_fn(*args): ] g_xs = g_xs[: len(g_xs) - num_additional_inputs] g_xs = [g for g, g_m in zip(g_xs, xs_mask) if g_m] - # g_c = [g if g_m else torch.ones_like(gi) for g, g_m, gi in zip(g_c, carry_mask, args[num_additional_inputs_masked:num_additional_inputs_masked+num_init]) ] + g_c = [g if g_m else torch.zeros_like(gi) for g, g_m, gi in zip(g_c, carry_mask, args[num_additional_inputs_masked:num_additional_inputs_masked+num_init]) ] return [*new_g_additional_inputs, *g_c, *g_xs] From f36237f8b2118518c0957e67364fb25c99ef8a05 Mon Sep 17 00:00:00 2001 From: Thomas Bohnstingl Date: Sat, 21 Sep 2024 14:52:50 +0200 Subject: [PATCH 06/12] Working version of partial gradients --- test/functorch/test_control_flow.py | 63 ++++++++++++++-------------- torch/_higher_order_ops/scan.py | 64 ++++++++++++++++------------- 2 files changed, 69 insertions(+), 58 deletions(-) diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py index 0389fb259aa4a..b0ec02ace93c8 100644 --- a/test/functorch/test_control_flow.py +++ b/test/functorch/test_control_flow.py @@ -2917,12 +2917,12 @@ def RNN(x: torch.Tensor, y: torch.Tensor): @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda - # @parametrize("reverse", [False, True]) - @parametrize("reverse", [False]) - # @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) - @parametrize("compile_mode", ["eager"]) - # @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) - @parametrize("device", [torch.device("cpu")]) + @parametrize("reverse", [False, True]) + # @parametrize("reverse", [False]) + @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) + # @parametrize("compile_mode", ["eager"]) + @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) + # @parametrize("device", [torch.device("cpu")]) def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): import random torch._dynamo.reset() @@ -2936,36 +2936,35 @@ def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): # The last four are the additional_inputs autograds = [] - # # Basic test - # autograds.append([True, True, True, True, True, True, True, True]) + # Basic test + autograds.append([True, True, True, True, True, True, True, True]) - # # xs tests - # autograds.append([True, False, True, True, True, True, True, True]) - # autograds.append([False, False, True, True, True, True, True, True]) + # xs tests + autograds.append([True, False, True, True, True, True, True, True]) + autograds.append([False, False, True, True, True, True, True, True]) - # # init tests - # autograds.append([True, True, False, True, True, True, True, True]) - # autograds.append([True, True, False, False, True, True, True, True]) + # init tests + autograds.append([True, True, False, True, True, True, True, True]) + autograds.append([True, True, False, False, True, True, True, True]) - # # additional input tests - # autograds.append([True, True, True, True, False, True, False, True]) - # autograds.append([True, True, True, True, False, False, False, False]) + # additional input tests + autograds.append([True, True, True, True, False, True, False, True]) + autograds.append([True, True, True, True, False, False, False, False]) # # Complex cases - # autograds.append([True, False, False, False, False, False, False, True]) - autograds.append([False, False, False, True, False, False, False, True]) + autograds.append([True, False, False, False, False, False, False, True]) + autograds.append([False, False, True, True, False, False, False, True]) - # # Random tests - # for _ in range(10): - # autograds.append([bool(random.randint(0, 1)) for _ in range(8)]) + # Random tests + for _ in range(10): + autograds.append([bool(random.randint(0, 1)) for _ in range(8)]) for autograd in autograds: - print(autograd) - x = torch.randn(3, 10, 5, device=device, requires_grad=autograd[0]) x1 = torch.randn(3, 10, 5, device=device, requires_grad=autograd[1]) h = torch.randn(3, 7, device=device, requires_grad=autograd[2]) h_1 = torch.randn(3, 7, device=device, requires_grad=autograd[3]) + h_2 = torch.randn(3, 7, device=device) W_ih = torch.randn(5, 7, device=device, requires_grad=autograd[4]) b_ih = torch.randn(7, device=device, requires_grad=autograd[5]) W_hh = torch.randn(7, 7, device=device, requires_grad=autograd[6]) @@ -2974,19 +2973,23 @@ def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): params = [p for p, a in zip([x, x1, h, h_1, W_ih, b_ih, W_hh, b_hh], autograd) if a] def RNN(x: torch.Tensor, y: torch.Tensor): - c_new_1 = y[0] @ W_ih + y[1] @ W_ih - c_new_2 = b_ih + y[1] @ W_ih + x[1] + c_new_0 = torch.zeros(3, 7, device=device) + c_new_1 = y[0] @ W_ih + y[1] @ W_ih + b_ih + x[1] + c_new_2 = b_ih + y[1] @ W_ih h_new = torch.tanh(c_new_1 + x[0] @ W_hh + b_hh) - return (c_new_1, c_new_2), h_new + return (c_new_0, c_new_1, c_new_2), h_new - result = scan_fct(RNN, (h, h_1), (x, x1), dim=dim, reverse=reverse) - result_exp = _fake_scan(RNN, (h, h_1), (x, x1), dim=dim, reverse=reverse) + result = scan_fct(RNN, (h, h_1, h_2), (x, x1), dim=dim, reverse=reverse) + result_exp = _fake_scan(RNN, (h, h_1, h_2), (x, x1), dim=dim, reverse=reverse) self.assertEqual(result, result_exp) if autograd: result_flat = pytree.tree_leaves(result) result_exp_flat = pytree.tree_leaves(result_exp) - self.check_autograd([r for r in result_flat if r.requires_grad], [r for r in result_exp_flat if r.requires_grad], params) + exp_grad_mask = [True if r.requires_grad else False for r in result_exp_flat] + self.check_autograd([r for r, m in zip(result_flat, exp_grad_mask) if m], + [r for r, m in zip(result_exp_flat, exp_grad_mask) if m], + params) @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda diff --git a/torch/_higher_order_ops/scan.py b/torch/_higher_order_ops/scan.py index c8602eacdd316..98461f47e9cdb 100644 --- a/torch/_higher_order_ops/scan.py +++ b/torch/_higher_order_ops/scan.py @@ -111,18 +111,33 @@ def wrapper_fwd_combine_fn(*args): new_carry, y = _extract_carry_and_out(combine_fn(*args), num_init) return [*new_carry, *[n_c.clone().detach() for n_c in new_carry], *y] - carry, outs = _extract_carry_and_out( + # Get gradient masks of inits, xs and additional_inputs + carry_mask = get_gradient_mask(init) + xs_mask = get_gradient_mask(xs) + additional_inputs_mask = get_gradient_mask(additional_inputs) + + # Enforce that the gradients of the inits are traced + fw_init_force_grad_track = [pytree.tree_map(_from_fun, x) for x in fw_init] + for el in fw_init_force_grad_track: + el.requires_grad = True + + fw_xs_force_grad_track = [pytree.tree_map(_from_fun, x) for x in fw_xs] + for el in fw_xs_force_grad_track: + el.requires_grad = True + + fw_additional_inputs_force_grad_track = [pytree.tree_map(_from_fun, x) for x in fw_additional_inputs] + for el in fw_additional_inputs_force_grad_track: + el.requires_grad = True + + carry_unforced_grad, outs_unforced_grad = _extract_carry_and_out( wrapper_fwd_combine_fn(*fw_init, *fw_xs, *fw_additional_inputs), num_init, ) - # # TODO: Support this in the future - # if pytree.tree_any( - # lambda t: not t.requires_grad, # type: ignore[union-attr] - # combine_fn(*fw_init, *fw_xs, *fw_additional_inputs), - # ): - # raise RuntimeError( - # "scan currently only supports Autograd if all init, xs and lifted parameters require gradients." - # ) + + carry, outs = _extract_carry_and_out( + wrapper_fwd_combine_fn(*fw_init_force_grad_track, *fw_xs_force_grad_track, *fw_additional_inputs_force_grad_track), + num_init, + ) fw_carry, fw_outputs = [pytree.tree_map(_from_fun, c) for c in carry], [ pytree.tree_map(_from_fun, o) for o in outs @@ -153,11 +168,11 @@ def wrapper_fwd_combine_fn(*args): # xs_mask = get_gradient_mask(fw_xs) # additional_inputs_mask = get_gradient_mask(fw_additional_inputs) - # Enforce that the gradients of the inits are traced - for el in fw_init: - el.requires_grad = True - for el in fw_carry: - el.requires_grad = True + # # Enforce that the gradients of the inits are traced + # for el in fw_init: + # el.requires_grad = True + # for el in fw_carry: + # el.requires_grad = True # fw_xs = [pytree.tree_map(_from_fun, x) for x in fw_xs] # for el in fw_xs: @@ -166,7 +181,8 @@ def wrapper_fwd_combine_fn(*args): _, joint_graph = create_fw_bw_graph( combine_fn, False, - (*fw_init, *fw_xs, *fw_additional_inputs), + # (*fw_init, *fw_xs, *fw_additional_inputs), + (*fw_init_force_grad_track, *fw_xs_force_grad_track, *fw_additional_inputs_force_grad_track), (*fw_carry, *fw_outputs[num_init:]), ) @@ -178,9 +194,10 @@ def wrapper_fwd_combine_fn(*args): *fw_xs, *fw_additional_inputs,), num_init ) - carry_mask = get_gradient_mask(g_c) - xs_mask = get_gradient_mask(g_xs[: len(g_xs) - num_additional_inputs]) - additional_inputs_mask = get_gradient_mask(g_xs[len(g_xs) - num_additional_inputs :]) + # carry_mask = [unforced_g & forced_g for unforced_g, forced_g in zip(carry_mask, get_gradient_mask(g_c))] + carry_mask_wrapper = get_gradient_mask(g_c) + xs_mask = [unforced_g & forced_g for unforced_g, forced_g in zip(xs_mask, get_gradient_mask(g_xs[: len(g_xs) - num_additional_inputs]))] + additional_inputs_mask = [unforced_g & forced_g for unforced_g, forced_g in zip(additional_inputs_mask, get_gradient_mask(g_xs[len(g_xs) - num_additional_inputs :]))] bw_additional_inputs = [add_inp for add_inp, add_inp_m in zip(bw_additional_inputs, additional_inputs_mask) if add_inp_m] num_additional_inputs_masked = len(bw_additional_inputs) @@ -203,7 +220,7 @@ def wrapper_bwd_combine_fn(*args): ] g_xs = g_xs[: len(g_xs) - num_additional_inputs] g_xs = [g for g, g_m in zip(g_xs, xs_mask) if g_m] - g_c = [g if g_m else torch.zeros_like(gi) for g, g_m, gi in zip(g_c, carry_mask, args[num_additional_inputs_masked:num_additional_inputs_masked+num_init]) ] + g_c = [g if g_m else torch.zeros_like(gi) for g, g_m, gi in zip(g_c, carry_mask_wrapper, args[num_additional_inputs_masked:num_additional_inputs_masked+num_init]) ] return [*new_g_additional_inputs, *g_c, *g_xs] @@ -758,15 +775,6 @@ def scan_autograd(combine_fn, init, xs, dim, reverse, additional_inputs): with torch._C._AutoDispatchBelowAutograd(): return scan_op(combine_fn, init, xs, dim, reverse, additional_inputs) - # # TODO: Support this in the future - # if pytree.tree_any( - # lambda t: not t.requires_grad, # type: ignore[union-attr] - # (init, xs, additional_inputs), - # ): - # raise RuntimeError( - # "scan currently only supports Autograd if all init, xs and lifted parameters 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. From fdc5ce4f974e8ad748407c6afd0ac1c788464d55 Mon Sep 17 00:00:00 2001 From: Thomas Bohnstingl Date: Sun, 22 Sep 2024 00:13:01 +0200 Subject: [PATCH 07/12] Working version of partial gradients --- test/functorch/test_control_flow.py | 62 ++++---- torch/_higher_order_ops/scan.py | 219 ++++++++++++++++------------ torch/_higher_order_ops/utils.py | 4 +- 3 files changed, 158 insertions(+), 127 deletions(-) diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py index b0ec02ace93c8..60aedcfc91a3a 100644 --- a/test/functorch/test_control_flow.py +++ b/test/functorch/test_control_flow.py @@ -2890,13 +2890,13 @@ def RNN(x: torch.Tensor, y: torch.Tensor): grads = grads[:2] self.assertEqual(grads, expected_grads) self.assertEqual(add_input_grads, expected_add_input_grads) - + @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda @parametrize("reverse", [False, True]) @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) @parametrize("autograd", [False, True]) - def test_scan_closure_RNN_parameters_as_inputs(self, reverse, compile_mode, device, autograd): + def test_scan_closure_RNN_parameters_as_inputs(self, reverse, device, autograd): x = torch.randn(3, 5, 10, device=device, requires_grad=autograd) h = torch.randn(3, 7, device=device, requires_grad=autograd) W_ih = torch.randn(5, 7, device=device, requires_grad=autograd) @@ -2910,51 +2910,51 @@ def RNN(x: torch.Tensor, y: torch.Tensor): return c_new, h_new with self.assertRaisesRegex( - RuntimeError, - "All xs leaves must at least have.*", + # Should be: RuntimeError, + # "All xs leaves must have a scan dimension > 0", + torch._dynamo.exc.Unsupported, + "Observed exception.*", ): result = scan(RNN, h, [x, W_ih, b_ih, W_hh, b_hh], dim=2, reverse=reverse) @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda @parametrize("reverse", [False, True]) - # @parametrize("reverse", [False]) @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) - # @parametrize("compile_mode", ["eager"]) @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) - # @parametrize("device", [torch.device("cpu")]) def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): import random + torch._dynamo.reset() torch._dynamo.config.cache_size_limit = 100 dim = 1 scan_fct = compile_mode_helper(scan, compile_mode) - + # The first two booleans are the xs # The second two are the inits # The last four are the additional_inputs autograds = [] - + # Basic test autograds.append([True, True, True, True, True, True, True, True]) - + # xs tests autograds.append([True, False, True, True, True, True, True, True]) autograds.append([False, False, True, True, True, True, True, True]) - + # init tests autograds.append([True, True, False, True, True, True, True, True]) autograds.append([True, True, False, False, True, True, True, True]) - + # additional input tests autograds.append([True, True, True, True, False, True, False, True]) autograds.append([True, True, True, True, False, False, False, False]) - + # # Complex cases autograds.append([True, False, False, False, False, False, False, True]) autograds.append([False, False, True, True, False, False, False, True]) - + # Random tests for _ in range(10): autograds.append([bool(random.randint(0, 1)) for _ in range(8)]) @@ -2970,7 +2970,11 @@ def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): W_hh = torch.randn(7, 7, device=device, requires_grad=autograd[6]) b_hh = torch.randn(7, device=device, requires_grad=autograd[7]) - params = [p for p, a in zip([x, x1, h, h_1, W_ih, b_ih, W_hh, b_hh], autograd) if a] + params = [ + p + for p, a in zip([x, x1, h, h_1, W_ih, b_ih, W_hh, b_hh], autograd) + if a + ] def RNN(x: torch.Tensor, y: torch.Tensor): c_new_0 = torch.zeros(3, 7, device=device) @@ -2980,27 +2984,29 @@ def RNN(x: torch.Tensor, y: torch.Tensor): return (c_new_0, c_new_1, c_new_2), h_new result = scan_fct(RNN, (h, h_1, h_2), (x, x1), dim=dim, reverse=reverse) - result_exp = _fake_scan(RNN, (h, h_1, h_2), (x, x1), dim=dim, reverse=reverse) + result_exp = _fake_scan( + RNN, (h, h_1, h_2), (x, x1), dim=dim, reverse=reverse + ) self.assertEqual(result, result_exp) if autograd: result_flat = pytree.tree_leaves(result) result_exp_flat = pytree.tree_leaves(result_exp) - exp_grad_mask = [True if r.requires_grad else False for r in result_exp_flat] - self.check_autograd([r for r, m in zip(result_flat, exp_grad_mask) if m], - [r for r, m in zip(result_exp_flat, exp_grad_mask) if m], - params) + exp_grad_mask = [ + True if r.requires_grad else False for r in result_exp_flat + ] + self.check_autograd( + [r for r, m in zip(result_flat, exp_grad_mask) if m], + [r for r, m in zip(result_exp_flat, exp_grad_mask) if m], + params, + ) @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda @parametrize("reverse", [False, True]) - # @parametrize("reverse", [False]) @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) - # @parametrize("compile_mode", ["eager"]) @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) - # @parametrize("device", [torch.device("cpu")]) @parametrize("autograd", [False, True]) - # @parametrize("autograd", [True]) def test_scan_closure_combine_fn_with_no_grad( self, reverse, compile_mode, device, autograd ): @@ -3018,13 +3024,13 @@ def RNN_1(x: torch.Tensor, y: torch.Tensor): with torch.no_grad(): h_new = torch.tanh(c_new + x @ W_hh + b_hh) return c_new, h_new - + def RNN_2(x: torch.Tensor, y: torch.Tensor): h_new = torch.tanh(x + x @ W_hh + b_hh) with torch.no_grad(): c_new = y @ W_ih + b_ih return c_new, h_new - + def RNN_3(x: torch.Tensor, y: torch.Tensor): with torch.no_grad(): c_new = y @ W_ih + b_ih @@ -3037,14 +3043,14 @@ def RNN_3(x: torch.Tensor, y: torch.Tensor): if autograd: self.check_autograd(result[0], result_exp[0], (x, W_ih, b_ih)) - + result = scan_fct(RNN_2, h, x, dim=dim, reverse=reverse) result_exp = _fake_scan(RNN_2, h, x, dim=dim, reverse=reverse) self.assertEqual(result, result_exp) if autograd: self.check_autograd(result[1], result_exp[1], (h, W_hh, b_hh)) - + result = scan_fct(RNN_3, h, x, dim=dim, reverse=reverse) result_exp = _fake_scan(RNN_3, h, x, dim=dim, reverse=reverse) self.assertEqual(result, result_exp) diff --git a/torch/_higher_order_ops/scan.py b/torch/_higher_order_ops/scan.py index 98461f47e9cdb..54da723bcd9c8 100644 --- a/torch/_higher_order_ops/scan.py +++ b/torch/_higher_order_ops/scan.py @@ -65,10 +65,15 @@ def shift_source_dim_to_target_dim(t, from_dim: int, to_dim: int): dims.insert(to_dim, from_dim) return t.permute(*dims) + 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] + + # Internal functions for scan.py def wrap_combine_fn_flat( *args, combine_fn, spec_init, spec_xs, num_init_leaves, num_inp_leaves @@ -111,31 +116,34 @@ def wrapper_fwd_combine_fn(*args): new_carry, y = _extract_carry_and_out(combine_fn(*args), num_init) return [*new_carry, *[n_c.clone().detach() for n_c in new_carry], *y] - # Get gradient masks of inits, xs and additional_inputs + # The forward graph needs to be constructed from a different combine_fn than the joint_graph + fw_graph = _maybe_reenter_make_fx(wrapper_fwd_combine_fn)( + *fw_init, *fw_xs, *fw_additional_inputs + ) + + # Get gradient masks of inits, these masks are used during the backward path to + # return gradients for inits with requires_grad=True and None for the others carry_mask = get_gradient_mask(init) xs_mask = get_gradient_mask(xs) additional_inputs_mask = get_gradient_mask(additional_inputs) - - # Enforce that the gradients of the inits are traced - fw_init_force_grad_track = [pytree.tree_map(_from_fun, x) for x in fw_init] - for el in fw_init_force_grad_track: - el.requires_grad = True - - fw_xs_force_grad_track = [pytree.tree_map(_from_fun, x) for x in fw_xs] - for el in fw_xs_force_grad_track: - el.requires_grad = True - - fw_additional_inputs_force_grad_track = [pytree.tree_map(_from_fun, x) for x in fw_additional_inputs] - for el in fw_additional_inputs_force_grad_track: - el.requires_grad = True - - carry_unforced_grad, outs_unforced_grad = _extract_carry_and_out( - wrapper_fwd_combine_fn(*fw_init, *fw_xs, *fw_additional_inputs), - num_init, - ) + + # Enforce that the gradients of the inits, xs and additional_inputs are traced + fw_init_force_grad_track = [ + pytree.tree_map(_from_fun, x, True) for x in fw_init + ] + fw_xs_force_grad_track = [ + pytree.tree_map(_from_fun, x, True) for x in fw_xs + ] + fw_additional_inputs_force_grad_track = [ + pytree.tree_map(_from_fun, x, True) for x in fw_additional_inputs + ] carry, outs = _extract_carry_and_out( - wrapper_fwd_combine_fn(*fw_init_force_grad_track, *fw_xs_force_grad_track, *fw_additional_inputs_force_grad_track), + wrapper_fwd_combine_fn( + *fw_init_force_grad_track, + *fw_xs_force_grad_track, + *fw_additional_inputs_force_grad_track, + ), num_init, ) @@ -158,59 +166,65 @@ def wrapper_fwd_combine_fn(*args): f"Got types {[type(out) for out in fw_outputs]}." ) - # The forward graph needs to be constructed from a different combine_fn than the joint_graph - fw_graph = _maybe_reenter_make_fx(wrapper_fwd_combine_fn)( - *fw_init, *fw_xs, *fw_additional_inputs - ) - - # # Get mask to mask None gradients - # carry_mask = get_gradient_mask(fw_init) - # xs_mask = get_gradient_mask(fw_xs) - # additional_inputs_mask = get_gradient_mask(fw_additional_inputs) - - # # Enforce that the gradients of the inits are traced - # for el in fw_init: - # el.requires_grad = True - # for el in fw_carry: - # el.requires_grad = True - - # fw_xs = [pytree.tree_map(_from_fun, x) for x in fw_xs] - # for el in fw_xs: - # el.requires_grad = True - + # The joint graph is constructed with the requires_grad forced to True for the init, xs and additional_inputs + # This is necessary because during the backward scan, we need the g_init for the other gradients, even if + # we don't directly need g_init _, joint_graph = create_fw_bw_graph( combine_fn, False, - # (*fw_init, *fw_xs, *fw_additional_inputs), - (*fw_init_force_grad_track, *fw_xs_force_grad_track, *fw_additional_inputs_force_grad_track), + ( + *fw_init_force_grad_track, + *fw_xs_force_grad_track, + *fw_additional_inputs_force_grad_track, + ), (*fw_carry, *fw_outputs[num_init:]), ) - - # Get mask to mask None gradients + g_c, g_xs = _extract_carry_and_out( - joint_graph(*fw_carry, - *fw_outputs[num_init:], - *fw_init, - *fw_xs, - *fw_additional_inputs,), num_init + joint_graph( + *fw_carry, + *fw_outputs[num_init:], + *fw_init, + *fw_xs, + *fw_additional_inputs, + ), + num_init, ) - # carry_mask = [unforced_g & forced_g for unforced_g, forced_g in zip(carry_mask, get_gradient_mask(g_c))] + # The gradient masks are combined with the initial ones. + # The reason is that the combine_fn might contain ``with torch.no_grad()`` statements + # Thus, even if the gradients of init, xs or additional_inputs should be tracked, + # The ``torch.no_grad()`` statements may break the gradient tracking carry_mask_wrapper = get_gradient_mask(g_c) - xs_mask = [unforced_g & forced_g for unforced_g, forced_g in zip(xs_mask, get_gradient_mask(g_xs[: len(g_xs) - num_additional_inputs]))] - additional_inputs_mask = [unforced_g & forced_g for unforced_g, forced_g in zip(additional_inputs_mask, get_gradient_mask(g_xs[len(g_xs) - num_additional_inputs :]))] - - bw_additional_inputs = [add_inp for add_inp, add_inp_m in zip(bw_additional_inputs, additional_inputs_mask) if add_inp_m] + xs_mask = [ + unforced_g & forced_g + for unforced_g, forced_g in zip( + xs_mask, + get_gradient_mask(g_xs[: len(g_xs) - num_additional_inputs]), + ) + ] + additional_inputs_mask = [ + unforced_g & forced_g + for unforced_g, forced_g in zip( + additional_inputs_mask, + get_gradient_mask(g_xs[len(g_xs) - num_additional_inputs :]), + ) + ] + + bw_additional_inputs = mask_gradient( + bw_additional_inputs, additional_inputs_mask + ) num_additional_inputs_masked = len(bw_additional_inputs) def wrapper_bwd_combine_fn(*args): carried_g_additional_input = args[:num_additional_inputs_masked] - # carried_g_additional_input = [add_inp for add_inp, add_inp_m in zip(carried_g_additional_input, additional_inputs_mask) if add_inp_m] g_c, g_xs = _extract_carry_and_out( joint_graph(*args[num_additional_inputs_masked:]), num_init ) current_g_additional_inputs = g_xs[len(g_xs) - num_additional_inputs :] - current_g_additional_inputs = [add_inp for add_inp, add_inp_m in zip(current_g_additional_inputs, additional_inputs_mask) if add_inp_m] + current_g_additional_inputs = mask_gradient( + current_g_additional_inputs, additional_inputs_mask + ) new_g_additional_inputs = [ carr_g + curr_g @@ -219,9 +233,19 @@ def wrapper_bwd_combine_fn(*args): ) ] g_xs = g_xs[: len(g_xs) - num_additional_inputs] - g_xs = [g for g, g_m in zip(g_xs, xs_mask) if g_m] - g_c = [g if g_m else torch.zeros_like(gi) for g, g_m, gi in zip(g_c, carry_mask_wrapper, args[num_additional_inputs_masked:num_additional_inputs_masked+num_init]) ] - + g_xs = mask_gradient(g_xs, xs_mask) + g_c = [ + g if g_m else torch.zeros_like(gi) + for g, g_m, gi in zip( + g_c, + carry_mask_wrapper, + args[ + num_additional_inputs_masked : num_additional_inputs_masked + + num_init + ], + ) + ] + return [*new_g_additional_inputs, *g_c, *g_xs] new_joint_graph = _maybe_reenter_make_fx(wrapper_bwd_combine_fn)( @@ -326,8 +350,14 @@ def add(x: torch.Tensor, y: torch.Tensor): raise RuntimeError("All init leaves must be a Tensor") if any(not isinstance(x, torch.Tensor) for x in leaves_xs): raise RuntimeError("All xs leaves must be a Tensor") + if any(x.ndim < dim for x in leaves_xs): + raise RuntimeError( + "All xs leaves must at least have 'dim' number of dimensions and scan dimension > 0" + ) if any(x.shape[dim] == 0 for x in leaves_xs): - raise RuntimeError("All xs leaves must have a scan dimension > 0") + raise RuntimeError( + "All xs leaves must at least have 'dim' number of dimensions and scan dimension > 0" + ) if len(leaves_xs) > 0: shape = leaves_xs[0].shape @@ -409,7 +439,11 @@ def _scan(init, xs): # Compute dummy shapes for the pre-allocation num_init_leaves = len(init) dummy_carry, dummy_out = _extract_carry_and_out( - operator(*carry, *[first_slice_copy(elem, dim) for elem in xs], *additional_inputs,), + operator( + *carry, + *[first_slice_copy(elem, dim) for elem in xs], + *additional_inputs, + ), num_init_leaves, ) @@ -418,22 +452,18 @@ def _scan(init, xs): # idxs -> Index matrix for scatter_ # out: (num_elems, M, N, ...) # idx: (1, M, N) - if len(dummy_out) > 0: - outs, idxs = zip( - *[ - [ - torch.zeros( - [num_elems] + list(e.size()), - dtype=e.dtype, - device=e.device, - ), - torch.ones_like(e, dtype=torch.int64).unsqueeze(0), - ] - for i, e in enumerate(dummy_out) - ] + outs = [ + torch.zeros( + [num_elems] + list(e.size()), + dtype=e.dtype, + device=e.device, ) - else: - outs, idxs = [], [] + for i, e in enumerate(dummy_out) + ] + idxs = [ + torch.ones_like(e, dtype=torch.int64).unsqueeze(0) + for i, e in enumerate(dummy_out) + ] def store_out_in_outs(out, ind): # Store the intermediate out in the outs matrix @@ -568,7 +598,7 @@ def forward( list(flat_args), num_leaves_init, num_leaves_xs ) ctx._num_leaves_additional_inputs = len(additional_inputs) - + ctx._carry_mask = carry_mask ctx._xs_mask = xs_mask ctx._additional_inputs_mask = additional_inputs_mask @@ -665,8 +695,8 @@ def prepare_initial_gradients( # The initial gradients for the additional_inputs are all zeros g_additional_inputs = [torch.zeros_like(ai) for ai in additional_inputs] return g_c_T, g_ys, g_additional_inputs - - def fill_grads_with_mask(real_grads, mask): + + def expand_grads_with_None(real_grads, mask): g_list = [] true_cnt = 0 for m in mask: @@ -680,13 +710,11 @@ def fill_grads_with_mask(real_grads, mask): joint_graph = ctx._joint_graph dim = ctx._dim reverse = ctx._reverse - num_leaves_init = ctx._num_leaves_init + num_leaves_init = ctx._num_leaves_init num_leaves_xs = ctx._num_leaves_xs - num_leaves_additional_inputs = ctx._num_leaves_additional_inputs num_leaves_ys = ctx._num_leaves_ys - + carry_mask = ctx._carry_mask - num_carry_mask = sum(carry_mask) xs_mask = ctx._xs_mask num_xs_mask = sum(xs_mask) additional_inputs_mask = ctx._additional_inputs_mask @@ -721,8 +749,10 @@ def fill_grads_with_mask(real_grads, mask): num_leaves_ys, bwd_scan_dim, ) - - g_additional_inputs = [add_inp for add_inp, add_inp_m in zip(g_additional_inputs, additional_inputs_mask) if add_inp_m] + + g_additional_inputs = mask_gradient( + g_additional_inputs, additional_inputs_mask + ) # Prepare the inputs for the backward scan. # This involves flipping the input xs if needed as well as @@ -732,7 +762,6 @@ def fill_grads_with_mask(real_grads, mask): ) xs_bwd = [*g_ys, *carries, *xs] - import pdb g_outs = scan_op( joint_graph, [*g_additional_inputs, *g_c_T], @@ -741,24 +770,20 @@ def fill_grads_with_mask(real_grads, mask): True, additional_inputs, ) - # pdb.set_trace() new_g_additional_inputs = g_outs[:num_additional_inputs_mask] g_init = g_outs[ num_additional_inputs_mask : num_additional_inputs_mask + num_leaves_init ] - # pdb.set_trace() - g_init = [g for g, g_m in zip(g_init, carry_mask) if g_m] - # pdb.set_trace() - g_xs = g_outs[len(g_outs) - num_xs_mask:] + g_init = mask_gradient(g_init, carry_mask) + g_xs = g_outs[len(g_outs) - num_xs_mask :] g_xs = prepare_final_gradients_xs(g_xs, dim, reverse) - - new_g_additional_inputs = fill_grads_with_mask(new_g_additional_inputs, additional_inputs_mask) - # pdb.set_trace() - g_xs = fill_grads_with_mask(g_xs, xs_mask) - g_init = fill_grads_with_mask(g_init, carry_mask) - # pdb.set_trace() + new_g_additional_inputs = expand_grads_with_None( + new_g_additional_inputs, additional_inputs_mask + ) + g_xs = expand_grads_with_None(g_xs, xs_mask) + g_init = expand_grads_with_None(g_init, carry_mask) return *[None] * 9, *g_init, *g_xs, *new_g_additional_inputs diff --git a/torch/_higher_order_ops/utils.py b/torch/_higher_order_ops/utils.py index 139e9a160cbe2..d252eb616178c 100644 --- a/torch/_higher_order_ops/utils.py +++ b/torch/_higher_order_ops/utils.py @@ -213,7 +213,7 @@ def unique_graph_id(proxy_mode, prefix): return i, next_name -def _from_fun(t): +def _from_fun(t, force_requires_grad=False): from torch._functorch.aot_autograd import from_fun from torch._subclasses.functional_tensor import FunctionalTensor @@ -223,7 +223,7 @@ def _from_fun(t): t.size(), t.stride(), dtype=t.dtype, - requires_grad=t.requires_grad, + requires_grad=t.requires_grad if not force_requires_grad else True, device=t.device, ) else: From a5100209b322f6c520d5c81cdf65a0e432761d0c Mon Sep 17 00:00:00 2001 From: Thomas Bohnstingl Date: Sun, 22 Sep 2024 10:47:24 +0200 Subject: [PATCH 08/12] Minor simplification --- torch/_higher_order_ops/scan.py | 25 +++++++++---------------- 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/torch/_higher_order_ops/scan.py b/torch/_higher_order_ops/scan.py index 54da723bcd9c8..ea1249507dd0d 100644 --- a/torch/_higher_order_ops/scan.py +++ b/torch/_higher_order_ops/scan.py @@ -194,21 +194,14 @@ def wrapper_fwd_combine_fn(*args): # The reason is that the combine_fn might contain ``with torch.no_grad()`` statements # Thus, even if the gradients of init, xs or additional_inputs should be tracked, # The ``torch.no_grad()`` statements may break the gradient tracking - carry_mask_wrapper = get_gradient_mask(g_c) - xs_mask = [ - unforced_g & forced_g - for unforced_g, forced_g in zip( - xs_mask, - get_gradient_mask(g_xs[: len(g_xs) - num_additional_inputs]), - ) - ] - additional_inputs_mask = [ - unforced_g & forced_g - for unforced_g, forced_g in zip( - additional_inputs_mask, - get_gradient_mask(g_xs[len(g_xs) - num_additional_inputs :]), - ) - ] + # carry_mask_wrapper = get_gradient_mask(g_c) + carry_mask = carry_mask and get_gradient_mask(g_c) + xs_mask = xs_mask and get_gradient_mask( + g_xs[: len(g_xs) - num_additional_inputs] + ) + additional_inputs_mask = additional_inputs_mask and get_gradient_mask( + g_xs[len(g_xs) - num_additional_inputs :] + ) bw_additional_inputs = mask_gradient( bw_additional_inputs, additional_inputs_mask @@ -238,7 +231,7 @@ def wrapper_bwd_combine_fn(*args): g if g_m else torch.zeros_like(gi) for g, g_m, gi in zip( g_c, - carry_mask_wrapper, + carry_mask, args[ num_additional_inputs_masked : num_additional_inputs_masked + num_init From 01ca2cf1808241dd5e5209f8a094e677d81815f3 Mon Sep 17 00:00:00 2001 From: Thomas Bohnstingl Date: Thu, 10 Oct 2024 00:30:27 +0200 Subject: [PATCH 09/12] Pre-cleanup --- test/functorch/test_control_flow.py | 176 ++++++++++++++++++++++------ torch/_higher_order_ops/scan.py | 168 ++++++++++++++++++-------- 2 files changed, 260 insertions(+), 84 deletions(-) diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py index 60aedcfc91a3a..64342a35cfc5c 100644 --- a/test/functorch/test_control_flow.py +++ b/test/functorch/test_control_flow.py @@ -2920,8 +2920,12 @@ def RNN(x: torch.Tensor, y: torch.Tensor): @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda @parametrize("reverse", [False, True]) + # @parametrize("reverse", [False]) @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) + # @parametrize("compile_mode", ["none"]) + # @parametrize("compile_mode", ["compile"]) @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) + # @parametrize("device", [torch.device("cpu")]) def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): import random @@ -2964,7 +2968,6 @@ def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): x1 = torch.randn(3, 10, 5, device=device, requires_grad=autograd[1]) h = torch.randn(3, 7, device=device, requires_grad=autograd[2]) h_1 = torch.randn(3, 7, device=device, requires_grad=autograd[3]) - h_2 = torch.randn(3, 7, device=device) W_ih = torch.randn(5, 7, device=device, requires_grad=autograd[4]) b_ih = torch.randn(7, device=device, requires_grad=autograd[5]) W_hh = torch.randn(7, 7, device=device, requires_grad=autograd[6]) @@ -2977,15 +2980,17 @@ def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): ] def RNN(x: torch.Tensor, y: torch.Tensor): - c_new_0 = torch.zeros(3, 7, device=device) - c_new_1 = y[0] @ W_ih + y[1] @ W_ih + b_ih + x[1] - c_new_2 = b_ih + y[1] @ W_ih - h_new = torch.tanh(c_new_1 + x[0] @ W_hh + b_hh) - return (c_new_0, c_new_1, c_new_2), h_new - - result = scan_fct(RNN, (h, h_1, h_2), (x, x1), dim=dim, reverse=reverse) + c_new_0 = x[0] + 1 + c_new_1 = x[1] + 1 + h_new = torch.tanh(c_new_1 + x[0] @ W_hh + b_hh) + y[0] @ W_ih + y[1] @ W_ih + b_ih + x[1] + return (c_new_0, c_new_1), h_new + + inits = (h, h_1) + result = scan_fct(RNN, inits, (x, x1), dim=dim, reverse=reverse) + # print([val.requires_grad for val in inits]) + # print([val.requires_grad for val in pytree.tree_leaves(result[0])]) result_exp = _fake_scan( - RNN, (h, h_1, h_2), (x, x1), dim=dim, reverse=reverse + RNN, (h, h_1), (x, x1), dim=dim, reverse=reverse ) self.assertEqual(result, result_exp) @@ -3006,55 +3011,152 @@ def RNN(x: torch.Tensor, y: torch.Tensor): @parametrize("reverse", [False, True]) @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) - @parametrize("autograd", [False, True]) - def test_scan_closure_combine_fn_with_no_grad( + @parametrize("autograd", [True]) + def test_scan_closure_combine_fn_with_no_grad_init_carries_unequal_grad( self, reverse, compile_mode, device, autograd ): dim = 1 scan_fct = compile_mode_helper(scan, compile_mode) - x = torch.randn(3, 10, 5, device=device, requires_grad=autograd) - h = torch.randn(3, 7, device=device, requires_grad=autograd) - W_ih = torch.randn(5, 7, device=device, requires_grad=autograd) - b_ih = torch.randn(7, device=device, requires_grad=autograd) - W_hh = torch.randn(7, 7, device=device, requires_grad=autograd) - b_hh = torch.randn(7, device=device, requires_grad=autograd) + x = torch.randn(3, 10, 7, device=device, requires_grad=autograd) + h1 = torch.randn(3, 7, device=device, requires_grad=autograd) + h2 = torch.randn(3, 7, device=device, requires_grad=autograd) - def RNN_1(x: torch.Tensor, y: torch.Tensor): - c_new = y @ W_ih + b_ih + def fct_c1_no_grad(x: torch.Tensor, y: torch.Tensor): + h_new = torch.tanh(x[0] + x[1] + y) + c2 = x[1] + y with torch.no_grad(): - h_new = torch.tanh(c_new + x @ W_hh + b_hh) - return c_new, h_new + c1 = x[0] + y + return (c1, c2), h_new + + with self.assertRaisesRegex( + RuntimeError, + "The init and carries need to have the same require_grad structure!.*", + ): + result = scan_fct(fct_c1_no_grad, (h1, h2), x, dim=dim, reverse=reverse) + + @unittest.skipIf(not SM70OrLater, "triton") + @requires_cuda + @parametrize("reverse", [False, True]) + @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) + @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) + @parametrize("autograd", [False, True]) + def test_scan_closure_combine_fn_with_no_grad_init_carries_equal_grad( + self, reverse, compile_mode, device, autograd + ): + dim = 1 + scan_fct = compile_mode_helper(scan, compile_mode) + x = torch.randn(3, 10, 7, device=device, requires_grad=autograd) + h1 = torch.randn(3, 7, device=device, requires_grad=False) + h2 = torch.randn(3, 7, device=device, requires_grad=autograd) - def RNN_2(x: torch.Tensor, y: torch.Tensor): - h_new = torch.tanh(x + x @ W_hh + b_hh) + def fct_c1_no_grad(x: torch.Tensor, y: torch.Tensor): + h_new = torch.tanh(x[0] + x[1] + y) + c2 = x[1] + y with torch.no_grad(): - c_new = y @ W_ih + b_ih - return c_new, h_new + c1 = x[0] + y + return (c1, c2), h_new + + result = scan_fct(fct_c1_no_grad, (h1, h2), x, dim=dim, reverse=reverse) + result_exp = _fake_scan(fct_c1_no_grad, (h1, h2), x, dim=dim, reverse=reverse) + self.assertEqual(result, result_exp) - def RNN_3(x: torch.Tensor, y: torch.Tensor): + if autograd: + # TODO: Ideally we should be able to select the results that require gradients like this + # [leaf for leaf in pytree.tree_leaves(result) if leaf.requires_grad == True] + # However, for the scan operator this does not work, as all outputs always have + # grad_fn= + res_req_grad_flat = pytree.tree_leaves(result)[1:] + res_exp_req_grad_flat = pytree.tree_leaves(result_exp)[1:] + self.check_autograd(res_req_grad_flat, res_exp_req_grad_flat, (x, h2)) + + @unittest.skipIf(not SM70OrLater, "triton") + @requires_cuda + @parametrize("reverse", [False, True]) + @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) + @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) + @parametrize("autograd", [False, True]) + def test_scan_closure_combine_fn_with_no_grad_for_out( + self, reverse, compile_mode, device, autograd + ): + dim = 1 + scan_fct = compile_mode_helper(scan, compile_mode) + x = torch.randn(3, 10, 7, device=device, requires_grad=autograd) + h1 = torch.randn(3, 7, device=device, requires_grad=autograd) + h2 = torch.randn(3, 7, device=device, requires_grad=autograd) + + def fct_ys_no_grad(x: torch.Tensor, y: torch.Tensor): + c1 = x[0] + y + c2 = x[1] + y with torch.no_grad(): - c_new = y @ W_ih + b_ih - h_new = torch.tanh(c_new + x @ W_hh + b_hh) - return c_new, h_new + h_new = torch.tanh(x[0] + x[1] + y) + return (c1, c2), h_new - result = scan_fct(RNN_1, h, x, dim=dim, reverse=reverse) - result_exp = _fake_scan(RNN_1, h, x, dim=dim, reverse=reverse) + result = scan_fct(fct_ys_no_grad, (h1, h2), x, dim=dim, reverse=reverse) + result_exp = _fake_scan(fct_ys_no_grad, (h1, h2), x, dim=dim, reverse=reverse) self.assertEqual(result, result_exp) if autograd: - self.check_autograd(result[0], result_exp[0], (x, W_ih, b_ih)) + self.check_autograd(result[0], result_exp[0], (x, h1, h2)) - result = scan_fct(RNN_2, h, x, dim=dim, reverse=reverse) - result_exp = _fake_scan(RNN_2, h, x, dim=dim, reverse=reverse) + + @unittest.skipIf(not SM70OrLater, "triton") + @requires_cuda + @parametrize("reverse", [False, True]) + # @parametrize("reverse", [False]) + @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) + # @parametrize("compile_mode", ["none"]) + # @parametrize("compile_mode", ["compile"]) + @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) + @parametrize("autograd", [False, True]) + # @parametrize("autograd", [True]) + def test_scan_closure_combine_fn_with_no_grad_additional_inputs( + self, reverse, compile_mode, device, autograd + ): + dim = 1 + scan_fct = compile_mode_helper(scan, compile_mode) + x = torch.randn(3, 10, 7, device=device, requires_grad=autograd) + h = torch.randn(3, 7, device=device, requires_grad=autograd) + W_ih = torch.randn(7, 7, device=device, requires_grad=autograd) + b_ih = torch.randn(7, device=device, requires_grad=autograd) + W_hh = torch.randn(7, 7, device=device, requires_grad=autograd) + b_hh = torch.randn(7, device=device, requires_grad=autograd) + + def fct_no_grad_bhh_Whh(x: torch.Tensor, y: torch.Tensor): + c_new = y @ W_ih + b_ih + x + + h_new = c_new + 1 + with torch.no_grad(): + h_new_no_grad = torch.tanh(x @ W_hh + b_hh) + h_new2 = h_new + h_new_no_grad + + # h_new2 = torch.tanh(x @ W_hh + b_hh) + + return c_new, h_new2 + + result = scan_fct(fct_no_grad_bhh_Whh, h, x, dim=dim, reverse=reverse) + result_exp = _fake_scan(fct_no_grad_bhh_Whh, h, x, dim=dim, reverse=reverse) self.assertEqual(result, result_exp) if autograd: - self.check_autograd(result[1], result_exp[1], (h, W_hh, b_hh)) + self.check_autograd(result[1], result_exp[1], (h, x, W_ih, b_ih)) - result = scan_fct(RNN_3, h, x, dim=dim, reverse=reverse) - result_exp = _fake_scan(RNN_3, h, x, dim=dim, reverse=reverse) + def fct_no_grad_bih_Wih_bhh_Whh(x: torch.Tensor, y: torch.Tensor): + c_new = x + y# + x @ W_hh + h_new = c_new + 1 + with torch.no_grad(): + c_new_no_grad = y @ W_ih + b_ih + h_new_no_grad = torch.tanh(x @ W_hh + b_hh) + c_new2 = c_new + c_new_no_grad + h_new2 = h_new + h_new_no_grad + return c_new2, h_new2 + + result = scan_fct(fct_no_grad_bih_Wih_bhh_Whh, h, x, dim=dim, reverse=reverse) + result_exp = _fake_scan(fct_no_grad_bih_Wih_bhh_Whh, h, x, dim=dim, reverse=reverse) self.assertEqual(result, result_exp) + if autograd: + self.check_autograd(result[1], result_exp[1], (h, x)) + @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda @parametrize("reverse", [False, True]) diff --git a/torch/_higher_order_ops/scan.py b/torch/_higher_order_ops/scan.py index ea1249507dd0d..c9e2338f368e1 100644 --- a/torch/_higher_order_ops/scan.py +++ b/torch/_higher_order_ops/scan.py @@ -123,26 +123,29 @@ def wrapper_fwd_combine_fn(*args): # Get gradient masks of inits, these masks are used during the backward path to # return gradients for inits with requires_grad=True and None for the others - carry_mask = get_gradient_mask(init) + init_mask = get_gradient_mask(init) xs_mask = get_gradient_mask(xs) additional_inputs_mask = get_gradient_mask(additional_inputs) # Enforce that the gradients of the inits, xs and additional_inputs are traced - fw_init_force_grad_track = [ - pytree.tree_map(_from_fun, x, True) for x in fw_init - ] - fw_xs_force_grad_track = [ - pytree.tree_map(_from_fun, x, True) for x in fw_xs - ] - fw_additional_inputs_force_grad_track = [ - pytree.tree_map(_from_fun, x, True) for x in fw_additional_inputs - ] + # fw_init_force_grad_track = [ + # pytree.tree_map(_from_fun, x, True) for x in fw_init + # ] + # fw_xs_force_grad_track = [ + # pytree.tree_map(_from_fun, x, True) for x in fw_xs + # ] + # fw_additional_inputs_force_grad_track = [ + # pytree.tree_map(_from_fun, x, True) for x in fw_additional_inputs + # ] carry, outs = _extract_carry_and_out( wrapper_fwd_combine_fn( - *fw_init_force_grad_track, - *fw_xs_force_grad_track, - *fw_additional_inputs_force_grad_track, + # *fw_init_force_grad_track, + *fw_init, + # *fw_xs_force_grad_track, + *fw_xs, + # *fw_additional_inputs_force_grad_track, + *fw_additional_inputs, ), num_init, ) @@ -173,9 +176,12 @@ def wrapper_fwd_combine_fn(*args): combine_fn, False, ( - *fw_init_force_grad_track, - *fw_xs_force_grad_track, - *fw_additional_inputs_force_grad_track, + # *fw_init_force_grad_track, + *fw_init, + # *fw_xs_force_grad_track, + *fw_xs, + # *fw_additional_inputs_force_grad_track, + *fw_additional_inputs, ), (*fw_carry, *fw_outputs[num_init:]), ) @@ -190,50 +196,75 @@ def wrapper_fwd_combine_fn(*args): ), num_init, ) + # xs_mask = get_gradient_mask(xs) + # additional_inputs_mask = get_gradient_mask(additional_inputs) + # The gradient masks are combined with the initial ones. # The reason is that the combine_fn might contain ``with torch.no_grad()`` statements # Thus, even if the gradients of init, xs or additional_inputs should be tracked, # The ``torch.no_grad()`` statements may break the gradient tracking - # carry_mask_wrapper = get_gradient_mask(g_c) - carry_mask = carry_mask and get_gradient_mask(g_c) + # carry_mask = carry_mask and get_gradient_mask(g_c) + if any(cm is not im or g_cm is not im for g_cm, cm, im in zip(get_gradient_mask(g_c), get_gradient_mask(carry), init_mask)): + raise RuntimeError( + "The init and carries need to have the same require_grad structure! \ + E.g., check for `with torch.no_gard` statements in the combine_fn." + ) + xs_mask = xs_mask and get_gradient_mask( g_xs[: len(g_xs) - num_additional_inputs] ) + # xs_mask = get_gradient_mask( + # g_xs[: len(g_xs) - num_additional_inputs] + # ) additional_inputs_mask = additional_inputs_mask and get_gradient_mask( g_xs[len(g_xs) - num_additional_inputs :] ) + # additional_inputs_mask = get_gradient_mask( + # g_xs[len(g_xs) - num_additional_inputs :] + # ) - bw_additional_inputs = mask_gradient( - bw_additional_inputs, additional_inputs_mask - ) - num_additional_inputs_masked = len(bw_additional_inputs) + # bw_additional_inputs = mask_gradient( + # bw_additional_inputs, additional_inputs_mask + # ) + # num_additional_inputs_masked = len(bw_additional_inputs) def wrapper_bwd_combine_fn(*args): - carried_g_additional_input = args[:num_additional_inputs_masked] + # carried_g_additional_input = args[:num_additional_inputs_masked] + carried_g_additional_input = args[:num_additional_inputs] g_c, g_xs = _extract_carry_and_out( - joint_graph(*args[num_additional_inputs_masked:]), num_init + # joint_graph(*args[num_additional_inputs_masked:]), num_init + joint_graph(*args[num_additional_inputs:]), num_init ) current_g_additional_inputs = g_xs[len(g_xs) - num_additional_inputs :] - current_g_additional_inputs = mask_gradient( - current_g_additional_inputs, additional_inputs_mask - ) + # current_g_additional_inputs = mask_gradient( + # current_g_additional_inputs, additional_inputs_mask + # ) new_g_additional_inputs = [ - carr_g + curr_g - for carr_g, curr_g in zip( - carried_g_additional_input, current_g_additional_inputs + # carr_g + curr_g + # The clone().detach() is required to avoid aliasing inputs + carr_g + curr_g if add_inp_m else carr_g.clone().detach() + # for carr_g, curr_g in zip( + for add_inp_m, carr_g, curr_g in zip( + # carried_g_additional_input, current_g_additional_inputs + additional_inputs_mask, carried_g_additional_input, current_g_additional_inputs ) ] g_xs = g_xs[: len(g_xs) - num_additional_inputs] + + # We need to mask the g_xs so that no None values are returned + # The reason being that in the backward implementation, we store + # The gradients in a matrix and thus, None values are problematic g_xs = mask_gradient(g_xs, xs_mask) g_c = [ g if g_m else torch.zeros_like(gi) for g, g_m, gi in zip( g_c, - carry_mask, + init_mask, args[ - num_additional_inputs_masked : num_additional_inputs_masked + # num_additional_inputs_masked : num_additional_inputs_masked + num_additional_inputs : num_additional_inputs + num_init ], ) @@ -243,13 +274,14 @@ def wrapper_bwd_combine_fn(*args): new_joint_graph = _maybe_reenter_make_fx(wrapper_bwd_combine_fn)( *bw_additional_inputs, + # *fw_additional_inputs, *fw_carry, *fw_outputs[num_init:], *fw_init, *fw_xs, *fw_additional_inputs, ) - return fw_graph, new_joint_graph, carry_mask, xs_mask, additional_inputs_mask + return fw_graph, new_joint_graph, init_mask, xs_mask, additional_inputs_mask def scan( @@ -590,7 +622,7 @@ def forward( init, xs, additional_inputs = ScanAutogradOp.extract_init_xs_additional_inputs( list(flat_args), num_leaves_init, num_leaves_xs ) - ctx._num_leaves_additional_inputs = len(additional_inputs) + ctx._num_additional_inputs = len(additional_inputs) ctx._carry_mask = carry_mask ctx._xs_mask = xs_mask @@ -689,13 +721,22 @@ def prepare_initial_gradients( g_additional_inputs = [torch.zeros_like(ai) for ai in additional_inputs] return g_c_T, g_ys, g_additional_inputs + def mask_grads_with_None(real_grads, mask): + g_list = [] + for m, g in zip(mask, real_grads): + if m: + g_list.append(g) + else: + g_list.append(None) + return g_list + def expand_grads_with_None(real_grads, mask): g_list = [] - true_cnt = 0 + real_grads_cnt = 0 for m in mask: if m: - g_list.append(real_grads[true_cnt]) - true_cnt += 1 + g_list.append(real_grads[real_grads_cnt]) + real_grads_cnt += 1 else: g_list.append(None) return g_list @@ -711,7 +752,8 @@ def expand_grads_with_None(real_grads, mask): xs_mask = ctx._xs_mask num_xs_mask = sum(xs_mask) additional_inputs_mask = ctx._additional_inputs_mask - num_additional_inputs_mask = sum(additional_inputs_mask) + # num_additional_inputs_mask = sum(additional_inputs_mask) + num_additional_inputs = ctx._num_additional_inputs # The results from the forward scan are always stacked on dim 0 # The gradients though need to be provided with the correct scan dimension dim @@ -733,6 +775,8 @@ def expand_grads_with_None(real_grads, mask): for o in flat_args[num_leaves_init : num_leaves_init + num_leaves_xs] ] + import pdb + # pdb.set_trace() with torch._C._AutoDispatchBelowAutograd(): # Prepare the initial gradients for the backward scan g_c_T, g_ys, g_additional_inputs = prepare_initial_gradients( @@ -742,10 +786,10 @@ def expand_grads_with_None(real_grads, mask): num_leaves_ys, bwd_scan_dim, ) - - g_additional_inputs = mask_gradient( - g_additional_inputs, additional_inputs_mask - ) + # pdb.set_trace() + # g_additional_inputs = mask_gradient( + # g_additional_inputs, additional_inputs_mask + # ) # Prepare the inputs for the backward scan. # This involves flipping the input xs if needed as well as @@ -755,6 +799,7 @@ def expand_grads_with_None(real_grads, mask): ) xs_bwd = [*g_ys, *carries, *xs] + # pdb.set_trace() g_outs = scan_op( joint_graph, [*g_additional_inputs, *g_c_T], @@ -763,20 +808,27 @@ def expand_grads_with_None(real_grads, mask): True, additional_inputs, ) - new_g_additional_inputs = g_outs[:num_additional_inputs_mask] + # pdb.set_trace() + # new_g_additional_inputs = g_outs[:num_additional_inputs_mask] + new_g_additional_inputs = g_outs[:num_additional_inputs] g_init = g_outs[ - num_additional_inputs_mask : num_additional_inputs_mask + # num_additional_inputs_mask : num_additional_inputs_mask + num_additional_inputs : num_additional_inputs + num_leaves_init ] - g_init = mask_gradient(g_init, carry_mask) + # pdb.set_trace() + # g_init = mask_gradient(g_init, carry_mask) g_xs = g_outs[len(g_outs) - num_xs_mask :] g_xs = prepare_final_gradients_xs(g_xs, dim, reverse) - new_g_additional_inputs = expand_grads_with_None( + # pdb.set_trace() + # new_g_additional_inputs = expand_grads_with_None( + new_g_additional_inputs = mask_grads_with_None( new_g_additional_inputs, additional_inputs_mask ) + # pdb.set_trace() g_xs = expand_grads_with_None(g_xs, xs_mask) - g_init = expand_grads_with_None(g_init, carry_mask) + g_init = mask_grads_with_None(g_init, carry_mask) return *[None] * 9, *g_init, *g_xs, *new_g_additional_inputs @@ -817,6 +869,24 @@ def scan_autograd(combine_fn, init, xs, dim, reverse, additional_inputs): additional_inputs_mask, ) = create_fw_bw_graph_combinefn(combine_fn, init, xs, dim, additional_inputs) + # class Fn(torch.autograd.Function): + # @staticmethod + # def forward(ctx, fwd_graph, *x): + # with torch._C._AutoDispatchBelowAutograd(): + # # outs = fwd_graph(*x) + # # ctx.save_for_backward(*outs,) + # # return *outs, + # return *x, + # def backward(ctx, grad): + # return grad.cos() + # # a = Fn.apply(fw_graph, *(init + [first_slice_copy(x, dim) for x in xs] + additional_inputs)) + # # a = Fn.apply(fw_graph, *(init)) + # # a = Fn.apply(fw_graph, *init) + # a = Fn.apply(fw_graph, *[torch.randn(3, 3, requires_grad=False), torch.randn(3, 3, requires_grad=False)]) + # print([el.requires_grad for el in a]) + # a = Fn.apply(fw_graph, *[torch.randn(3, 3, requires_grad=False), torch.randn(3, 3, requires_grad=True)]) + # print([el.requires_grad for el in a]) + flat_out = ScanAutogradOp.apply( fw_graph, joint_graph, @@ -843,6 +913,7 @@ def scan_proxy_mode(mode, combine_fn, init, xs, dim, reverse, additional_inputs) def scan_fake_tensor_mode(mode, combine_fn, init, xs, dim, reverse, additional_inputs): with mode: scan_length = xs[0].shape[dim] + # try: carry, outputs = _extract_carry_and_out( combine_fn( *init, @@ -851,6 +922,9 @@ def scan_fake_tensor_mode(mode, combine_fn, init, xs, dim, reverse, additional_i ), len(init), ) + # except: + # import pdb + # pdb.set_trace() out = [ *carry, *[stack_y(t, scan_length) for t in outputs], From 59ce6dd27ad312945ecfa4a5eba705d931b9e1ee Mon Sep 17 00:00:00 2001 From: Thomas Bohnstingl Date: Thu, 10 Oct 2024 00:49:07 +0200 Subject: [PATCH 10/12] After clean-up --- test/functorch/test_control_flow.py | 41 ++++----- torch/_higher_order_ops/scan.py | 125 ++++++---------------------- 2 files changed, 41 insertions(+), 125 deletions(-) diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py index 64342a35cfc5c..087f0f64cf754 100644 --- a/test/functorch/test_control_flow.py +++ b/test/functorch/test_control_flow.py @@ -2920,12 +2920,8 @@ def RNN(x: torch.Tensor, y: torch.Tensor): @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda @parametrize("reverse", [False, True]) - # @parametrize("reverse", [False]) @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) - # @parametrize("compile_mode", ["none"]) - # @parametrize("compile_mode", ["compile"]) @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) - # @parametrize("device", [torch.device("cpu")]) def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): import random @@ -2982,16 +2978,18 @@ def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): def RNN(x: torch.Tensor, y: torch.Tensor): c_new_0 = x[0] + 1 c_new_1 = x[1] + 1 - h_new = torch.tanh(c_new_1 + x[0] @ W_hh + b_hh) + y[0] @ W_ih + y[1] @ W_ih + b_ih + x[1] + h_new = ( + torch.tanh(c_new_1 + x[0] @ W_hh + b_hh) + + y[0] @ W_ih + + y[1] @ W_ih + + b_ih + + x[1] + ) return (c_new_0, c_new_1), h_new inits = (h, h_1) result = scan_fct(RNN, inits, (x, x1), dim=dim, reverse=reverse) - # print([val.requires_grad for val in inits]) - # print([val.requires_grad for val in pytree.tree_leaves(result[0])]) - result_exp = _fake_scan( - RNN, (h, h_1), (x, x1), dim=dim, reverse=reverse - ) + result_exp = _fake_scan(RNN, (h, h_1), (x, x1), dim=dim, reverse=reverse) self.assertEqual(result, result_exp) if autograd: @@ -3033,7 +3031,7 @@ def fct_c1_no_grad(x: torch.Tensor, y: torch.Tensor): "The init and carries need to have the same require_grad structure!.*", ): result = scan_fct(fct_c1_no_grad, (h1, h2), x, dim=dim, reverse=reverse) - + @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda @parametrize("reverse", [False, True]) @@ -3063,12 +3061,12 @@ def fct_c1_no_grad(x: torch.Tensor, y: torch.Tensor): if autograd: # TODO: Ideally we should be able to select the results that require gradients like this # [leaf for leaf in pytree.tree_leaves(result) if leaf.requires_grad == True] - # However, for the scan operator this does not work, as all outputs always have + # However, for the scan operator this does not work, as all outputs always have # grad_fn= res_req_grad_flat = pytree.tree_leaves(result)[1:] res_exp_req_grad_flat = pytree.tree_leaves(result_exp)[1:] self.check_autograd(res_req_grad_flat, res_exp_req_grad_flat, (x, h2)) - + @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda @parametrize("reverse", [False, True]) @@ -3098,17 +3096,12 @@ def fct_ys_no_grad(x: torch.Tensor, y: torch.Tensor): if autograd: self.check_autograd(result[0], result_exp[0], (x, h1, h2)) - @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda @parametrize("reverse", [False, True]) - # @parametrize("reverse", [False]) @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) - # @parametrize("compile_mode", ["none"]) - # @parametrize("compile_mode", ["compile"]) @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) @parametrize("autograd", [False, True]) - # @parametrize("autograd", [True]) def test_scan_closure_combine_fn_with_no_grad_additional_inputs( self, reverse, compile_mode, device, autograd ): @@ -3123,14 +3116,12 @@ def test_scan_closure_combine_fn_with_no_grad_additional_inputs( def fct_no_grad_bhh_Whh(x: torch.Tensor, y: torch.Tensor): c_new = y @ W_ih + b_ih + x - + h_new = c_new + 1 with torch.no_grad(): h_new_no_grad = torch.tanh(x @ W_hh + b_hh) h_new2 = h_new + h_new_no_grad - - # h_new2 = torch.tanh(x @ W_hh + b_hh) - + return c_new, h_new2 result = scan_fct(fct_no_grad_bhh_Whh, h, x, dim=dim, reverse=reverse) @@ -3141,7 +3132,7 @@ def fct_no_grad_bhh_Whh(x: torch.Tensor, y: torch.Tensor): self.check_autograd(result[1], result_exp[1], (h, x, W_ih, b_ih)) def fct_no_grad_bih_Wih_bhh_Whh(x: torch.Tensor, y: torch.Tensor): - c_new = x + y# + x @ W_hh + c_new = x + y h_new = c_new + 1 with torch.no_grad(): c_new_no_grad = y @ W_ih + b_ih @@ -3151,7 +3142,9 @@ def fct_no_grad_bih_Wih_bhh_Whh(x: torch.Tensor, y: torch.Tensor): return c_new2, h_new2 result = scan_fct(fct_no_grad_bih_Wih_bhh_Whh, h, x, dim=dim, reverse=reverse) - result_exp = _fake_scan(fct_no_grad_bih_Wih_bhh_Whh, h, x, dim=dim, reverse=reverse) + result_exp = _fake_scan( + fct_no_grad_bih_Wih_bhh_Whh, h, x, dim=dim, reverse=reverse + ) self.assertEqual(result, result_exp) if autograd: diff --git a/torch/_higher_order_ops/scan.py b/torch/_higher_order_ops/scan.py index c9e2338f368e1..d635f0c7fcd22 100644 --- a/torch/_higher_order_ops/scan.py +++ b/torch/_higher_order_ops/scan.py @@ -121,30 +121,17 @@ def wrapper_fwd_combine_fn(*args): *fw_init, *fw_xs, *fw_additional_inputs ) - # Get gradient masks of inits, these masks are used during the backward path to + # Get gradient masks of inits, xs and additional_inputs. + # These masks are used during the backward path to # return gradients for inits with requires_grad=True and None for the others init_mask = get_gradient_mask(init) xs_mask = get_gradient_mask(xs) additional_inputs_mask = get_gradient_mask(additional_inputs) - # Enforce that the gradients of the inits, xs and additional_inputs are traced - # fw_init_force_grad_track = [ - # pytree.tree_map(_from_fun, x, True) for x in fw_init - # ] - # fw_xs_force_grad_track = [ - # pytree.tree_map(_from_fun, x, True) for x in fw_xs - # ] - # fw_additional_inputs_force_grad_track = [ - # pytree.tree_map(_from_fun, x, True) for x in fw_additional_inputs - # ] - carry, outs = _extract_carry_and_out( wrapper_fwd_combine_fn( - # *fw_init_force_grad_track, *fw_init, - # *fw_xs_force_grad_track, *fw_xs, - # *fw_additional_inputs_force_grad_track, *fw_additional_inputs, ), num_init, @@ -176,11 +163,8 @@ def wrapper_fwd_combine_fn(*args): combine_fn, False, ( - # *fw_init_force_grad_track, *fw_init, - # *fw_xs_force_grad_track, *fw_xs, - # *fw_additional_inputs_force_grad_track, *fw_additional_inputs, ), (*fw_carry, *fw_outputs[num_init:]), @@ -196,65 +180,52 @@ def wrapper_fwd_combine_fn(*args): ), num_init, ) - # xs_mask = get_gradient_mask(xs) - # additional_inputs_mask = get_gradient_mask(additional_inputs) - - # The gradient masks are combined with the initial ones. - # The reason is that the combine_fn might contain ``with torch.no_grad()`` statements - # Thus, even if the gradients of init, xs or additional_inputs should be tracked, - # The ``torch.no_grad()`` statements may break the gradient tracking - # carry_mask = carry_mask and get_gradient_mask(g_c) - if any(cm is not im or g_cm is not im for g_cm, cm, im in zip(get_gradient_mask(g_c), get_gradient_mask(carry), init_mask)): + + # Check whether the init and the carries have the same requires_grad flags + # Scan enforces that the levaes of inits and of the carries require the same gradients + if any( + cm is not im or g_cm is not im + for g_cm, cm, im in zip( + get_gradient_mask(g_c), get_gradient_mask(carry), init_mask + ) + ): raise RuntimeError( "The init and carries need to have the same require_grad structure! \ E.g., check for `with torch.no_gard` statements in the combine_fn." ) - + + # The gradient masks are combined with the initial ones. + # The reason is that the combine_fn might contain ``with torch.no_grad()`` statements + # Thus, even if the gradients of xs or additional_inputs should be tracked, + # The ``torch.no_grad()`` statements may break the gradient tracking xs_mask = xs_mask and get_gradient_mask( g_xs[: len(g_xs) - num_additional_inputs] ) - # xs_mask = get_gradient_mask( - # g_xs[: len(g_xs) - num_additional_inputs] - # ) additional_inputs_mask = additional_inputs_mask and get_gradient_mask( g_xs[len(g_xs) - num_additional_inputs :] ) - # additional_inputs_mask = get_gradient_mask( - # g_xs[len(g_xs) - num_additional_inputs :] - # ) - - # bw_additional_inputs = mask_gradient( - # bw_additional_inputs, additional_inputs_mask - # ) - # num_additional_inputs_masked = len(bw_additional_inputs) def wrapper_bwd_combine_fn(*args): - # carried_g_additional_input = args[:num_additional_inputs_masked] carried_g_additional_input = args[:num_additional_inputs] g_c, g_xs = _extract_carry_and_out( - # joint_graph(*args[num_additional_inputs_masked:]), num_init joint_graph(*args[num_additional_inputs:]), num_init ) current_g_additional_inputs = g_xs[len(g_xs) - num_additional_inputs :] - # current_g_additional_inputs = mask_gradient( - # current_g_additional_inputs, additional_inputs_mask - # ) new_g_additional_inputs = [ - # carr_g + curr_g # The clone().detach() is required to avoid aliasing inputs carr_g + curr_g if add_inp_m else carr_g.clone().detach() - # for carr_g, curr_g in zip( for add_inp_m, carr_g, curr_g in zip( - # carried_g_additional_input, current_g_additional_inputs - additional_inputs_mask, carried_g_additional_input, current_g_additional_inputs + additional_inputs_mask, + carried_g_additional_input, + current_g_additional_inputs, ) ] g_xs = g_xs[: len(g_xs) - num_additional_inputs] - + # We need to mask the g_xs so that no None values are returned - # The reason being that in the backward implementation, we store + # The reason being that in the backward implementation, we store # The gradients in a matrix and thus, None values are problematic g_xs = mask_gradient(g_xs, xs_mask) g_c = [ @@ -262,11 +233,7 @@ def wrapper_bwd_combine_fn(*args): for g, g_m, gi in zip( g_c, init_mask, - args[ - # num_additional_inputs_masked : num_additional_inputs_masked - num_additional_inputs : num_additional_inputs - + num_init - ], + args[num_additional_inputs : num_additional_inputs + num_init], ) ] @@ -274,7 +241,6 @@ def wrapper_bwd_combine_fn(*args): new_joint_graph = _maybe_reenter_make_fx(wrapper_bwd_combine_fn)( *bw_additional_inputs, - # *fw_additional_inputs, *fw_carry, *fw_outputs[num_init:], *fw_init, @@ -729,7 +695,7 @@ def mask_grads_with_None(real_grads, mask): else: g_list.append(None) return g_list - + def expand_grads_with_None(real_grads, mask): g_list = [] real_grads_cnt = 0 @@ -752,7 +718,6 @@ def expand_grads_with_None(real_grads, mask): xs_mask = ctx._xs_mask num_xs_mask = sum(xs_mask) additional_inputs_mask = ctx._additional_inputs_mask - # num_additional_inputs_mask = sum(additional_inputs_mask) num_additional_inputs = ctx._num_additional_inputs # The results from the forward scan are always stacked on dim 0 @@ -775,8 +740,6 @@ def expand_grads_with_None(real_grads, mask): for o in flat_args[num_leaves_init : num_leaves_init + num_leaves_xs] ] - import pdb - # pdb.set_trace() with torch._C._AutoDispatchBelowAutograd(): # Prepare the initial gradients for the backward scan g_c_T, g_ys, g_additional_inputs = prepare_initial_gradients( @@ -786,20 +749,11 @@ def expand_grads_with_None(real_grads, mask): num_leaves_ys, bwd_scan_dim, ) - # pdb.set_trace() - # g_additional_inputs = mask_gradient( - # g_additional_inputs, additional_inputs_mask - # ) - - # Prepare the inputs for the backward scan. - # This involves flipping the input xs if needed as well as - # Prepending the init of the forward scan to the carries xs, carries = prepare_xs_carries_for_bwd( xs, init, carries, bwd_scan_dim, reverse ) xs_bwd = [*g_ys, *carries, *xs] - # pdb.set_trace() g_outs = scan_op( joint_graph, [*g_additional_inputs, *g_c_T], @@ -808,25 +762,16 @@ def expand_grads_with_None(real_grads, mask): True, additional_inputs, ) - # pdb.set_trace() - # new_g_additional_inputs = g_outs[:num_additional_inputs_mask] new_g_additional_inputs = g_outs[:num_additional_inputs] g_init = g_outs[ - # num_additional_inputs_mask : num_additional_inputs_mask - num_additional_inputs : num_additional_inputs - + num_leaves_init + num_additional_inputs : num_additional_inputs + num_leaves_init ] - # pdb.set_trace() - # g_init = mask_gradient(g_init, carry_mask) g_xs = g_outs[len(g_outs) - num_xs_mask :] g_xs = prepare_final_gradients_xs(g_xs, dim, reverse) - # pdb.set_trace() - # new_g_additional_inputs = expand_grads_with_None( new_g_additional_inputs = mask_grads_with_None( new_g_additional_inputs, additional_inputs_mask ) - # pdb.set_trace() g_xs = expand_grads_with_None(g_xs, xs_mask) g_init = mask_grads_with_None(g_init, carry_mask) return *[None] * 9, *g_init, *g_xs, *new_g_additional_inputs @@ -869,24 +814,6 @@ def scan_autograd(combine_fn, init, xs, dim, reverse, additional_inputs): additional_inputs_mask, ) = create_fw_bw_graph_combinefn(combine_fn, init, xs, dim, additional_inputs) - # class Fn(torch.autograd.Function): - # @staticmethod - # def forward(ctx, fwd_graph, *x): - # with torch._C._AutoDispatchBelowAutograd(): - # # outs = fwd_graph(*x) - # # ctx.save_for_backward(*outs,) - # # return *outs, - # return *x, - # def backward(ctx, grad): - # return grad.cos() - # # a = Fn.apply(fw_graph, *(init + [first_slice_copy(x, dim) for x in xs] + additional_inputs)) - # # a = Fn.apply(fw_graph, *(init)) - # # a = Fn.apply(fw_graph, *init) - # a = Fn.apply(fw_graph, *[torch.randn(3, 3, requires_grad=False), torch.randn(3, 3, requires_grad=False)]) - # print([el.requires_grad for el in a]) - # a = Fn.apply(fw_graph, *[torch.randn(3, 3, requires_grad=False), torch.randn(3, 3, requires_grad=True)]) - # print([el.requires_grad for el in a]) - flat_out = ScanAutogradOp.apply( fw_graph, joint_graph, @@ -913,7 +840,6 @@ def scan_proxy_mode(mode, combine_fn, init, xs, dim, reverse, additional_inputs) def scan_fake_tensor_mode(mode, combine_fn, init, xs, dim, reverse, additional_inputs): with mode: scan_length = xs[0].shape[dim] - # try: carry, outputs = _extract_carry_and_out( combine_fn( *init, @@ -922,9 +848,6 @@ def scan_fake_tensor_mode(mode, combine_fn, init, xs, dim, reverse, additional_i ), len(init), ) - # except: - # import pdb - # pdb.set_trace() out = [ *carry, *[stack_y(t, scan_length) for t in outputs], From 418b32368177d64e7d9ea1fe9cdf6cdd7b2cf2d3 Mon Sep 17 00:00:00 2001 From: Thomas Bohnstingl Date: Fri, 11 Oct 2024 11:47:32 +0200 Subject: [PATCH 11/12] Fixed issue with output aliasing of joint_graph --- test/functorch/test_control_flow.py | 308 ++++++++++++++++++---------- torch/_higher_order_ops/scan.py | 11 +- torch/_higher_order_ops/utils.py | 74 +++++-- 3 files changed, 272 insertions(+), 121 deletions(-) diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py index 087f0f64cf754..215ae5bd19397 100644 --- a/test/functorch/test_control_flow.py +++ b/test/functorch/test_control_flow.py @@ -129,7 +129,7 @@ def compile_mode_helper(fct, compile_mode): return fct -def get_scan_combine_fn(name, associative=True): +def get_scan_combine_fn(name, expand_with_carry=False, parameters=None): def add(x: torch.Tensor, y: torch.Tensor): return x + y @@ -163,6 +163,18 @@ def non_pointwise(x: torch.Tensor, y: torch.Tensor): W = torch.diag(torch.ones(2, device=x.device)) return x @ W + y @ W + def RNN(x: torch.Tensor, y: torch.Tensor): + c_new = y @ parameters[0] + parameters[1] + h_new = torch.tanh(c_new + x @ parameters[2] + parameters[3]) + return h_new, h_new + + def fct_c1_no_grad(x: torch.Tensor, y: torch.Tensor): + h_new = torch.tanh(x[0] + x[1] + y) + c2 = x[1] + y + with torch.no_grad(): + c1 = x[0] + y + return (c1, c2), h_new + if name == "add": fct = add elif name == "adds": @@ -179,10 +191,14 @@ def non_pointwise(x: torch.Tensor, y: torch.Tensor): fct = complex_pointwise elif name == "non_pointwise": fct = non_pointwise + elif name == "RNN": + fct = RNN + elif name == "fct_c1_no_grad": + fct = fct_c1_no_grad else: raise ValueError("Combine_fn name unknown!") - if not associative: + if expand_with_carry: return lambda x, y: (fct(x, y), fct(x, y)) else: return fct @@ -1316,8 +1332,8 @@ def test_associative_scan_compile( scan_fct = compile_mode_helper(associative_scan, compile_mode) for op, op_pt in [ - (get_scan_combine_fn("add", True), torch.cumsum), - (get_scan_combine_fn("mul", True), torch.cumprod), + (get_scan_combine_fn("add", False), torch.cumsum), + (get_scan_combine_fn("mul", False), torch.cumprod), ]: result = scan_fct(op, x, 0, reverse=reverse, combine_mode=combine_mode) result_exp = _fake_associative_scan(op, xs=x, dim=0, reverse=reverse) @@ -1329,14 +1345,14 @@ def test_associative_scan_compile( # Jax Examples x = torch.arange(0, 4, device=device) cumsum1 = scan_fct( - get_scan_combine_fn("add", True), + get_scan_combine_fn("add", False), x, 0, reverse=reverse, combine_mode=combine_mode, ) cumsum_exp = _fake_associative_scan( - get_scan_combine_fn("add", True), x, 0, reverse=reverse + get_scan_combine_fn("add", False), x, 0, reverse=reverse ) if not reverse: self.assertEqual( @@ -1375,12 +1391,12 @@ def add2(x: torch.Tensor, y: torch.Tensor): for op, op_pt, init in [ ( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), torch.cumsum, torch.zeros(2, 2, device=device, requires_grad=autograd), ), ( - get_scan_combine_fn("mul", False), + get_scan_combine_fn("mul", True), torch.cumprod, torch.ones(2, 2, device=device, requires_grad=autograd), ), @@ -1399,14 +1415,14 @@ def add2(x: torch.Tensor, y: torch.Tensor): x = torch.arange(0, 4, device=device, dtype=torch.int64) init = torch.zeros(1, device=device, dtype=torch.int64) cumsum1 = scan_fct( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, x, dim=0, reverse=reverse, ) cumsum_exp = _fake_scan( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init=init, xs=x, dim=0, @@ -1451,14 +1467,14 @@ def add2(x: torch.Tensor, y: torch.Tensor): ) init = torch.ones(1, device=device, dtype=torch.float32, requires_grad=autograd) result = scan_fct( - get_scan_combine_fn("div", False), + get_scan_combine_fn("div", True), init, x, dim=0, reverse=reverse, ) result_exp = _fake_scan( - get_scan_combine_fn("div", False), + get_scan_combine_fn("div", True), init=init, xs=x, dim=0, @@ -1502,7 +1518,7 @@ def test_scan_dtype(self, reverse, compile_mode, device, dtype): # Check all outputs and carries on the correct device and with torch.float32 x = torch.randn(3, 10, 2, device=device).to(dtype=dtype) op, init = ( - get_scan_combine_fn("adds"), + get_scan_combine_fn("adds", False), torch.zeros(10, 2, device=device, dtype=dtype), ) result = scan_fct(op, init, x, dim=0, reverse=reverse) @@ -1521,7 +1537,7 @@ def test_scan_dtype(self, reverse, compile_mode, device, dtype): # carry.dtype torch.float32 and output.dtype torch.float16 x = torch.randn(3, 10, 2, device=device).to(dtype=dtype) op, init = ( - get_scan_combine_fn("adds"), + get_scan_combine_fn("adds", False), torch.zeros(10, 2, device=device, dtype=torch.float32), ) result = scan_fct(op, init, x, dim=0, reverse=reverse) @@ -1539,7 +1555,7 @@ def test_scan_dtype(self, reverse, compile_mode, device, dtype): # carry.dtype torch.int64 and output.dtype torch.float32 x = torch.randn(3, 10, 2, device=device) op, init = ( - get_scan_combine_fn("adds"), + get_scan_combine_fn("adds", False), torch.zeros(10, 2, device=device, dtype=dtype), ) result = scan_fct(op, init, x, dim=0, reverse=reverse) @@ -1577,8 +1593,8 @@ def test_associative_scan_dim(self, combine_mode, reverse, device): x = torch.randn(*shapes, device=device) for op, op_pt in [ - (get_scan_combine_fn("add", True), torch.cumsum), - (get_scan_combine_fn("mul", True), torch.cumprod), + (get_scan_combine_fn("add", False), torch.cumsum), + (get_scan_combine_fn("mul", False), torch.cumprod), ]: result = associative_scan( op, x, rnd_scan_dim, reverse=reverse, combine_mode=combine_mode @@ -1612,12 +1628,12 @@ def test_scan_dim(self, reverse, compile_mode, device, autograd): for op, op_pt, init in [ ( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), torch.cumsum, torch.zeros(*init_shapes, device=device, requires_grad=autograd), ), ( - get_scan_combine_fn("mul", False), + get_scan_combine_fn("mul", True), torch.cumprod, torch.ones(*init_shapes, device=device, requires_grad=autograd), ), @@ -1663,14 +1679,14 @@ def test_associative_scan_binary_operator(self, combine_mode, reverse, device): elements = (A.repeat((timesteps, 1)), projected_inputs) result1 = associative_scan( - get_scan_combine_fn("s5_operator", True), + get_scan_combine_fn("s5_operator", False), elements, 0, combine_mode=combine_mode, reverse=reverse, ) expected_result = _fake_associative_scan( - get_scan_combine_fn("s5_operator", True), elements, 0, reverse=reverse + get_scan_combine_fn("s5_operator", False), elements, 0, reverse=reverse ) self.assertEqual( result1, @@ -1706,14 +1722,14 @@ def test_scan_binary_operator(self, reverse, device, autograd): ) result = scan( - get_scan_combine_fn("s5_operator", False), + get_scan_combine_fn("s5_operator", True), init, elements, dim=0, reverse=reverse, ) expected_result = _fake_scan( - get_scan_combine_fn("s5_operator", False), + get_scan_combine_fn("s5_operator", True), init=init, xs=elements, dim=0, @@ -1757,14 +1773,14 @@ def test_associative_scan_tuple(self, combine_mode, reverse, device): inp = (x, y) result1 = associative_scan( - get_scan_combine_fn("tuple_fct", True), + get_scan_combine_fn("tuple_fct", False), inp, 0, reverse=reverse, combine_mode=combine_mode, ) expected_result = _fake_associative_scan( - get_scan_combine_fn("tuple_fct", True), inp, 0, reverse=reverse + get_scan_combine_fn("tuple_fct", False), inp, 0, reverse=reverse ) self.assertEqual(result1, expected_result) @@ -1779,14 +1795,14 @@ def test_scan_tuple(self, reverse, device): init = tuple(torch._ops.ops.aten.slice(e, 0, 0, 1, 1) for e in inp) result_same = scan( - get_scan_combine_fn("tuple_fct", False), + get_scan_combine_fn("tuple_fct", True), init, inp, dim=0, reverse=reverse, ) expected_result = _fake_scan( - get_scan_combine_fn("tuple_fct", False), + get_scan_combine_fn("tuple_fct", True), init=init, xs=inp, dim=0, @@ -1864,14 +1880,14 @@ def fct_pointwise(x, y): inp = {"i": x, "j": ([y], [{"o": z}])} result = associative_scan( - get_scan_combine_fn("complex_pointwise", True), + get_scan_combine_fn("complex_pointwise", False), inp, 0, combine_mode=combine_mode, reverse=reverse, ) expected_result = _fake_associative_scan( - get_scan_combine_fn("complex_pointwise", True), inp, 0, reverse=reverse + get_scan_combine_fn("complex_pointwise", False), inp, 0, reverse=reverse ) self.assertEqual(result, expected_result) @@ -1925,14 +1941,14 @@ def test_scan_complex_pytree(self, reverse, device): init = pytree.tree_unflatten(init_flat, inp_spec) result = scan( - get_scan_combine_fn("complex_pointwise", False), + get_scan_combine_fn("complex_pointwise", True), init, inp, dim=0, reverse=reverse, ) expected_result = _fake_scan( - get_scan_combine_fn("complex_pointwise", False), + get_scan_combine_fn("complex_pointwise", True), init=init, xs=inp, dim=0, @@ -1962,7 +1978,7 @@ def test_associative_scan_downstream_scan_matmul( def chain_fct(inp): W = torch.ones(2, 5, device=device) o = associative_scan( - get_scan_combine_fn("add", True), + get_scan_combine_fn("add", False), inp, 1, reverse=reverse, @@ -1974,7 +1990,7 @@ def chain_fct(inp): inp = torch.randn(3, 10, 2, device=device) expected_result = _fake_associative_scan( - get_scan_combine_fn("add", True), inp, 1, reverse=reverse + get_scan_combine_fn("add", False), inp, 1, reverse=reverse ) @ torch.ones(2, 5, device=device) result1 = fct_cmp(inp) self.assertEqual(result1, expected_result) @@ -2000,14 +2016,14 @@ def test_associative_scan_downstream_scan_scan( # Chain with scan def chain_fct_same_dim(inp): o1 = associative_scan( - get_scan_combine_fn("add", True), + get_scan_combine_fn("add", False), inp, 1, combine_mode=combine_mode, reverse=reverse, ) o2 = associative_scan( - get_scan_combine_fn("add", True), + get_scan_combine_fn("add", False), o1, 1, combine_mode=combine_mode, @@ -2020,9 +2036,9 @@ def chain_fct_same_dim(inp): inp = torch.randn(3, 10, 2, device=device) expected_result = _fake_associative_scan( - get_scan_combine_fn("add", True), + get_scan_combine_fn("add", False), _fake_associative_scan( - get_scan_combine_fn("add", True), inp, 1, reverse=reverse + get_scan_combine_fn("add", False), inp, 1, reverse=reverse ), 1, reverse=reverse, @@ -2051,14 +2067,14 @@ def test_associative_scan_downstream_scan_scan_different_dim( # Chain with scan on different dim def chain_fct_different_dim(inp): o1 = associative_scan( - get_scan_combine_fn("add", True), + get_scan_combine_fn("add", False), inp, 1, combine_mode=combine_mode, reverse=reverse, ) o2 = associative_scan( - get_scan_combine_fn("add", True), + get_scan_combine_fn("add", False), o1, 0, combine_mode=combine_mode, @@ -2070,9 +2086,9 @@ def chain_fct_different_dim(inp): inp = torch.randn(3, 10, 2, device=device) expected_result = _fake_associative_scan( - get_scan_combine_fn("add", True), + get_scan_combine_fn("add", False), _fake_associative_scan( - get_scan_combine_fn("add", True), inp, 1, reverse=reverse + get_scan_combine_fn("add", False), inp, 1, reverse=reverse ), 0, reverse=reverse, @@ -2095,7 +2111,7 @@ def test_scan_downstream_scan_matmul(self, compile_mode, reverse, device, autogr def chain_fct(inp): W = torch.ones(2, 5, device=device) o = scan( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, inp, dim=1, @@ -2106,7 +2122,7 @@ def chain_fct(inp): fct_cmp = compile_mode_helper(chain_fct, compile_mode) expected_result = _fake_scan( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init=init, xs=inp, dim=1, @@ -2135,7 +2151,7 @@ def test_scan_downstream_scan_scan_dim( def chain_fct_different_dim(inp): o1 = scan( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, inp, dim=1, @@ -2143,7 +2159,7 @@ def chain_fct_different_dim(inp): ) o1 = pytree.tree_map(lambda t: shift_source_dim_to_target_dim(t, 0, 1), o1) o2 = scan( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init2, o1[1], dim=0, @@ -2154,7 +2170,7 @@ def chain_fct_different_dim(inp): fct_cmp = compile_mode_helper(chain_fct_different_dim, compile_mode) xs = _fake_scan( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init=init, xs=inp, dim=1, @@ -2162,7 +2178,7 @@ def chain_fct_different_dim(inp): )[1] xs = pytree.tree_map(lambda t: shift_source_dim_to_target_dim(t, 0, 1), xs) expected_result = _fake_scan( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init=init2, xs=xs, dim=0, @@ -2192,7 +2208,7 @@ def test_associative_scan_non_pointwise(self, reverse, device): "For combine_mode='pointwise', the combine_fn needs to be pointwise", ): out = associative_scan( - get_scan_combine_fn("non_pointwise", True), + get_scan_combine_fn("non_pointwise", False), x, 0, reverse=reverse, @@ -2212,10 +2228,10 @@ def test_associative_scan_non_pointwise(self, reverse, device): def test_associative_scan_non_pointwise_generic(self, reverse, device): x = torch.randn(3, 10, 2, device=device) result_expected = _fake_associative_scan( - get_scan_combine_fn("non_pointwise", True), x, 0, reverse=reverse + get_scan_combine_fn("non_pointwise", False), x, 0, reverse=reverse ) result1 = associative_scan( - get_scan_combine_fn("non_pointwise", True), + get_scan_combine_fn("non_pointwise", False), x, 0, reverse=reverse, @@ -2230,7 +2246,7 @@ def test_scan_non_pointwise(self, reverse, device): x = torch.randn(3, 10, 2, device=device) init = torch.randn(10, 2, device=device) result_expected = _fake_scan( - get_scan_combine_fn("non_pointwise", False), + get_scan_combine_fn("non_pointwise", True), init=init, xs=x, dim=0, @@ -2238,7 +2254,7 @@ def test_scan_non_pointwise(self, reverse, device): ) out = scan( - get_scan_combine_fn("non_pointwise", False), + get_scan_combine_fn("non_pointwise", True), init, x, dim=0, @@ -2261,7 +2277,7 @@ def test_scan_compile_cnt(self, reverse, device): init = torch.randn(3, 5, device=device) # First compilation step torch.compile(scan, backend=cnt)( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, x, dim=dim, @@ -2273,7 +2289,7 @@ def test_scan_compile_cnt(self, reverse, device): init = torch.randn(3, 5, device=device) # Recompilation due to first different size torch.compile(scan, backend=cnt)( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, x, dim=dim, @@ -2285,7 +2301,7 @@ def test_scan_compile_cnt(self, reverse, device): init = torch.randn(3, 5, device=device) # No recompilation, because of dynamic shape torch.compile(scan, backend=cnt)( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, x, dim=dim, @@ -2297,7 +2313,7 @@ def test_scan_compile_cnt(self, reverse, device): init = torch.randn(3, 40, device=device) # Recompilation because of dim change torch.compile(scan, backend=cnt)( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, x, dim=2, @@ -2309,7 +2325,7 @@ def test_scan_compile_cnt(self, reverse, device): init = torch.randn(3, 40, device=device) # Recompilation due to first different size on new dim torch.compile(scan, backend=cnt)( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, x, dim=2, @@ -2321,7 +2337,7 @@ def test_scan_compile_cnt(self, reverse, device): init = torch.randn(3, 40, device=device) # No recompilation, because of dynamic shape on new dim torch.compile(scan, backend=cnt)( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, x, dim=2, @@ -2333,7 +2349,7 @@ def test_scan_compile_cnt(self, reverse, device): init = torch.randn(3, 40, device=device) # Recompilation because of dim change torch.compile(scan, backend=cnt)( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, x, dim=1, @@ -2345,7 +2361,7 @@ def test_scan_compile_cnt(self, reverse, device): init = torch.randn(3, 40, device=device) # Recompilation because of reverse change torch.compile(scan, backend=cnt)( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, x, dim=1, @@ -2357,7 +2373,7 @@ def test_scan_compile_cnt(self, reverse, device): init = torch.randn(3, 40, device=device) # No recompilation, as nothing changed torch.compile(scan, backend=cnt)( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, x, dim=1, @@ -2369,7 +2385,7 @@ def test_scan_compile_cnt(self, reverse, device): init = torch.randn(3, 80, device=device) # No recompilation, final test torch.compile(scan, backend=cnt)( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, x, dim=1, @@ -2399,7 +2415,7 @@ def test_scan_init_scanned_0(self, reverse, compile_mode, device): "Observed exception.*", ): result_init = scan_fct( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, inp, dim=dim, @@ -2425,7 +2441,7 @@ def test_scan_init_non_tensor(self, reverse, compile_mode, device): "Observed exception.*", ): result_init = scan_fct( - get_scan_combine_fn("add", False), init, x, dim=dim, reverse=reverse + get_scan_combine_fn("add", True), init, x, dim=dim, reverse=reverse ) @unittest.skipIf(not SM70OrLater, "triton") @@ -2450,7 +2466,7 @@ def test_scan_init_wrong_shape(self, reverse, compile_mode, device): "Observed exception.*", ): result_init = scan_fct( - get_scan_combine_fn("add", False), + get_scan_combine_fn("add", True), init, inp, dim=dim, @@ -2499,7 +2515,7 @@ def test_scan_init(self, reverse, compile_mode, device, autograd): # Only init and no input x = torch.randn(3, 1, 2, device=device, requires_grad=autograd) dim = 1 - op, op_pt = (get_scan_combine_fn("add", False), torch.cumsum) + op, op_pt = (get_scan_combine_fn("add", True), torch.cumsum) # Only init given init = torch._ops.ops.aten.slice(x, dim, 0, 1, 1) @@ -2516,7 +2532,7 @@ def test_scan_init(self, reverse, compile_mode, device, autograd): x = torch.randn(3, 5, 2, device=device, requires_grad=autograd) dim = 0 - op, op_pt = (get_scan_combine_fn("add", False), torch.cumsum) + op, op_pt = (get_scan_combine_fn("add", True), torch.cumsum) inp = torch._ops.ops.aten.slice(x, dim, 1, None, 1) # Init tensor scalar @@ -2580,7 +2596,7 @@ def add_scalar_carry_sliced_out(x: torch.Tensor, y: torch.Tensor): self.check_autograd(result_init, result_exp, (init, inp)) # Correct case - op, op_pt = (get_scan_combine_fn("add", False), torch.cumsum) + op, op_pt = (get_scan_combine_fn("add", True), torch.cumsum) x = torch.randn(3, 2, 2, device=device, requires_grad=autograd) init = torch.zeros(3, 2, device=device, requires_grad=autograd) dim = 2 @@ -2677,7 +2693,7 @@ def test_scan_init_wrong_pytree_complex(self, reverse, device): ".*", ): result = scan( - get_scan_combine_fn("complex_pointwise", False), + get_scan_combine_fn("complex_pointwise", True), init, inp, dim=0, @@ -2763,14 +2779,14 @@ def fct_pointwise_different_carry(x, y): ), } result = scan_fct( - get_scan_combine_fn("complex_pointwise", False), + get_scan_combine_fn("complex_pointwise", True), init, inp, dim=0, reverse=reverse, ) expected_result = _fake_scan( - get_scan_combine_fn("complex_pointwise", False), + get_scan_combine_fn("complex_pointwise", True), init, inp, dim=0, @@ -2847,15 +2863,19 @@ def test_scan_closure_RNN(self, autograd): W_hh = W_hh.detach() b_hh = b_hh.detach() - def RNN(x: torch.Tensor, y: torch.Tensor): - c_new = y @ W_ih + b_ih - h_new = torch.tanh(c_new + x @ W_hh + b_hh) - return h_new, h_new - expected_result = rnn(x, torch.unsqueeze(h, 0)) expected_result_out = expected_result[0] expected_result_state = expected_result[1][0, :] - result = scan(RNN, h, x, dim=dim, reverse=False) + + result = scan( + get_scan_combine_fn( + "RNN", expand_with_carry=False, parameters=[W_ih, b_ih, W_hh, b_hh] + ), + h, + x, + dim=dim, + reverse=False, + ) result_cmp = [result[0], shift_source_dim_to_target_dim(result[1], 0, dim)] self.assertEqual(result_cmp[0], expected_result_state) self.assertEqual(result_cmp[1], expected_result_out) @@ -2904,18 +2924,21 @@ def test_scan_closure_RNN_parameters_as_inputs(self, reverse, device, autograd): W_hh = torch.randn(7, 7, device=device, requires_grad=autograd) b_hh = torch.randn(7, device=device, requires_grad=autograd) - def RNN(x: torch.Tensor, y: torch.Tensor): - c_new = y[0] @ W_ih[1] + b_ih[2] - h_new = torch.tanh(c_new + x @ W_hh[3] + b_hh[4]) - return c_new, h_new - with self.assertRaisesRegex( # Should be: RuntimeError, # "All xs leaves must have a scan dimension > 0", torch._dynamo.exc.Unsupported, "Observed exception.*", ): - result = scan(RNN, h, [x, W_ih, b_ih, W_hh, b_hh], dim=2, reverse=reverse) + result = scan( + get_scan_combine_fn( + "RNN", expand_with_carry=False, parameters=[W_ih, b_ih, W_hh, b_hh] + ), + h, + [x, W_ih, b_ih, W_hh, b_hh], + dim=2, + reverse=reverse, + ) @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda @@ -3019,18 +3042,17 @@ def test_scan_closure_combine_fn_with_no_grad_init_carries_unequal_grad( h1 = torch.randn(3, 7, device=device, requires_grad=autograd) h2 = torch.randn(3, 7, device=device, requires_grad=autograd) - def fct_c1_no_grad(x: torch.Tensor, y: torch.Tensor): - h_new = torch.tanh(x[0] + x[1] + y) - c2 = x[1] + y - with torch.no_grad(): - c1 = x[0] + y - return (c1, c2), h_new - with self.assertRaisesRegex( RuntimeError, "The init and carries need to have the same require_grad structure!.*", ): - result = scan_fct(fct_c1_no_grad, (h1, h2), x, dim=dim, reverse=reverse) + result = scan_fct( + get_scan_combine_fn("fct_c1_no_grad", False), + (h1, h2), + x, + dim=dim, + reverse=reverse, + ) @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda @@ -3047,15 +3069,20 @@ def test_scan_closure_combine_fn_with_no_grad_init_carries_equal_grad( h1 = torch.randn(3, 7, device=device, requires_grad=False) h2 = torch.randn(3, 7, device=device, requires_grad=autograd) - def fct_c1_no_grad(x: torch.Tensor, y: torch.Tensor): - h_new = torch.tanh(x[0] + x[1] + y) - c2 = x[1] + y - with torch.no_grad(): - c1 = x[0] + y - return (c1, c2), h_new - - result = scan_fct(fct_c1_no_grad, (h1, h2), x, dim=dim, reverse=reverse) - result_exp = _fake_scan(fct_c1_no_grad, (h1, h2), x, dim=dim, reverse=reverse) + result = scan_fct( + get_scan_combine_fn("fct_c1_no_grad", False), + (h1, h2), + x, + dim=dim, + reverse=reverse, + ) + result_exp = _fake_scan( + get_scan_combine_fn("fct_c1_no_grad", False), + (h1, h2), + x, + dim=dim, + reverse=reverse, + ) self.assertEqual(result, result_exp) if autograd: @@ -3102,7 +3129,7 @@ def fct_ys_no_grad(x: torch.Tensor, y: torch.Tensor): @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) @parametrize("autograd", [False, True]) - def test_scan_closure_combine_fn_with_no_grad_additional_inputs( + def test_scan_closure_combine_fn_with_no_grad_additional_inputs_partial( self, reverse, compile_mode, device, autograd ): dim = 1 @@ -3131,9 +3158,27 @@ def fct_no_grad_bhh_Whh(x: torch.Tensor, y: torch.Tensor): if autograd: self.check_autograd(result[1], result_exp[1], (h, x, W_ih, b_ih)) + @unittest.skipIf(not SM70OrLater, "triton") + @requires_cuda + @parametrize("reverse", [False, True]) + @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) + @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) + @parametrize("autograd", [False, True]) + def test_scan_closure_combine_fn_with_no_grad_additional_inputs_all( + self, reverse, compile_mode, device, autograd + ): + dim = 1 + scan_fct = compile_mode_helper(scan, compile_mode) + x = torch.randn(3, 10, 7, device=device, requires_grad=autograd) + h = torch.randn(3, 7, device=device, requires_grad=autograd) + W_ih = torch.randn(7, 7, device=device, requires_grad=autograd) + b_ih = torch.randn(7, device=device, requires_grad=autograd) + W_hh = torch.randn(7, 7, device=device, requires_grad=autograd) + b_hh = torch.randn(7, device=device, requires_grad=autograd) + def fct_no_grad_bih_Wih_bhh_Whh(x: torch.Tensor, y: torch.Tensor): c_new = x + y - h_new = c_new + 1 + h_new = c_new + x with torch.no_grad(): c_new_no_grad = y @ W_ih + b_ih h_new_no_grad = torch.tanh(x @ W_hh + b_hh) @@ -3150,6 +3195,59 @@ def fct_no_grad_bih_Wih_bhh_Whh(x: torch.Tensor, y: torch.Tensor): if autograd: self.check_autograd(result[1], result_exp[1], (h, x)) + @unittest.skipIf(not SM70OrLater, "triton") + @requires_cuda + @parametrize("reverse", [False, True]) + @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) + @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) + @parametrize("autograd", [False, True]) + def test_scan_closure_combine_fn_carries_ys_same_grad( + self, reverse, compile_mode, device, autograd + ): + dim = 1 + scan_fct = compile_mode_helper(scan, compile_mode) + x = torch.randn(3, 10, 7, device=device, requires_grad=autograd) + h = torch.randn(3, 7, device=device, requires_grad=autograd) + W_ih = torch.randn(7, 7, device=device, requires_grad=autograd) + b_ih = torch.randn(7, device=device, requires_grad=autograd) + W_hh = torch.randn(7, 7, device=device, requires_grad=autograd) + b_hh = torch.randn(7, device=device, requires_grad=autograd) + + def fct_no_grad_bih_Wih_bhh_Whh(x: torch.Tensor, y: torch.Tensor): + c_new = x + y + h_new = c_new + 1 + with torch.no_grad(): + c_new_no_grad = y @ W_ih + b_ih + h_new_no_grad = torch.tanh(x @ W_hh + b_hh) + c_new2 = c_new + c_new_no_grad + h_new2 = h_new + h_new_no_grad + return c_new2, h_new2 + + # The gradient of fct_no_grad_bih_Wih_bhh_Whh has aliased outputs + # which is not currently supported in inductor. + if compile_mode not in ["none", "eager"] and autograd: + with self.assertRaisesRegex( + # Should be + # UnsupportedAliasMutationException, + # "Combine_fn might be aliasing its outputs!.*", + torch._dynamo.exc.BackendCompilerFailed, + "backend='inductor' raised.*", + ): + result = scan_fct( + fct_no_grad_bih_Wih_bhh_Whh, h, x, dim=dim, reverse=reverse + ) + else: + result = scan_fct( + fct_no_grad_bih_Wih_bhh_Whh, h, x, dim=dim, reverse=reverse + ) + result_exp = _fake_scan( + fct_no_grad_bih_Wih_bhh_Whh, h, x, dim=dim, reverse=reverse + ) + self.assertEqual(result, result_exp) + + if autograd: + self.check_autograd(result[1], result_exp[1], (h, x)) + @unittest.skipIf(not SM70OrLater, "triton") @requires_cuda @parametrize("reverse", [False, True]) @@ -3264,7 +3362,7 @@ def f(fct, init, xs): "Observed exception.*", ): gm = make_fx(f, tracing_mode="symbolic")( - get_scan_combine_fn("add", True), init, x + get_scan_combine_fn("add", False), init, x ) @skipIfNoDynamoSupport @@ -3321,7 +3419,7 @@ def f(fct, init, xs): # Correct case gm = make_fx(f, tracing_mode="symbolic")( - get_scan_combine_fn("add", False), init, x + get_scan_combine_fn("add", True), init, x ) self.assertExpectedInline( gm.code.strip(), @@ -3341,7 +3439,7 @@ def forward(self, fct_1, init_1, xs_1): # Check graph backend = EagerAndRecordGraphs() - torch.compile(f, backend=backend)(get_scan_combine_fn("add", False), init, x) + torch.compile(f, backend=backend)(get_scan_combine_fn("add", True), init, x) gm = backend.graphs[0] self.assertExpectedInline( diff --git a/torch/_higher_order_ops/scan.py b/torch/_higher_order_ops/scan.py index d635f0c7fcd22..dd117667393cc 100644 --- a/torch/_higher_order_ops/scan.py +++ b/torch/_higher_order_ops/scan.py @@ -13,6 +13,7 @@ from torch._higher_order_ops.utils import ( _has_potential_branch_input_alias, _has_potential_branch_input_mutation, + _has_potential_branch_output_alias, _set_compilation_env, reenter_make_fx, unique_graph_id, @@ -877,13 +878,19 @@ def scan_functionalize(ctx, combine_fn, init, xs, dim, reverse, additional_input functional_combine_fn, sample_inputs, pre_dispatch=pre_dispatch ): raise UnsupportedAliasMutationException( - "Combine_fn might be modifying the input!" + "Combine_fn might be modifying the input! Please also check the gradient of combine_fn for modifying the input" ) if _has_potential_branch_input_alias( functional_combine_fn, sample_inputs, pre_dispatch=pre_dispatch ): raise UnsupportedAliasMutationException( - "Combine_fn might be aliasing the input!" + "Combine_fn might be aliasing the input! Please also check the gradient of combine_fn for modifying the input" + ) + if _has_potential_branch_output_alias( + functional_combine_fn, sample_inputs, pre_dispatch=pre_dispatch + ): + raise UnsupportedAliasMutationException( + "Combine_fn might be aliasing its outputs! Please also check the gradient of combine_fn for aliasing its outputs" ) ret = scan_op( functional_combine_fn, diff --git a/torch/_higher_order_ops/utils.py b/torch/_higher_order_ops/utils.py index d252eb616178c..7822b67f77b8a 100644 --- a/torch/_higher_order_ops/utils.py +++ b/torch/_higher_order_ops/utils.py @@ -114,12 +114,7 @@ def _set_compilation_env(): torch.fx._symbolic_trace._is_fx_tracing_flag = _old_is_tracing -def _has_potential_branch_input_mutation(branch, inputs, pre_dispatch=False): - """ - Dispatch-trace the branch with inputs and check if - producing graph has mutable op on the input. This is - bit restrictive as the branch must be traceable. - """ +def _create_gm_from_branch(branch, inputs, pre_dispatch): try: gm = make_fx(branch, pre_dispatch=pre_dispatch)(*inputs) except UnsupportedAliasMutationException: @@ -129,6 +124,21 @@ def _has_potential_branch_input_mutation(branch, inputs, pre_dispatch=False): except Exception as e: raise e + return gm + + +def _has_potential_branch_input_mutation(branch, inputs, pre_dispatch=False): + """ + Dispatch-trace the branch with inputs and check if + producing graph has mutable op on the input. This is + bit restrictive as the branch must be traceable. + """ + gm = _create_gm_from_branch(branch, inputs, pre_dispatch) + # this can happen when nested cond_op is + # functionalized + if isinstance(gm, bool): + return gm + def _detect_input_mutation(gm): input_nodes = set() for node in gm.graph.nodes: @@ -160,14 +170,11 @@ def _has_potential_branch_input_alias(branch, inputs, pre_dispatch=False): producing graph has output aliasing the branch input. This is bit restrictive as the branch must be traceable. """ - try: - gm = make_fx(branch, pre_dispatch=pre_dispatch)(*inputs) - except UnsupportedAliasMutationException: - # this can happen when nested cond_op is - # functionalized - return True - except Exception as e: - raise e + gm = _create_gm_from_branch(branch, inputs, pre_dispatch) + # this can happen when nested cond_op is + # functionalized + if isinstance(gm, bool): + return gm def _detect_input_alias(gm): input_storages = set() @@ -197,6 +204,45 @@ def check_alias(out): return _detect_input_alias(gm) +def _has_potential_branch_output_alias(branch, inputs, pre_dispatch=False): + """ + Dispatch-trace the branch with inputs and check if + producing graph has aliasing the branch output. This is + bit restrictive as the branch must be traceable. + """ + gm = _create_gm_from_branch(branch, inputs, pre_dispatch) + # this can happen when nested cond_op is + # functionalized + if isinstance(gm, bool): + return gm + + def _detect_output_alias(gm): + output_storages = set() + for node in gm.graph.nodes: + if node.op == "output": + + def check_alias(out): + if out is not None and "val" in out.meta: + out_storage = StorageWeakRef(out.meta["val"]._typed_storage()) + alias = out_storage in output_storages + output_storages.add(out_storage) + return alias + return False + + if any(pytree.tree_leaves(pytree.tree_map(check_alias, node.args))): + return True + + for _, module in gm.named_children(): + if isinstance(module, torch.fx.GraphModule) and _detect_output_alias( + module + ): + return True + + return False + + return _detect_output_alias(gm) + + def unique_graph_id(proxy_mode, prefix): """Returns a unique name and id for a graph to be added to a proxy_mode tracer""" # There are probably better ways - I know that create_arg has some self incrementing name From ddbdb662adbe2a4c73726713e7819a67428ac410 Mon Sep 17 00:00:00 2001 From: Thomas Bohnstingl Date: Sun, 13 Oct 2024 00:53:41 +0200 Subject: [PATCH 12/12] *) Combined input alias check with output alias check *) Parametrized test to avoid manually increasing compiler cache size --- test/functorch/test_control_flow.py | 56 +++++++++++++-------------- torch/_dynamo/trace_rules.py | 2 +- torch/_higher_order_ops/cond.py | 4 +- torch/_higher_order_ops/hints_wrap.py | 4 +- torch/_higher_order_ops/map.py | 4 +- torch/_higher_order_ops/scan.py | 11 +----- torch/_higher_order_ops/utils.py | 48 ++++------------------- torch/_higher_order_ops/while_loop.py | 4 +- 8 files changed, 46 insertions(+), 87 deletions(-) diff --git a/test/functorch/test_control_flow.py b/test/functorch/test_control_flow.py index 215ae5bd19397..aa3ea646c1ff3 100644 --- a/test/functorch/test_control_flow.py +++ b/test/functorch/test_control_flow.py @@ -2944,13 +2944,13 @@ def test_scan_closure_RNN_parameters_as_inputs(self, reverse, device, autograd): @requires_cuda @parametrize("reverse", [False, True]) @parametrize("compile_mode", ["none", "eager", "compile", "compile_dynamic_shape"]) + @parametrize( + "partial_grad", ["xs", "init", "additional_inputs", "complex", "random"] + ) @parametrize("device", [torch.device("cpu"), torch.device("cuda")]) - def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): - import random - - torch._dynamo.reset() - torch._dynamo.config.cache_size_limit = 100 - + def test_scan_closure_RNN_partial_autograd( + self, reverse, compile_mode, partial_grad, device + ): dim = 1 scan_fct = compile_mode_helper(scan, compile_mode) @@ -2959,28 +2959,28 @@ def test_scan_closure_RNN_partial_autograd(self, reverse, compile_mode, device): # The last four are the additional_inputs autograds = [] - # Basic test - autograds.append([True, True, True, True, True, True, True, True]) - - # xs tests - autograds.append([True, False, True, True, True, True, True, True]) - autograds.append([False, False, True, True, True, True, True, True]) - - # init tests - autograds.append([True, True, False, True, True, True, True, True]) - autograds.append([True, True, False, False, True, True, True, True]) - - # additional input tests - autograds.append([True, True, True, True, False, True, False, True]) - autograds.append([True, True, True, True, False, False, False, False]) - - # # Complex cases - autograds.append([True, False, False, False, False, False, False, True]) - autograds.append([False, False, True, True, False, False, False, True]) - - # Random tests - for _ in range(10): - autograds.append([bool(random.randint(0, 1)) for _ in range(8)]) + if partial_grad == "xs": + # xs tests + autograds.append([True, False, True, True, True, True, True, True]) + autograds.append([False, False, True, True, True, True, True, True]) + elif partial_grad == "init": + # init tests + autograds.append([True, True, False, True, True, True, True, True]) + autograds.append([True, True, False, False, True, True, True, True]) + elif partial_grad == "additional_inputs": + # additional input tests + autograds.append([True, True, True, True, False, True, False, True]) + autograds.append([True, True, True, True, False, False, False, False]) + elif partial_grad == "complex": + # complex cases + autograds.append([True, False, False, False, False, False, False, True]) + autograds.append([False, False, True, True, False, False, False, True]) + elif partial_grad == "random": + # random tests + import random + + for _ in range(5): + autograds.append([bool(random.randint(0, 1)) for _ in range(8)]) for autograd in autograds: x = torch.randn(3, 10, 5, device=device, requires_grad=autograd[0]) diff --git a/torch/_dynamo/trace_rules.py b/torch/_dynamo/trace_rules.py index ed81b26191b4d..bd1af47063464 100644 --- a/torch/_dynamo/trace_rules.py +++ b/torch/_dynamo/trace_rules.py @@ -2284,7 +2284,7 @@ "torch._guards.compile_context", "torch._guards.detect_fake_mode", "torch._guards.tracing", - "torch._higher_order_ops.map._has_potential_branch_input_alias", + "torch._higher_order_ops.map._has_potential_branch_input_output_alias", "torch._higher_order_ops.map._has_potential_branch_input_mutation", "torch._higher_order_ops.map._stack_pytree", "torch._higher_order_ops.map._unstack_pytree", diff --git a/torch/_higher_order_ops/cond.py b/torch/_higher_order_ops/cond.py index 0467e2899adc2..aa820356948c7 100644 --- a/torch/_higher_order_ops/cond.py +++ b/torch/_higher_order_ops/cond.py @@ -16,8 +16,8 @@ from torch._functorch.utils import exposed_in from torch._guards import detect_fake_mode from torch._higher_order_ops.utils import ( - _has_potential_branch_input_alias, _has_potential_branch_input_mutation, + _has_potential_branch_input_output_alias, _maybe_run_with_interpreter, _set_compilation_env, reenter_make_fx, @@ -469,7 +469,7 @@ def cond_func(ctx, pred, true_fn, false_fn, inputs): "One of torch.cond branch might be modifying the input!" ) for branch in [true_fn, false_fn]: - if _has_potential_branch_input_alias( + if _has_potential_branch_input_output_alias( branch, unwrapped_inputs, pre_dispatch=pre_dispatch ): raise UnsupportedAliasMutationException( diff --git a/torch/_higher_order_ops/hints_wrap.py b/torch/_higher_order_ops/hints_wrap.py index c211d40561463..bd6faa9aedf28 100644 --- a/torch/_higher_order_ops/hints_wrap.py +++ b/torch/_higher_order_ops/hints_wrap.py @@ -3,8 +3,8 @@ import torch.utils._pytree as pytree from torch._C import DispatchKey from torch._higher_order_ops.utils import ( - _has_potential_branch_input_alias, _has_potential_branch_input_mutation, + _has_potential_branch_input_output_alias, autograd_not_implemented, reenter_make_fx, unique_graph_id, @@ -103,7 +103,7 @@ def hints_wrapper_functionalize(ctx, body_fn, args, kwargs, hints): raise UnsupportedAliasMutationException( "body_fn of hints_wrapper might be modifying the input!" ) - if _has_potential_branch_input_alias( + if _has_potential_branch_input_output_alias( functional_body_fn, unwrapped_args, pre_dispatch=pre_dispatch ): raise UnsupportedAliasMutationException( diff --git a/torch/_higher_order_ops/map.py b/torch/_higher_order_ops/map.py index d57d68d5e473f..244568395362f 100644 --- a/torch/_higher_order_ops/map.py +++ b/torch/_higher_order_ops/map.py @@ -5,8 +5,8 @@ from torch._dispatch.python import suspend_functionalization from torch._functorch.aot_autograd import AOTConfig, create_joint from torch._higher_order_ops.utils import ( - _has_potential_branch_input_alias, _has_potential_branch_input_mutation, + _has_potential_branch_input_output_alias, _maybe_run_with_interpreter, reenter_make_fx, UnsupportedAliasMutationException, @@ -255,7 +255,7 @@ def map_functionalize(ctx, f, xs, pos_args): ): raise UnsupportedAliasMutationException("torch.map is mutating the input!") - if _has_potential_branch_input_alias( + if _has_potential_branch_input_output_alias( f, example_inputs, pre_dispatch=pre_dispatch ): raise UnsupportedAliasMutationException("torch.map is aliasing the input!") diff --git a/torch/_higher_order_ops/scan.py b/torch/_higher_order_ops/scan.py index dd117667393cc..2df8b22199cec 100644 --- a/torch/_higher_order_ops/scan.py +++ b/torch/_higher_order_ops/scan.py @@ -11,9 +11,8 @@ from torch._C import DispatchKey from torch._dispatch.python import suspend_functionalization from torch._higher_order_ops.utils import ( - _has_potential_branch_input_alias, _has_potential_branch_input_mutation, - _has_potential_branch_output_alias, + _has_potential_branch_input_output_alias, _set_compilation_env, reenter_make_fx, unique_graph_id, @@ -880,18 +879,12 @@ def scan_functionalize(ctx, combine_fn, init, xs, dim, reverse, additional_input raise UnsupportedAliasMutationException( "Combine_fn might be modifying the input! Please also check the gradient of combine_fn for modifying the input" ) - if _has_potential_branch_input_alias( + if _has_potential_branch_input_output_alias( functional_combine_fn, sample_inputs, pre_dispatch=pre_dispatch ): raise UnsupportedAliasMutationException( "Combine_fn might be aliasing the input! Please also check the gradient of combine_fn for modifying the input" ) - if _has_potential_branch_output_alias( - functional_combine_fn, sample_inputs, pre_dispatch=pre_dispatch - ): - raise UnsupportedAliasMutationException( - "Combine_fn might be aliasing its outputs! Please also check the gradient of combine_fn for aliasing its outputs" - ) ret = scan_op( functional_combine_fn, unwrapped_init, diff --git a/torch/_higher_order_ops/utils.py b/torch/_higher_order_ops/utils.py index 7822b67f77b8a..0479a5b78eef3 100644 --- a/torch/_higher_order_ops/utils.py +++ b/torch/_higher_order_ops/utils.py @@ -164,7 +164,7 @@ def _detect_input_mutation(gm): return _detect_input_mutation(gm) -def _has_potential_branch_input_alias(branch, inputs, pre_dispatch=False): +def _has_potential_branch_input_output_alias(branch, inputs, pre_dispatch=False): """ Dispatch-trace the branch with inputs and check if producing graph has output aliasing the branch input. This is @@ -185,46 +185,14 @@ def _detect_input_alias(gm): if node.op == "placeholder" and "val" in node.meta: input_storages.add(StorageWeakRef(node.meta["val"]._typed_storage())) if node.op == "output": + output_storages = set() def check_alias(out): if out is not None and "val" in out.meta: out_storage = StorageWeakRef(out.meta["val"]._typed_storage()) - return out_storage in input_storages - return False - - if any(pytree.tree_leaves(pytree.tree_map(check_alias, node.args))): - return True - - for _, module in gm.named_children(): - if isinstance(module, torch.fx.GraphModule) and _detect_input_alias(module): - return True - - return False - - return _detect_input_alias(gm) - - -def _has_potential_branch_output_alias(branch, inputs, pre_dispatch=False): - """ - Dispatch-trace the branch with inputs and check if - producing graph has aliasing the branch output. This is - bit restrictive as the branch must be traceable. - """ - gm = _create_gm_from_branch(branch, inputs, pre_dispatch) - # this can happen when nested cond_op is - # functionalized - if isinstance(gm, bool): - return gm - - def _detect_output_alias(gm): - output_storages = set() - for node in gm.graph.nodes: - if node.op == "output": - - def check_alias(out): - if out is not None and "val" in out.meta: - out_storage = StorageWeakRef(out.meta["val"]._typed_storage()) - alias = out_storage in output_storages + alias = (out_storage in input_storages) or ( + out_storage in output_storages + ) output_storages.add(out_storage) return alias return False @@ -233,14 +201,12 @@ def check_alias(out): return True for _, module in gm.named_children(): - if isinstance(module, torch.fx.GraphModule) and _detect_output_alias( - module - ): + if isinstance(module, torch.fx.GraphModule) and _detect_input_alias(module): return True return False - return _detect_output_alias(gm) + return _detect_input_alias(gm) def unique_graph_id(proxy_mode, prefix): diff --git a/torch/_higher_order_ops/while_loop.py b/torch/_higher_order_ops/while_loop.py index f14321842f40b..04cfc8c152a75 100644 --- a/torch/_higher_order_ops/while_loop.py +++ b/torch/_higher_order_ops/while_loop.py @@ -5,8 +5,8 @@ import torch.utils._pytree as pytree from torch._C import DispatchKey from torch._higher_order_ops.utils import ( - _has_potential_branch_input_alias, _has_potential_branch_input_mutation, + _has_potential_branch_input_output_alias, _maybe_run_with_interpreter, _set_compilation_env, autograd_not_implemented, @@ -266,7 +266,7 @@ def while_loop_func(ctx, cond_fn, body_fn, carried_inputs, additional_inputs): f"torch.while_loop's {fn_name} might be modifying the input!" ) - if _has_potential_branch_input_alias( + if _has_potential_branch_input_output_alias( fn, unwrapped_inputs, pre_dispatch=pre_dispatch ): raise UnsupportedAliasMutationException(