[dynamo] AutogradFunctionMethodHigherOrderVariable check for new guards is broken
Labels: module: dynamo

AutogradFunctionMethodHigherOrderVariable has a check for new guards being added in the following places:
https://github.com/pytorch/pytorch/blob/f0cde8613c4c8814e157c0a742187a91aa72a009/torch/_dynamo/variables/higher_order_ops.py#L1091
https://github.com/pytorch/pytorch/blob/f0cde8613c4c8814e157c0a742187a91aa72a009/torch/_dynamo/variables/higher_order_ops.py#L1115
https://github.com/pytorch/pytorch/blob/f0cde8613c4c8814e157c0a742187a91aa72a009/torch/_dynamo/variables/higher_order_ops.py#L1122-L1123

As written, this check does nothing because we are storing a *reference* in pre_guards that just gets mutated.  So `pre_guards is post_guards` evaluates `True`.

The following will make the code do as intended:
```patch
diff --git a/torch/_dynamo/variables/higher_order_ops.py b/torch/_dynamo/variables/higher_order_ops.py
index 2292c71f048..5d7d87f810c 100644
--- a/torch/_dynamo/variables/higher_order_ops.py
+++ b/torch/_dynamo/variables/higher_order_ops.py
@@ -1088,7 +1088,7 @@ class AutogradFunctionMethodHigherOrderVariable(TorchHigherOrderOperatorVariable
         else:
             fn = TorchVariable(self.value)
         checkpoint = tx.copy_graphstate()
-        pre_guards = tx.output.guards
+        pre_guards = tx.output.guards.clone()
         graph_checkpoint = tx.output.graph
 
         # TODO: Support kwargs
diff --git a/torch/_guards.py b/torch/_guards.py
index e532a32cdd2..4f9e874476e 100644
--- a/torch/_guards.py
+++ b/torch/_guards.py
@@ -483,6 +483,8 @@ class GuardsSet:
         for o in others:
             for g in o:
                 self.add(g, skip=1)

+    def clone(self):
+        return GuardsSet(set(self.inner))
 
 
 class GuardsContext(Checkpointable[GuardsCheckpointState]):
```

However, it causes a test to fail:
```
____________________________________________ ReproTests.test_hf_xsoftmax_training ____________________________________________
Traceback (most recent call last):
  File "/home/jansel/conda/envs/pytorch/lib/python3.10/unittest/case.py", line 59, in testPartExecutor
    yield
  File "/home/jansel/conda/envs/pytorch/lib/python3.10/unittest/case.py", line 591, in run
    self._callTestMethod(testMethod)
  File "/home/jansel/conda/envs/pytorch/lib/python3.10/unittest/case.py", line 549, in _callTestMethod
    method()
  File "/home/jansel/pytorch/torch/testing/_internal/common_utils.py", line 2453, in wrapper
    method(*args, **kwargs)
  File "/home/jansel/pytorch/test/dynamo/test_repros.py", line 3163, in test_hf_xsoftmax_training
    self.assertEqual(dict(counters["frames"]), {"total": 1, "ok": 1})
  File "/home/jansel/pytorch/torch/testing/_internal/common_utils.py", line 3356, in assertEqual
    raise error_metas.pop()[0].to_error(
AssertionError: Scalars are not equal!

Expected 1 but got 2.
Absolute difference: 1
Relative difference: 1.0

The failure occurred for item ['ok']

To execute this test, run the following from the base repo dir:
     python test/dynamo/test_repros.py -k test_hf_xsoftmax_training
```

I think this piece of code needs to be revisited, as I am not sure if the graph break it adds is correct.




cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng @anijain2305