🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/136184

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit f8ed313 with merge base 0a4bcbf ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@zou3519 Let's discuss if this is the proper way at here. I think we only need to update the example_value of the autograd.Function.appy output, which is inferred from speculating fwd graph. Now I'm switching to a separate fake tensor prop to get the example_value, and it works well in common cases. However, there are several cases don't work well:

• The fake tensor prop still need to prop under torch.no_grad mode, check this unit test:
  

  pytorch/test/dynamo/test_autograd_function.py

  Line 442 in e248c1d

  def test_fwd_no_grad(self):

• The fake tensor prop can't handle arbitrary user defined objects as input, check this unit test:
  

  pytorch/test/dynamo/test_autograd_function.py

  Line 515 in e248c1d

  def test_user_defined_object_as_input(self):

• We generate more node in the current FX graph, which seems fine. The following unit test is failed because of cnt.op_count doesn't match anymore.
  

  pytorch/test/dynamo/test_autograd_function.py

  Line 1221 in e248c1d

  def test_tuple_arg(self):

• The fake tensor prop doesn't work with Triton, since Triton was wrapped as another HOP.
  

  pytorch/test/dynamo/test_autograd_function.py

  Line 1249 in e248c1d

  def test_triton_kernel_basic(self):

I think all of these issues are resolved by speculate_subgraph, and probably it's the reason that we need it. So if we leverage on fake tensor prop to get example_value, seems we need to duplicate these features again.

And especially the first issue above, seems we still need to fake prop under no_grad mode, which seems we are back to the starting point.

@zou3519 I checked all of these failures and found they are from two reasons:

• The added fake prop doesn't handle a pathological scenario well: If users define an autograd function with a static forward method but also pass in self as the first argument, e.g, the failed DebertaForQuestionAnswering. This is because the new added fake prop using nn.Module._call_impl which offsets self and causes positional arguments mismatch.

• Dynamic shapes related failure

torch._dynamo.exc.InternalTorchDynamoError: PendingUnbackedSymbolNotFound: Pending unbacked symbols {u5, u4} not in returned outputs FakeTensor(..., size=(u0 + u1,), grad_fn=<ApplyTemplateBackward>) ((1,), 0).
Did you accidentally call new_dynamic_size() or item() more times than you needed to in your fake implementation?
For more help, see https://docs.google.com/document/d/1RWrH-3wLEpzR9kCS6gGBNen_-Fs-8PVbWWFE5AcgeWE/edit

I think this is because we create symbols (u0 + u1) during speculate_subgraph, but when do the fake prop I added, we are creating new symbols (u5, u4) rather than directly refer to the original one. This seems only happens when capture_scalar_outputs=True. cc @ezyang

• The other flex attention related failures are trunk issues, which is not relevant to this PR.

I'm not terribly familiar with the context behind this PR, but if you are doing a fake prop and then throwing away the results, you might want a pattern similar to what we did for aot autograd metadata analysis:

        fake_mode = detect_fake_mode()
        if fake_mode and (shape_env := fake_mode.shape_env):
            suppress_pending = shape_env.ignore_fresh_unbacked_symbols()
        with disable_above, mode, suppress_pending:
            # precondition: The passed in function already handles unflattening inputs + flattening outputs
            flat_f_args = pytree.tree_map(_to_fun, flat_args)
            flat_f_outs = f(*flat_f_args)
            # We didn't do any tracing, so we don't need to process the
            # unbacked symbols, they will just disappear into the ether.
            # Also, prevent memoization from applying.
            if fake_mode:
                fake_mode.epoch += 1
                fake_mode.reset_nt_tensor_id_counter()

Ignore fresh unbacked symbols is the important part, but I'm not sure if you need to rev the epoch or not.

If these new unbacked symbols are going into the FX graph, though, you need to do something like PropagateUnbackedSymInts instead

Context: the way Dynamo handles autograd.Function isn't faithful to eager-mode PyTorch when the inputs have requires_grad=True. What Dynamo does today is:

• we trace out the autograd.Function forward and backward
• we put these into a HOP
• we put a call to the HOP in the graph.
• we set the example values on the output of the HOP to be the FakeTensor outputs of tracing the forward

The FakeTensor outputs of tracing the forward are not actually faithful to eager-mode PyTorch. The problem is we trace the forward under torch.no_grad(), so the FakeTensor outputs have requires_grad=False. In eager-mode, the eager outputs of the autograd.Function have requires_grad=True.

There's two approaches we can do to resolve this:

1. Add a new step that does FakeTensor prop using the autograd.Function, so that the outputs have the correct required_grad-ness, and use these as the FakeTensor outputs
2. Somehow "fix up" the output of the HOP by emulating what eager-mode autograd.Function does

We tried to do (1) in this PR, because it seemed the easiest, but are running into the unbacked symint thing because we do FakeTensor prop twice: once when tracing out the autograd.Function forward, and once in a new step that does FakeTensor prop using the autograd.Function.

If these new unbacked symbols are going into the FX graph, though, you need to do something like PropagateUnbackedSymInts instead

They are going into the FX graph. If we keep going down this design -- if there are symints in both traces, then we probably need to indicate that they are the same

Is help still needed on this? Might need to hop on a VC

@ezyang yes, we still need help. Yanbo's on recharge so we can wait until he comes back.

@yanboliang @zou3519 Hi, thanks for these findings and any updates here?

@xinyu-intel I will update this PR and get this in soon.

@zou3519 @ezyang This is ready for reviewing.

Endorsing test changes

@pytorchbot merge

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team