Repository: pytorch
Issue: [export] Remove_effect_tokens_pass does not work well with other HOO
Labels: triaged,oncall: pt2,module: higher order operators,export-triaged,oncall: export,module: pt2-dispatcher

### 🐛 Describe the bug

```python
    def test_with_effects_no_grad(self):
        class M(torch.nn.Module):
            def forward(self, x, y):
                with torch.no_grad():
                    torch._print("moo")
                    z = x + y
                return x * z 
            
        ep = torch.export.export(M(), (torch.ones(3, 3), torch.ones(3, 3)))
        print(ep)
        inp = (torch.randn(3, 3), torch.randn(3, 3))
        m = ep.module()
        self.assertTrue(torch.allclose(M()(*inp), m(*inp)))
```

Results in the following exported program:
```python
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, token: "f32[0]", x: "f32[3, 3]", y: "f32[3, 3]"):
            # No stacktrace found for following nodes
            submod_3 = self.submod_1
            wrap_with_set_grad_enabled = torch._higher_order_ops.wrap.wrap_with_set_grad_enabled(False, submod_3, token, x, y);  submod_3 = token = y = None
            
             # File: /data/users/angelayi/pytorch/moo.py:518 in forward, code: z = x + y
            add: "f32[3, 3]" = wrap_with_set_grad_enabled[0]
            
             # File: /data/users/angelayi/pytorch/moo.py:517 in forward, code: torch._print("moo")
            getitem_2: "f32[0]" = wrap_with_set_grad_enabled[1];  wrap_with_set_grad_enabled = None
            
             # File: /data/users/angelayi/pytorch/moo.py:519 in forward, code: return x * z
            mul: "f32[3, 3]" = torch.ops.aten.mul.Tensor(x, add);  x = add = None
            return (getitem_2, mul)
            
        class GraphModule(torch.nn.Module):
            def forward(self, token: "f32[0]", x: "f32[3, 3]", y: "f32[3, 3]"):
                 # File: /data/users/angelayi/pytorch/moo.py:517 in forward, code: torch._print("moo")
                with_effects = torch._higher_order_ops.effects.with_effects(token, torch.ops.aten._print.default, 'moo');  token = None
                getitem: "f32[0]" = with_effects[0];  with_effects = None
                
                 # File: /data/users/angelayi/pytorch/moo.py:518 in forward, code: z = x + y
                add: "f32[3, 3]" = torch.ops.aten.add.Tensor(x, y);  x = y = None
                return (add, getitem)
                
Graph signature: ExportGraphSignature(input_specs=[InputSpec(kind=<InputKind.TOKEN: 6>, arg=TokenArgument(name='token'), target=None, persistent=None), InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='x'), target=None, persistent=None), InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='y'), target=None, persistent=None)], output_specs=[OutputSpec(kind=<OutputKind.TOKEN: 7>, arg=TokenArgument(name='getitem_2'), target=None), OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='mul'), target=None)])
Range constraints: {}
```

But when we call ep.module(), it fails with:
```
======================================================================
FAIL: test_with_effects_no_grad (__main__.TestMoo)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/data/users/angelayi/pytorch/torch/testing/_internal/common_utils.py", line 2918, in wrapper
    method(*args, **kwargs)
  File "/data/users/angelayi/pytorch/moo.py", line 524, in test_with_effects_no_grad
    ep.module()
  File "/data/users/angelayi/pytorch/torch/export/exported_program.py", line 945, in module
    module = _unlift_exported_program_lifted_states(self)
  File "/data/users/angelayi/pytorch/torch/export/_unlift.py", line 291, in _unlift_exported_program_lifted_states
    ep = _remove_effect_tokens(ep)
  File "/data/users/angelayi/pytorch/torch/export/_remove_effect_tokens_pass.py", line 170, in _remove_effect_tokens
    _remove_effect_tokens_from_graph_helper(
  File "/data/users/angelayi/pytorch/torch/export/_remove_effect_tokens_pass.py", line 84, in _remove_effect_tokens_from_graph_helper
    with ep.graph.inserting_before(node):
  File "/data/users/angelayi/pytorch/torch/fx/graph.py", line 1095, in inserting_before
    assert n.graph == self, "Node to insert before is not in graph."
AssertionError: Node to insert before is not in graph.
```

because we try to remove `with_effects` from the toplevel module, but are unable to do so since it's part of the set_grad_enabled subgraph. 

### Versions

main

cc @ezyang @chauhang @penguinwu @zou3519 @ydwu4 @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @suo @bdhirsh