[AOTAutograd] add export entrypoints

[bdhirsh]

The main addition in this PR is two new API's in AOTAutograd.

APIs

aot_export_module: Given a module, exports it into a functionalized FX graph. Returns an fx.GraphModule, GraphSignature pair. The GraphSignature tells you various information about the graph, such as which graph inputs correspond to module params/buffers (and their fqn's), how to pytree-ify the inputs and the outputs of the graph. If you specify trace_joint=True, then you'll get back a joint forward-backward graph, that also returns parameter gradients in addition to the user outputs.

There are several restrictions on this API, detailed in the comments. The most notable one is probably that this API does not handle partial graphs: If you want a backward graph, then you module's forward function is required to return a scalar loss that we can backprop through. It also does not support capturing the optimizer step.

I (gratefully) used @SherlockNoMad and @suo's internal version of the GraphSignature object for this API, with a few minor changes in order to integrate it into AOTAutograd.

aot_export_joint_simple: Given a function, we'll trace it into a joint forward-backward graph and return it. Unlike the above API, the function is not required to return a scalar loss. However, this API makes the guarantee that you do not need to make any calling convention changes between the original function, and the exported one, provided that you do that you do the following:

• If you pass trace_joint=False, no work is needed: we'll export a functionalized forward graph with the same set of inputs as the original function
• If you pass trace_joint=True, then you will need to manually use the default_partitioner or min_cut_partitioner from functorch. If you do, and get back a fw and bw graph, then the forward graph will be runnable identically to the original user function.

The main use case for this API is higher order ops: a higher order op like torch.cond() can implement its derivative formula by using this API to export a joint graph (for both the true subgraph and the false subgraph), partition it into a fw/bw graph, and run cond on the true_bw, false_bw subgraphs. cc @zou3519 @Chillee

Implementation Strategy

A lot of the work in this PR went in to trying to find a reasonable way to re-use existing AOTAutograd components to expose these API's. Concretely:

• The two new API's are both thin wrappers around _aot_export_function: this is a general purpose export API, that just re-uses create_aot_dispatcher_function. If we want to add e.g. an export API that includes the optimizer step in the future, we could probably implement it using _aot_export_function.
• aot_export_module works extra hard to re-use as much of AOTAutograd as possible. For example, when tracing an inference graph, I perform the export under torch.no_grad() to make sure we don't accidentally trace out a backwards graph. When exporting a joint graph, I manually .detach() all user outputs except the loss, to make sure that we don't accidentally compute gradients for any other user outputs (even if the user forgot to manually detach them).
• A large portion of aot_export_module comes from parsing out and creating a GraphSignature object. We discussed a few weeks ago that there's potentially a lot more information that we could stuff into this object (see doc). For now, I ended up deciding to support the more limited use case of exporting a fwd-bwd full graph, without some of the extra annotations in that doc (for example, if we were to export partial graphs, we would need annotations for saved activations). My thought is that once a more concrete use case comes up that the existing API doesn't satisfy, we can revisit the annotations then.
• I factored out create_functional_call() and create_tree_flattened_fn() for pytree-flattening and lifting-params-and-buffers, since I also need them in the export code
• I added an AOTConfig.is_export flag. The export API re-uses all of the same code paths as the rest of AOTAutograd, but there are a few points where we need to either exit early (and avoid making a runtime epilogue), or add extra error checking, that is only valuable for export.
• aot_dispatch_autograd() now exits early if it's being called in an export context, so it returns the full graph instead of also trying to create an autograd.Function. I think we probably want to factor this out, although I figured it would be safer to wait a bit for clarity on how functional RNG works with export.

Stack from ghstack (oldest at bottom):

• -> [AOTAutograd] add export entrypoints #100587
• aot_autograd: factor out runtime epilogue from aot_dispatch_base #100586

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 59a4374:
💚 Looks good so far! There are no failures yet. 💚

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

[bdhirsh]

TODO: I added one large test for aot_export_module, but I still need to add some tests for aot_export_joint_simplified, and for all of the error cases where we ban certain functions/modules (stubbed out for now).

[ezyang]

TODO I suppose

[bdhirsh]

yep, will add these in before landing

[ezyang]

worth saying explicitly what the str denotes (node.target, amirite?)

[bdhirsh]

node.name 😛 yep, will update (it came from here: https://github.com/pytorch/pytorch/pull/100587/files#diff-df954bbf954d2dcb81f687876053267ffa4ddb36ed86b7d2bd76319ff2b94416R3300)

[ezyang]

Confusing. Maybe worth a type alias?

[SherlockNoMad]

inputs_to_parameters are recording the name mapping.
key is graph input name
value is parameter fqn

[bdhirsh]

I ended up making some type aliases. FQN, InputName and OutputName

[ezyang]

nit: If you actually want to docblock it, the docblock goes in the function, not on top.

[ezyang]

Any difference in code movement here?

[bdhirsh]

nope, copy paste

[ezyang]

Any change from code motion here?

[bdhirsh]

nope, just a straight move out into its own function

[ezyang]

OK to do this inplace?

[bdhirsh]

I now error instead of detaching

[ezyang]

If trace_joint = True but output_loss_index is None, what is the behavior?

[SherlockNoMad]

This should be invalid input.

[bdhirsh]

fixed

[ezyang]

Use of asserts here is inappropriate, because users can violate these asserts via inputs. Do real RuntimeErrors with proper error messages.

[bdhirsh]

updated them all to RuntimeErrors, thanks for the callout

[ezyang]

I like to preemptively file bugs and link 'em in the error message, so that I know I will get notified if someone does need it.

[bdhirsh]

Good call - updated, issue here #101192

[ezyang]

uhh, is args fake or not here?

[bdhirsh]

good catch... we fakeify later, so they're not guaranteed to be fake here. I'll ensure fakeness

[ezyang]

Looks great! Ship it!

[SherlockNoMad]

Let's add a comment for this field, (or have a better name make it explicit)
key is buffer graph output name
value is buffer fqn

[SherlockNoMad]

love it!

[SherlockNoMad]

loss could be 1d scalar? A more robust condition is

needed_tangents[0].numel() == 1

[SherlockNoMad]

slightly concern about returning a different type for aot_dispatch_autograd.
The function return signature will be confused by this...

I wonder if it can be refactored into sth like this

def aot_dispatch_autograd_graph_version -> GraphModule:
      ....
      return fx_g

def aot_dispatch_autograd -> Callable:
     fx_g = aot_dispatch_autograd_graph_version(...)
     fn = create_runtime_wrapper(fx_g)
     return fn

Please also feel free to ignore this nitpick, coz I know it might be hard to do so in practice.

[bdhirsh]

Agreed - I originally planned to push this off and wait for more clarity on how functionalized RNG fits into the export picture. But I may as well do this refactor now (and I'll leave functionalized RNG only into the callable version for now).

[bdhirsh]

updated

[SherlockNoMad]

nitpick... this should be be graph module API....

[SherlockNoMad]

same here... this should be be graph module API....

[SherlockNoMad]

if trace_joint is provided?

[SherlockNoMad]

This is great!

[SherlockNoMad]

This should throw? We only allow loss to require_grad.
calling detach for user might be surprising behavior.

[bdhirsh]

fair, I'll error instead.

[SherlockNoMad]

Just trying to understand the use case here...

Say we have a torch.cond op, that we wish to export. Both branches will be passed into this aot_export_joint_simple function?

I am also having a hard time imagine how trace_joint will behave for high-order ops...
e.g. The backward subgraph for torch.cond's branch is not directly connected to the forward subgraph... so there are not strictly "jointed" ?

[bdhirsh]

So the main use for this API isn't really strictly for export - it's so that we can generate the backward formulas for higher order ops, like torch.cond, map, etc.

I have a test coming soon, but I'm picturing that this API will be used by higher order ops roughly like this:

I'm open to feedback on exactly what this API should return (cc @zou3519). I went off the thought process of:
(1) partitioning a joint into a fw/bw is only a few lines of code (see code snippet above)
(2) returning a joint graph is a bit more general, and alleviates AOTAutograd from having to tell the caller how to use / call the partitioned backward graph (although maybe this isn't a real concern).

[SherlockNoMad]

Thanks a lot @bdhirsh for pushing this to finish line!!!

[bdhirsh]

hmm. I added this logic to create_joint():

                # If our output is a scalar loss, we don't need to pass in tangents.
                if len(needed_tangents) == 1 and needed_tangents[0].numel() == 1:
                    backward_out = torch.autograd.grad(
                        needed_outs,
                        grad_primals,
                        allow_unused=True,
                    )
                else:
                    backward_out = torch.autograd.grad(
                        needed_outs,
                        grad_primals,
                        grad_outputs=needed_tangents,
                        allow_unused=True,
                    )

Where the idea was that for export, we want to make sure that there are no tangent inputs to our joint graph, since we know we're just calling .backward() (which will implicitly pass a scalar 1 as the tangent).

This doesn't actually work though, because we don't know what the value of the scalar is. If the scalar is a 1, then the above is ok, but if it's something else, then we need to actually use the tangent value. And we don't know this info at trace time, since we don't know value of the tangent.

The easiest thing to do is probably to plumb aot_config into create_joint() and branch off of aot_config.is_export

[bdhirsh]

@ezyang I don't love this but let me know what you think. My goal was to avoid having to duplicate the logic for calling autograd.grad() - if we want a joint graph with no tangent inputs, we need to make sure we don't specify grad_outputs=.... I originally thought I could detect that automatically by checking to see if our tangent is a single scalar, but that's not true - if the scalar is not 1, then it should still be an input to our graph (but we don't know its value at trace time).

[ezyang]

@pytorchbot merge

[pytorchmergebot]

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

Advanced Debugging

Check the merge workflow status
here